From 351e08d1360b6a1b2a05ad460a2ce6fd342e92cf Mon Sep 17 00:00:00 2001 From: Ricky Stewart Date: Wed, 11 Dec 2024 16:25:05 -0600 Subject: [PATCH] teamcity: add script to run PGO build and collect profiles This script ties together the pieces of the workflow: start a TC job, wait for it to finish, downloads the artifacts, parses the profiles, merges them, and produces a final result. Part of: CRDB-44692 Epic: CRDB-41952 Release note: None --- pkg/BUILD.bazel | 2 + pkg/cmd/run-pgo-build/BUILD.bazel | 15 ++ pkg/cmd/run-pgo-build/main.go | 267 ++++++++++++++++++++++++++++++ 3 files changed, 284 insertions(+) create mode 100644 pkg/cmd/run-pgo-build/BUILD.bazel create mode 100644 pkg/cmd/run-pgo-build/main.go diff --git a/pkg/BUILD.bazel b/pkg/BUILD.bazel index f9d6aff33541..17e3ea1eaf08 100644 --- a/pkg/BUILD.bazel +++ b/pkg/BUILD.bazel @@ -1220,6 +1220,8 @@ GO_TARGETS = [ "//pkg/cmd/roachtest:roachtest_test", "//pkg/cmd/roachvet:roachvet", "//pkg/cmd/roachvet:roachvet_lib", + "//pkg/cmd/run-pgo-build:run-pgo-build", + "//pkg/cmd/run-pgo-build:run-pgo-build_lib", "//pkg/cmd/skip-test:skip-test", "//pkg/cmd/skip-test:skip-test_lib", "//pkg/cmd/skiperrs:skiperrs", diff --git a/pkg/cmd/run-pgo-build/BUILD.bazel b/pkg/cmd/run-pgo-build/BUILD.bazel new file mode 100644 index 000000000000..bcbf6d798f4d --- /dev/null +++ b/pkg/cmd/run-pgo-build/BUILD.bazel @@ -0,0 +1,15 @@ +load("@io_bazel_rules_go//go:def.bzl", "go_binary", "go_library") + +go_library( + name = "run-pgo-build_lib", + srcs = ["main.go"], + importpath = "github.com/cockroachdb/cockroach/pkg/cmd/run-pgo-build", + visibility = ["//visibility:private"], + deps = ["@com_github_google_pprof//profile"], +) + +go_binary( + name = "run-pgo-build", + embed = [":run-pgo-build_lib"], + visibility = ["//visibility:public"], +) diff --git a/pkg/cmd/run-pgo-build/main.go b/pkg/cmd/run-pgo-build/main.go new file mode 100644 index 000000000000..013edab14a08 --- /dev/null +++ b/pkg/cmd/run-pgo-build/main.go @@ -0,0 +1,267 @@ +// Copyright 2024 The Cockroach Authors. +// +// Use of this software is governed by the CockroachDB Software License +// included in the /LICENSE file. +// +// run-pgo-build triggers a run of a special roachtest job in TeamCity on the +// current branch, waits for the job to complete, downloads the associated +// artifacts, extracts all the CPU (.pprof) profiles, merges them, and places +// the merged result in a location on disk. +package main + +import ( + "archive/zip" + "bytes" + "encoding/json" + "flag" + "fmt" + "io" + "net/http" + "os" + "path/filepath" + "strings" + "sync" + "time" + + "github.com/google/pprof/profile" +) + +const ( + buildConfigID = "Cockroach_Nightlies_RoachtestGceForceProfile" +) + +var ( + branch = os.Getenv("TC_BUILD_BRANCH") + serverURL = os.Getenv("TC_SERVER_URL") + username = os.Getenv("TC_API_USER") + password = os.Getenv("TC_API_PASSWORD") + + outFile = flag.String("out", "", "where to store the result pprof profile") +) + +type Build struct { + ID int64 + FinishDate string + Status string +} + +type ReadAtCloser interface { + io.ReaderAt + io.ReadCloser +} + +func doRequest(req *http.Request) (*Build, error) { + req.SetBasicAuth(username, password) + req.Header.Add("Accept", "application/json") + resp, err := http.DefaultClient.Do(req) + if err != nil { + return nil, err + } + var buf bytes.Buffer + if _, err := io.Copy(&buf, resp.Body); err != nil { + return nil, err + } + var ret Build + if err := json.Unmarshal(buf.Bytes(), &ret); err != nil { + return nil, err + } + return &ret, nil +} + +func queueBuild(buildConfigID, branch string) (*Build, error) { + reqStruct := struct { + BuildTypeID string `json:"buildTypeId,omitempty"` + BranchName string `json:"branchName,omitempty"` + }{ + BuildTypeID: buildConfigID, + BranchName: branch, + } + reqJson, err := json.Marshal(reqStruct) + if err != nil { + return nil, err + } + req, err := http.NewRequest("POST", fmt.Sprintf("https://%s/httpAuth/app/rest/buildQueue", serverURL), bytes.NewReader(reqJson)) + if err != nil { + return nil, err + } + req.Header.Add("Content-Type", "application/json") + return doRequest(req) +} + +func getBuild(id int64) (*Build, error) { + req, err := http.NewRequest("GET", fmt.Sprintf("https://%s/httpAuth/app/rest/builds/id:%d?fields=id,finishDate,status", serverURL, id), nil) + if err != nil { + return nil, err + } + return doRequest(req) +} + +// downloadArtifacts downloads the artifacts.zip archive for the given build and +// returns the ReadAtCloser corresponding to it. The file will be stored on +// disk in the given tmpDir. The caller is responsible for closing the file. +// The returned int64 is the length of the file. +func downloadArtifacts(buildID int64, tmpDir string) (ReadAtCloser, int64, error) { + req, err := http.NewRequest("GET", fmt.Sprintf("https://%s/repository/downloadAll/%s/%d:id", serverURL, buildConfigID, buildID), nil) + if err != nil { + return nil, 0, err + } + req.SetBasicAuth(username, password) + resp, err := http.DefaultClient.Do(req) + if err != nil { + return nil, 0, err + } + defer resp.Body.Close() + + zipFile, err := os.Create(filepath.Join(tmpDir, "artifacts.zip")) + if err != nil { + return nil, 0, err + } + zipLength, err := io.Copy(zipFile, resp.Body) + if err != nil { + zipFile.Close() + return nil, 0, err + } + _, err = zipFile.Seek(0, 0) + if err != nil { + zipFile.Close() + return nil, 0, err + } + + return zipFile, zipLength, nil +} + +// processArtifactsZip processes a file named artifacts.zip inside the larger +// artifacts zip archive. This is confusing, but the `.pprof` files for any +// given node are inside a sub-zip archive named `artifacts.zip`, which is one +// of the files in the giant zip archive that we fetch for the build. This +// function processes that sub-zip archive and extracts all the CPU (.pprof) +// files. +// +// The .pprof files will be parsed and put into the pprofFilesChan. +// wg.Done() will be called when this function completes. +func processArtifactsZip(f *zip.File, profilesChan chan *profile.Profile, wg *sync.WaitGroup) { + archive, err := f.Open() + if err != nil { + panic(err) + } + defer archive.Close() + var archiveBuf bytes.Buffer + _, err = io.Copy(&archiveBuf, archive) + zipReader, err := zip.NewReader(bytes.NewReader(archiveBuf.Bytes()), int64(f.UncompressedSize64)) + if err != nil { + panic(err) + } + for _, file := range zipReader.File { + if strings.HasSuffix(file.FileHeader.Name, ".pprof") && + strings.Contains(file.FileHeader.Name, "/cpuprof.") && + file.UncompressedSize64 > 0 { + pprofFile, err := file.Open() + if err != nil { + panic(err) + } + defer pprofFile.Close() + prof, err := profile.Parse(pprofFile) + if err != nil { + panic(err) + } + profilesChan <- prof + } + } + wg.Done() +} + +// processPprofFiles reads all the parsed profiles from the given channel, +// merges all of the profiles, and dumps the results profile to a final +// location. wg.Done() will be called at the end of this function. +func processPprofFiles(profilesChan chan *profile.Profile, wg *sync.WaitGroup) { + var profiles []*profile.Profile + for prof := range profilesChan { + profiles = append(profiles, prof) + } + if len(profiles) == 0 { + panic("expected to find a profile; found none") + } + res, err := profile.Merge(profiles) + if err != nil { + panic(err) + } + w, err := os.Create(*outFile) + if err != nil { + panic(err) + } + defer w.Close() + err = res.Write(w) + if err != nil { + panic(err) + } + wg.Done() +} + +func main() { + flag.Parse() + + if branch == "" || serverURL == "" || username == "" || password == "" { + panic("ensure credentials, server URL, and branch are supplied") + } + if *outFile == "" { + panic("must supply -out") + } + + tmpDir, err := os.MkdirTemp("", "run-pgo-build") + if err != nil { + panic(err) + } + defer func() { + _ = os.RemoveAll(tmpDir) + }() + + build, err := queueBuild(buildConfigID, branch) + if err != nil { + panic(err) + } + fmt.Printf("queued build with ID %d", build.ID) + + for { + if build.FinishDate != "" { + break + } + time.Sleep(3 * time.Minute) + lookupID := build.ID + build, err = getBuild(lookupID) + if err != nil { + fmt.Printf("failed to get build %d; got error %+v\n", lookupID, err) + } + } + + if build.Status != "SUCCESS" { + panic(fmt.Sprintf("expected build to succeed; got status %s for build %+v", build.Status, build)) + } + + zipFile, zipLength, err := downloadArtifacts(build.ID, tmpDir) + if err != nil { + panic(err) + } + defer zipFile.Close() + + zipReader, err := zip.NewReader(zipFile, zipLength) + if err != nil { + panic(err) + } + + // wg tracks the progress of loading the profiles, readWg tracks the + // progress of reading them. + var wg, readWg sync.WaitGroup + // All parsed profiles will be sent to profilesChan. + profilesChan := make(chan *profile.Profile) + for _, file := range zipReader.File { + if strings.HasSuffix(file.FileHeader.Name, "/artifacts.zip") { + wg.Add(1) + go processArtifactsZip(file, profilesChan, &wg) + } + } + readWg.Add(1) + go processPprofFiles(profilesChan, &readWg) + wg.Wait() + close(profilesChan) + readWg.Wait() +}