Skip to content

Commit

Permalink
Merge #103307
Browse files Browse the repository at this point in the history
103307: roachprod: add command to self update the roachprod binary r=srosenberg,rail,herkolategan a=smg260

This PR contains 2 ways to get the latest `roachprod`.

1. `update` command as part of roachprod itself: \
\
`roachprod update` will check and download the latest binary for the current platform, and optionally update the running roachprod.
\
This uses the [TeamCity REST API](https://www.jetbrains.com/help/teamcity/rest/teamcity-rest-api-documentation.html) with guest authentication to find the latest successful build (on `master`) from which to download the binary.
\
When proceeding with an update, the existing binary is renamed with a `.bak` prefix so that it may be reverted, either manually or via `roachprod update --revert`. Permissions are copied from the existing binary.

2. `scripts/roachprod-get-latest.sh` shell script \
\
Downloads the latest roachprod binary from TeamCity to the specified (or default current) directory. Has basic checks for `curl` and confirming whether to overwrite any existing roachprod.

The builds used by both the binary, and the script are [here](https://teamcity.cockroachdb.com/project.html?projectId=Cockroach_Ci_Builds&branch_Cockroach_Ci_Builds=%3Cdefault%3E)


Note:  
- The linter prevents direct use of http.Get, so that meant creating proto files for the TeamCity rest API responses.
- The existing `httputil` has no accommodations for unmarshalling a subset of fields in a json response, hence the addition of an `IgnoreUnknownFields` option.
- This should work as long as our build remain public

Epic: none
Fixes: #97311

Release note: None

Co-authored-by: Miral Gadani <[email protected]>
  • Loading branch information
craig[bot] and Miral Gadani committed Jun 1, 2023
2 parents 6eebd78 + 73bff11 commit bf25b7b
Show file tree
Hide file tree
Showing 12 changed files with 389 additions and 6 deletions.
2 changes: 2 additions & 0 deletions pkg/BUILD.bazel
Original file line number Diff line number Diff line change
Expand Up @@ -1067,6 +1067,7 @@ GO_TARGETS = [
"//pkg/cmd/roachprod-stress:roachprod-stress",
"//pkg/cmd/roachprod-stress:roachprod-stress_lib",
"//pkg/cmd/roachprod/grafana:grafana",
"//pkg/cmd/roachprod/upgrade:upgrade",
"//pkg/cmd/roachprod:roachprod",
"//pkg/cmd/roachprod:roachprod_lib",
"//pkg/cmd/roachtest/cluster:cluster",
Expand Down Expand Up @@ -2647,6 +2648,7 @@ GET_X_DATA_TARGETS = [
"//pkg/cmd/roachprod-microbench/google:get_x_data",
"//pkg/cmd/roachprod-stress:get_x_data",
"//pkg/cmd/roachprod/grafana:get_x_data",
"//pkg/cmd/roachprod/upgrade:get_x_data",
"//pkg/cmd/roachtest:get_x_data",
"//pkg/cmd/roachtest/cluster:get_x_data",
"//pkg/cmd/roachtest/clusterstats:get_x_data",
Expand Down
1 change: 1 addition & 0 deletions pkg/cmd/roachprod/BUILD.bazel
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@ go_library(
deps = [
"//pkg/build",
"//pkg/cmd/roachprod/grafana",
"//pkg/cmd/roachprod/upgrade",
"//pkg/roachprod",
"//pkg/roachprod/config",
"//pkg/roachprod/errors",
Expand Down
5 changes: 5 additions & 0 deletions pkg/cmd/roachprod/flags.go
Original file line number Diff line number Diff line change
Expand Up @@ -78,6 +78,8 @@ var (

// hostCluster is used for multi-tenant functionality.
hostCluster string

revertUpdate bool
)

func initFlags() {
Expand Down Expand Up @@ -303,6 +305,9 @@ Default is "RECURRING '*/15 * * * *' FULL BACKUP '@hourly' WITH SCHEDULE OPTIONS
" the provider chosen for the cluster. If no volume type is provided the provider default will be used. "+
"Note: This volume will be deleted once the VM is deleted.")

updateCmd.Flags().BoolVar(&revertUpdate, "revert", false, "restore roachprod to the previous version "+
"which would have been renamed to roachprod.bak during the update process")

for _, cmd := range []*cobra.Command{createCmd, destroyCmd, extendCmd, logsCmd} {
cmd.Flags().StringVarP(&username, "username", "u", os.Getenv("ROACHPROD_USER"),
"Username to run under, detect if blank")
Expand Down
42 changes: 41 additions & 1 deletion pkg/cmd/roachprod/main.go
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,7 @@ import (

"github.com/cockroachdb/cockroach/pkg/build"
"github.com/cockroachdb/cockroach/pkg/cmd/roachprod/grafana"
"github.com/cockroachdb/cockroach/pkg/cmd/roachprod/upgrade"
"github.com/cockroachdb/cockroach/pkg/roachprod"
"github.com/cockroachdb/cockroach/pkg/roachprod/config"
rperrors "github.com/cockroachdb/cockroach/pkg/roachprod/errors"
Expand Down Expand Up @@ -1283,6 +1284,45 @@ func validateAndConfigure(cmd *cobra.Command, args []string) {
}
}

var updateCmd = &cobra.Command{
Use: "update",
Short: "check TeamCity for a new roachprod binary and update if available",
Long: "Will attempt to download the latest master branch roachprod binary from teamcity" +
" and swap the current roachprod with it. The current roachprod binary will be backed up" +
" and can be restored via `roachprod update --revert`.",
Run: wrap(func(cmd *cobra.Command, args []string) error {
currentBinary, err := os.Executable()
if err != nil {
return err
}

if revertUpdate {
if upgrade.PromptYesNo("Revert to previous version? Note: this will replace the" +
" current roachprod binary with a previous roachprod.bak binary.") {
if err := upgrade.SwapBinary(currentBinary, currentBinary+".bak"); err != nil {
return err
}
fmt.Println("roachprod successfully reverted, run `roachprod -v` to confirm.")
}
return nil
}

newBinary := currentBinary + ".new"
if err := upgrade.DownloadLatestRoadprod(newBinary); err != nil {
return err
}

if upgrade.PromptYesNo("Continue with update? This will overwrite any existing roachprod.bak binary.") {
if err := upgrade.SwapBinary(currentBinary, newBinary); err != nil {
return errors.WithDetail(err, "unable to update binary")
}

fmt.Println("Update successful: run `roachprod -v` to confirm.")
}
return nil
}),
}

func main() {
_ = roachprod.InitProviders()
providerOptsContainer = vm.CreateProviderOptionsContainer()
Expand All @@ -1299,7 +1339,6 @@ func main() {
syncCmd,
gcCmd,
setupSSHCmd,

statusCmd,
monitorCmd,
startCmd,
Expand Down Expand Up @@ -1333,6 +1372,7 @@ func main() {
rootStorageCmd,
snapshotCmd,
fixLongRunningAWSHostnamesCmd,
updateCmd,
)
setBashCompletionFunction()

Expand Down
37 changes: 37 additions & 0 deletions pkg/cmd/roachprod/upgrade/BUILD.bazel
Original file line number Diff line number Diff line change
@@ -0,0 +1,37 @@
load("@rules_proto//proto:defs.bzl", "proto_library")
load("//build/bazelutil/unused_checker:unused.bzl", "get_x_data")
load("@io_bazel_rules_go//go:def.bzl", "go_library")
load("@io_bazel_rules_go//proto:def.bzl", "go_proto_library")

go_library(
name = "upgrade",
srcs = [
"teamcity.go",
"util.go",
],
embed = [":upgrade_go_proto"],
importpath = "github.com/cockroachdb/cockroach/pkg/cmd/roachprod/upgrade",
visibility = ["//visibility:public"],
deps = [
"//pkg/util/httputil",
"@com_github_cockroachdb_errors//:errors",
"@com_github_cockroachdb_errors//oserror",
],
)

proto_library(
name = "upgrade_proto",
srcs = ["teamcity.proto"],
strip_import_prefix = "/pkg",
visibility = ["//visibility:public"],
)

go_proto_library(
name = "upgrade_go_proto",
compilers = ["//pkg/cmd/protoc-gen-gogoroach:protoc-gen-gogoroach_compiler"],
importpath = "github.com/cockroachdb/cockroach/pkg/cmd/roachprod/upgrade",
proto = ":upgrade_proto",
visibility = ["//visibility:public"],
)

get_x_data(name = "get_x_data")
118 changes: 118 additions & 0 deletions pkg/cmd/roachprod/upgrade/teamcity.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,118 @@
// Copyright 2023 The Cockroach Authors.
//
// Use of this software is governed by the Business Source License
// included in the file licenses/BSL.txt.
//
// As of the Change Date specified in that file, in accordance with
// the Business Source License, use of this software will be governed
// by the Apache License, Version 2.0, included in the file
// licenses/APL.txt.

package upgrade

import (
"context"
"fmt"
"io"
"net/http"
"os"
"runtime"
"time"

"github.com/cockroachdb/cockroach/pkg/util/httputil"
)

var (
buildIDs = map[string]string{
"linux-amd64": "Cockroach_Ci_Builds_BuildLinuxX8664",
"linux-arm64": "Cockroach_UnitTests_BazelBuildLinuxArmCross",
"darwin-amd64": "Cockroach_UnitTests_BazelBuildMacOSCross",
"darwin-arm64": "Cockroach_Ci_Builds_BuildMacOSArm64",
"windows-amd64": "Cockroach_UnitTests_BazelBuildWindowsCross",
}
apiBase = "https://teamcity.cockroachdb.com/guestAuth/app/rest"
)

// DownloadLatestRoadprod attempts to download the latest binary to the
// current binary's directory. It returns the path to the downloaded binary.
// toFile is the path to the file to download to.
func DownloadLatestRoadprod(toFile string) error {
buildType, ok := buildIDs[runtime.GOOS+"-"+runtime.GOARCH]
if !ok {
fmt.Println("Supported platforms:")
for k := range buildIDs {
fmt.Printf("\t%s\n", k)
}
return fmt.Errorf("unable to find build type for this platform")
}

// Build are sorted by build date desc, so limiting to 1 will get the latest.
builds, err := getBuilds("count:1,status:SUCCESS,branch:master,buildType:" + buildType)
if err != nil {
return err
}

if len(builds.Build) == 0 {
return fmt.Errorf("no builds found")
}

out, err := os.Create(toFile)
if err != nil {
return err
}

defer out.Close()
err = downloadRoachprod(builds.Build[0].Id, out)
if err != nil {
return err
}
fmt.Printf("Downloaded latest roachprod to:\t%s\n", toFile)
return nil
}

// getBuilds returns a list of builds matching the locator
// See https://www.jetbrains.com/help/teamcity/rest/buildlocator.html
func getBuilds(locator string) (TCBuildResponse, error) {
urlWithLocator := fmt.Sprintf("%s/builds?locator=%s", apiBase, locator)
buildResp := &TCBuildResponse{}
err := httputil.GetJSONWithOptions(*httputil.DefaultClient.Client, urlWithLocator, buildResp,
httputil.IgnoreUnknownFields())
return *buildResp, err
}

// downloadRoachprod downloads the roachprod binary from the build
// using the specified writer. It is the caller's responsibility to close the writer.
func downloadRoachprod(buildID int32, destWriter io.Writer) error {
if buildID <= 0 {
return fmt.Errorf("invalid build id: %v", buildID)
}

url := roachprodDownloadURL(buildID)
fmt.Printf("Downloading roachprod from:\t%s\n", url)

// Set a long timeout here because the download can take a while.
httpClient := httputil.NewClientWithTimeouts(httputil.StandardHTTPTimeout, 10*time.Minute)
resp, err := httpClient.Get(context.Background(), url)
if err != nil {
return err
}

defer resp.Body.Close()
if resp.StatusCode != http.StatusOK {
return fmt.Errorf("bad status downloading roachprod: %s", resp.Status)
}

_, err = io.Copy(destWriter, resp.Body)
if err != nil {
return err
}

return nil
}

func roachprodDownloadURL(buildID int32) string {
url := fmt.Sprintf("%s%s",
apiBase,
fmt.Sprintf("/builds/id:%v/artifacts/content/bazel-bin/pkg/cmd/roachprod/roachprod_/roachprod", buildID))
return url
}
25 changes: 25 additions & 0 deletions pkg/cmd/roachprod/upgrade/teamcity.proto
Original file line number Diff line number Diff line change
@@ -0,0 +1,25 @@
// Copyright 2023 The Cockroach Authors.
//
// Use of this software is governed by the Business Source License
// included in the file licenses/BSL.txt.
//
// As of the Change Date specified in that file, in accordance with
// the Business Source License, use of this software will be governed
// by the Apache License, Version 2.0, included in the file
// licenses/APL.txt.

syntax = "proto3";

package upgrade;

message TCBuildResponse {
int32 count = 1;
repeated TCBuild build = 2;
}

message TCBuild {
int32 id = 1;
string webUrl = 2;
string branchName = 3;
string finishOnAgentDate = 4;
}
72 changes: 72 additions & 0 deletions pkg/cmd/roachprod/upgrade/util.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,72 @@
// Copyright 2023 The Cockroach Authors.
//
// Use of this software is governed by the Business Source License
// included in the file licenses/BSL.txt.
//
// As of the Change Date specified in that file, in accordance with
// the Business Source License, use of this software will be governed
// by the Apache License, Version 2.0, included in the file
// licenses/APL.txt.

package upgrade

import (
"fmt"
"os"
"strings"

"github.com/cockroachdb/errors"
"github.com/cockroachdb/errors/oserror"
)

func PromptYesNo(msg string) bool {
fmt.Printf("%s y[default]/n: ", msg)
var answer string
_, _ = fmt.Scanln(&answer)
answer = strings.TrimSpace(answer)

return answer == "y" || answer == "Y" || answer == ""
}

// SwapBinary attempts to swap the `old` file with the `new` file. Used to
// update a running roachprod binary.
// Note: there is special handling if `new` points to a file ending in `.bak`.
// In this case, it is assumed to be a `revert` operation, in which case we
// do *not* backup the old/current file.
func SwapBinary(old, new string) error {
destInfo, err := os.Stat(new)

if err != nil {
if oserror.IsNotExist(err) {
return errors.WithDetail(err, "binary does not exist: "+new)
}
return err
}

if destInfo.IsDir() {
return errors.Newf("binary path is a directory, not a file: %s", new)
}

oldInfo, err := os.Stat(old)
if err != nil {
return err
}

// Copy the current file permissions to the new binary and ensure it is executable.
err = os.Chmod(new, oldInfo.Mode())
if err != nil {
return err
}

// Backup only for upgrading, not when reverting which is assumed if the new binary ends in `.bak`.
if !strings.HasSuffix(new, ".bak") {
// Backup the current binary, so that it may be restored via `roachprod update --revert`.
err = os.Rename(old, old+".bak")
if err != nil {
return errors.WithDetail(err, "unable to backup current binary")
}
}

// Move the new binary into place.
return os.Rename(new, old)
}
1 change: 1 addition & 0 deletions pkg/gen/protobuf.bzl
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@ PROTOBUF_SRCS = [
"//pkg/cloud/cloudpb:cloudpb_go_proto",
"//pkg/cloud/externalconn/connectionpb:connectionpb_go_proto",
"//pkg/clusterversion:clusterversion_go_proto",
"//pkg/cmd/roachprod/upgrade:upgrade_go_proto",
"//pkg/config/zonepb:zonepb_go_proto",
"//pkg/config:config_go_proto",
"//pkg/geo/geoindex:geoindex_go_proto",
Expand Down
9 changes: 7 additions & 2 deletions pkg/util/httputil/client.go
Original file line number Diff line number Diff line change
Expand Up @@ -27,12 +27,17 @@ const StandardHTTPTimeout time.Duration = 3 * time.Second

// NewClientWithTimeout defines a http.Client with the given timeout.
func NewClientWithTimeout(timeout time.Duration) *Client {
return NewClientWithTimeouts(timeout, timeout)
}

// NewClientWithTimeouts defines a http.Client with the given dialer and client timeouts.
func NewClientWithTimeouts(dialerTimeout, clientTimeout time.Duration) *Client {
return &Client{&http.Client{
Timeout: timeout,
Timeout: clientTimeout,
Transport: &http.Transport{
// Don't leak a goroutine on OSX (the TCP level timeout is probably
// much higher than on linux).
DialContext: (&net.Dialer{Timeout: timeout}).DialContext,
DialContext: (&net.Dialer{Timeout: dialerTimeout}).DialContext,
DisableKeepAlives: true,
},
}}
Expand Down
Loading

0 comments on commit bf25b7b

Please sign in to comment.