[sk] Update release-branch command for SkCQ and Chrome Faster

Bug: skia:12271
Change-Id: If38e43e44e6613ba2cdb044d70430949ccf86658
Reviewed-on: https://skia-review.googlesource.com/c/buildbot/+/434176
Commit-Queue: Eric Boren <borenet@google.com>
Reviewed-by: Ravi Mistry <rmistry@google.com>
Reviewed-by: Heather Miller <hcm@google.com>
diff --git a/go/gerrit/change_edit_helpers.go b/go/gerrit/change_edit_helpers.go
index bd7434d..6285407 100644
--- a/go/gerrit/change_edit_helpers.go
+++ b/go/gerrit/change_edit_helpers.go
@@ -2,6 +2,9 @@
 
 import (
 	"context"
+	"fmt"
+	"io/ioutil"
+	"path/filepath"
 	"strings"
 	"time"
 
@@ -86,13 +89,22 @@
 
 // CreateCLWithChanges is a helper which creates a new Change in the given
 // project based on the given branch with the given commit message and the given
-// map of filepath to new file contents. If submit is true, the change is marked
-// with the self-approval label(s) and submitted.
+// map of filepath to new file contents. Empty file contents indicate deletion
+// of the file. If submit is true, the change is marked with the self-approval
+// label(s) and submitted.
 func CreateCLWithChanges(ctx context.Context, g GerritInterface, project, branch, commitMsg, baseCommit string, changes map[string]string, submit bool) (*ChangeInfo, error) {
 	ci, err := CreateAndEditChange(ctx, g, project, branch, commitMsg, baseCommit, func(ctx context.Context, g GerritInterface, ci *ChangeInfo) error {
 		for filepath, contents := range changes {
-			if err := g.EditFile(ctx, ci, filepath, contents); err != nil {
-				return skerr.Wrap(err)
+			if contents == "" {
+				if err := g.DeleteFile(ctx, ci, filepath); err != nil {
+					return skerr.Wrap(err)
+				}
+			} else {
+				if err := g.EditFile(ctx, ci, filepath, contents); err != nil {
+					if !strings.Contains(err.Error(), ErrNoChanges) {
+						return skerr.Wrap(err)
+					}
+				}
 			}
 		}
 		return nil
@@ -115,3 +127,30 @@
 	}
 	return ci, nil
 }
+
+// CreateCLFromLocalDiffs is a helper which creates a Change based on the
+// diffs from the provided branch in a local checkout. Note that the diff is
+// performed against the given branch on "origin", and not any local version.
+func CreateCLFromLocalDiffs(ctx context.Context, g GerritInterface, project, branch, commitMsg string, submit bool, co *git.Checkout) (*ChangeInfo, error) {
+	baseCommit, err := co.Git(ctx, "rev-parse", fmt.Sprintf("origin/%s", branch))
+	if err != nil {
+		return nil, skerr.Wrap(err)
+	}
+	baseCommit = strings.TrimSpace(baseCommit)
+	diff, err := co.Git(ctx, "diff", "--name-only", baseCommit)
+	if err != nil {
+		return nil, skerr.Wrap(err)
+	}
+	diffSplit := strings.Split(diff, "\n")
+	changes := make(map[string]string, len(diffSplit))
+	for _, diffLine := range diffSplit {
+		if diffLine != "" {
+			contents, err := ioutil.ReadFile(filepath.Join(co.Dir(), diffLine))
+			if err != nil {
+				return nil, skerr.Wrap(err)
+			}
+			changes[diffLine] = string(contents)
+		}
+	}
+	return CreateCLWithChanges(ctx, g, project, branch, commitMsg, baseCommit, changes, submit)
+}
diff --git a/go/gerrit/gerrit.go b/go/gerrit/gerrit.go
index 09e3676..aff9f13 100644
--- a/go/gerrit/gerrit.go
+++ b/go/gerrit/gerrit.go
@@ -172,6 +172,10 @@
 	// merge conflict occurred.
 	ErrMergeConflict = "conflict during merge"
 
+	// ErrNoChanges as a substring of an error message indicates that there were
+	// no changes to apply. Generally we can ignore this error.
+	ErrNoChanges = "no changes were made"
+
 	// These were copied from the defaults used by gitfs:
 	// https://gerrit.googlesource.com/gitfs/+show/59c1163fd1737445281f2339399b2b986b0d30fe/gitiles/client.go#102
 	// Hopefully they apply to Gerrit as well.
diff --git a/sk/go/release-branch/BUILD.bazel b/sk/go/release-branch/BUILD.bazel
index f4f6b03..1b06435 100644
--- a/sk/go/release-branch/BUILD.bazel
+++ b/sk/go/release-branch/BUILD.bazel
@@ -8,12 +8,16 @@
     deps = [
         "//go/auth",
         "//go/common",
+        "//go/cq",
+        "//go/exec",
         "//go/gerrit",
         "//go/git",
         "//go/gitiles",
         "//go/httputils",
         "//go/skerr",
-        "//go/supported_branches/cmd/new-branch/helper",
+        "//go/supported_branches",
+        "//go/util",
+        "//task_scheduler/go/specs",
         "@com_github_urfave_cli_v2//:cli",
     ],
 )
diff --git a/sk/go/release-branch/release-branch.go b/sk/go/release-branch/release-branch.go
index d0a56d6..77af9e1 100644
--- a/sk/go/release-branch/release-branch.go
+++ b/sk/go/release-branch/release-branch.go
@@ -1,45 +1,57 @@
 package try
 
 import (
+	"bytes"
 	"context"
+	"encoding/json"
 	"fmt"
 	"io/ioutil"
 	"os"
+	"path/filepath"
 	"regexp"
 	"strconv"
+	"strings"
 
 	"github.com/urfave/cli/v2"
 
 	"go.skia.org/infra/go/auth"
 	"go.skia.org/infra/go/common"
+	"go.skia.org/infra/go/cq"
+	"go.skia.org/infra/go/exec"
 	"go.skia.org/infra/go/gerrit"
 	"go.skia.org/infra/go/git"
 	"go.skia.org/infra/go/gitiles"
 	"go.skia.org/infra/go/httputils"
 	"go.skia.org/infra/go/skerr"
-	"go.skia.org/infra/go/supported_branches/cmd/new-branch/helper"
+	"go.skia.org/infra/go/supported_branches"
+	"go.skia.org/infra/go/util"
+	"go.skia.org/infra/task_scheduler/go/specs"
 )
 
 const (
-	chromeBranchTmpl             = "chrome/m%d"
+	chromeBranchTmpl             = "chrome/m%s"
 	flagDryRun                   = "dry-run"
 	gerritProject                = "skia"
 	milestoneFile                = "include/core/SkMilestone.h"
 	milestoneTmpl                = "#define SK_MILESTONE %s"
 	supportedChromeBranches      = 4
 	updateMilestoneCommitMsgTmpl = "Update Skia milestone to %d"
+	jobsJSONFile                 = "infra/bots/jobs.json"
+	tasksJSONFile                = "infra/bots/tasks.json"
+	cqJSONFile                   = "infra/skcq.json"
 )
 
 var (
-	milestoneRegex = regexp.MustCompile(fmt.Sprintf(milestoneTmpl, `(\d+)`))
+	chromeBranchMilestoneRegex = regexp.MustCompile(fmt.Sprintf(chromeBranchTmpl, `(\d+)`))
+	milestoneRegex             = regexp.MustCompile(fmt.Sprintf(milestoneTmpl, `(\d+)`))
 
-	excludeTrybotsOnReleaseBranches = []string{
-		"chromium.*",
-		".*Android_Framework.*",
-		".*G3_Framework.*",
-		".*CanvasKit.*",
-		".*PathKit.*",
+	excludeTrybotsOnReleaseBranches = []*regexp.Regexp{
+		regexp.MustCompile(".*CanvasKit.*"),
+		regexp.MustCompile(".*PathKit.*"),
 	}
+
+	jobsJSONReplaceRegex    = regexp.MustCompile(`(?m)\{\n    "name": "(\S+)",\n    "cq_config": null\n  }`)
+	jobsJSONReplaceContents = []byte(`{"name": "$1"}`)
 )
 
 // Command returns a cli.Command instance which represents the "release-branch"
@@ -47,13 +59,13 @@
 func Command() *cli.Command {
 	return &cli.Command{
 		Name:        "release-branch",
-		Usage:       "release-branch <commit hash>",
-		Description: "Create a new Skia release branch at the given hash.",
+		Usage:       "release-branch [options] <newly-created branch name>",
+		Description: "Perform the necessary updates after creating a new release branch. The new branch must already exist.",
 		Flags: []cli.Flag{
 			&cli.BoolFlag{
 				Name:  flagDryRun,
 				Value: false,
-				Usage: "Create the branch CLs but do not submit them or create the branch itself.",
+				Usage: "Create the CLs but do not submit them.",
 			},
 		},
 		Action: func(ctx *cli.Context) error {
@@ -68,7 +80,20 @@
 
 // releaseBranch performs the actions necessary to create a new Skia release
 // branch.
-func releaseBranch(ctx context.Context, commit string, dryRun bool) error {
+func releaseBranch(ctx context.Context, newBranch string, dryRun bool) error {
+	// Derive the new milestone number and the newly-expired branch name from
+	// the new branch name.
+	m := chromeBranchMilestoneRegex.FindStringSubmatch(newBranch)
+	if len(m) != 2 {
+		return skerr.Fmt("Invalid branch name %q; expected format: %s", newBranch, chromeBranchMilestoneRegex.String())
+	}
+	currentMilestone, err := strconv.Atoi(m[1])
+	if err != nil {
+		return skerr.Wrap(err)
+	}
+	newMilestone := currentMilestone + 1
+	removeBranch := fmt.Sprintf(chromeBranchTmpl, strconv.Itoa(currentMilestone-supportedChromeBranches))
+
 	// Setup.
 	ts, err := auth.NewDefaultTokenSource(true, auth.SCOPE_GERRIT)
 	if err != nil {
@@ -82,28 +107,34 @@
 	}
 
 	// Retrieve the current milestone.
-	baseCommit, err := repo.ResolveRef(ctx, git.DefaultRef)
+	baseCommit, err := repo.ResolveRef(ctx, git.MainBranch)
 	if err != nil {
 		return skerr.Wrap(err)
 	}
-	milestone, milestoneFileContents, err := getCurrentMilestone(ctx, repo, baseCommit)
+	haveMilestone, milestoneFileContents, err := getCurrentMilestone(ctx, repo, baseCommit)
 	if err != nil {
 		return skerr.Wrap(err)
 	}
-	newBranch := fmt.Sprintf(chromeBranchTmpl, milestone)
-	oldBranch := fmt.Sprintf(chromeBranchTmpl, milestone-supportedChromeBranches)
-	fmt.Println(fmt.Sprintf("Creating branch %s and removing support (eg. CQ) for %s", newBranch, oldBranch))
+	fmt.Println(fmt.Sprintf("Adding support for branch %s and removing support for %s", newBranch, removeBranch))
 
-	fmt.Println(fmt.Sprintf("Creating branch %s...", newBranch))
-	if err := createNewBranch(ctx, newBranch, commit, dryRun); err != nil {
+	if haveMilestone == newMilestone {
+		fmt.Println(fmt.Sprintf("Milestone is up to date at %d; not updating.", newMilestone))
+	} else {
+		fmt.Println(fmt.Sprintf("Creating CL to update milestone from %d to %d...", haveMilestone, newMilestone))
+		if err := updateMilestone(ctx, g, baseCommit, newMilestone, milestoneFileContents, dryRun); err != nil {
+			return skerr.Wrap(err)
+		}
+	}
+	fmt.Println("Creating CL to update supported branches...")
+	if err := updateSupportedBranches(ctx, g, repo, removeBranch, newBranch, dryRun); err != nil {
 		return skerr.Wrap(err)
 	}
-	fmt.Println("Creating CL to update milestone...")
-	if err := updateMilestone(ctx, g, baseCommit, milestone+1, milestoneFileContents, dryRun); err != nil {
+	fmt.Println(fmt.Sprintf("Creating CL to filter out unsupported CQ try jobs on %s...", newBranch))
+	if err := updateTryjobs(ctx, g, repo, newBranch, dryRun); err != nil {
 		return skerr.Wrap(err)
 	}
-	fmt.Println("Creating CL to update CQ...")
-	if err := updateInfraConfig(ctx, g, oldBranch, newBranch, dryRun); err != nil {
+	fmt.Println(fmt.Sprintf("Creating CL to remove CQ on %s", removeBranch))
+	if err := removeCQ(ctx, g, repo, removeBranch, dryRun); err != nil {
 		return skerr.Wrap(err)
 	}
 	return nil
@@ -129,44 +160,6 @@
 	return milestone, string(contents), nil
 }
 
-// createNewBranch creates a branch with the given name at the given commit.
-func createNewBranch(ctx context.Context, name, commit string, dryRun bool) (rvErr error) {
-	// Create a temporary checkout of Skia.
-	tmp, err := ioutil.TempDir("", "")
-	if err != nil {
-		return skerr.Wrap(err)
-	}
-	if !dryRun {
-		defer func() {
-			if err := os.RemoveAll(tmp); err != nil {
-				if rvErr == nil {
-					rvErr = err
-				}
-			}
-		}()
-	}
-	co, err := git.NewCheckout(ctx, common.REPO_SKIA, tmp)
-	if err != nil {
-		return skerr.Wrap(err)
-	}
-
-	// Create and push the branch.
-	if _, err := co.Git(ctx, "checkout", "-b", name); err != nil {
-		return skerr.Wrap(err)
-	}
-	if _, err := co.Git(ctx, "reset", "--hard", commit); err != nil {
-		return skerr.Wrap(err)
-	}
-	if dryRun {
-		fmt.Fprintf(os.Stderr, "Branch %s created in %s; not pushing because --dry-run was specified.\n", name, co.Dir())
-	} else {
-		if _, err := co.Git(ctx, "push", "--set-upstream", "origin", name); err != nil {
-			return skerr.Wrap(err)
-		}
-	}
-	return nil
-}
-
 // updateMilestone creates a CL to update the Skia milestone.
 func updateMilestone(ctx context.Context, g gerrit.GerritInterface, baseCommit string, milestone int, oldMilestoneContents string, dryRun bool) error {
 	commitMsg := fmt.Sprintf(updateMilestoneCommitMsgTmpl, milestone)
@@ -181,12 +174,175 @@
 	return skerr.Wrap(err)
 }
 
-// updateInfraConfig updates the infra/config branch to edit the supported
+// updateSupportedBranches updates the infra/config branch to edit the supported
 // branches and commit queue config to add the new branch and remove the old.
-func updateInfraConfig(ctx context.Context, g gerrit.GerritInterface, oldBranch, newBranch string, dryRun bool) error {
+func updateSupportedBranches(ctx context.Context, g gerrit.GerritInterface, repo *gitiles.Repo, removeBranch string, newBranch string, dryRun bool) error {
 	owner, err := g.GetUserEmail(ctx)
 	if err != nil {
 		return skerr.Wrap(err)
 	}
-	return skerr.Wrap(helper.AddSupportedBranch(common.REPO_SKIA, newBranch, owner, oldBranch, excludeTrybotsOnReleaseBranches, !dryRun))
+	newRef := git.FullyQualifiedBranchName(newBranch)
+
+	// Setup.
+	baseCommitInfo, err := repo.Details(ctx, cq.CQ_CFG_REF)
+	if err != nil {
+		return skerr.Wrap(err)
+	}
+	baseCommit := baseCommitInfo.Hash
+
+	// Download and modify the supported-branches.json file.
+	branchesContents, err := repo.ReadFileAtRef(ctx, supported_branches.SUPPORTED_BRANCHES_FILE, baseCommit)
+	if err != nil {
+		return skerr.Wrap(err)
+	}
+	sbc, err := supported_branches.DecodeConfig(bytes.NewReader(branchesContents))
+	if err != nil {
+		return skerr.Wrap(err)
+	}
+	deleteRef := git.FullyQualifiedBranchName(removeBranch)
+	foundNewRef := false
+	deletedRef := false
+	newBranches := make([]*supported_branches.SupportedBranch, 0, len(sbc.Branches)+1)
+	for _, sb := range sbc.Branches {
+		if sb.Ref != deleteRef {
+			newBranches = append(newBranches, sb)
+		} else {
+			deletedRef = true
+		}
+		if sb.Ref == newRef {
+			foundNewRef = true
+		}
+	}
+	if foundNewRef {
+		_, _ = fmt.Fprintf(os.Stderr, "Already have %s in %s; not adding a duplicate.\n", newRef, supported_branches.SUPPORTED_BRANCHES_FILE)
+	} else {
+		newBranches = append(newBranches, &supported_branches.SupportedBranch{
+			Ref:   newRef,
+			Owner: owner,
+		})
+	}
+	if !deletedRef {
+		_, _ = fmt.Fprintf(os.Stderr, "%s not found in %s; not removing.\n", deleteRef, supported_branches.SUPPORTED_BRANCHES_FILE)
+	}
+	if !deletedRef && foundNewRef {
+		// We didn't change anything, so don't upload a CL.
+		return nil
+	}
+	sbc.Branches = newBranches
+	buf := bytes.Buffer{}
+	if err := supported_branches.EncodeConfig(&buf, sbc); err != nil {
+		return skerr.Wrap(err)
+	}
+
+	// Create the Gerrit CL.
+	commitMsg := fmt.Sprintf("Add supported branch %s, remove %s", newBranch, removeBranch)
+	repoSplit := strings.Split(repo.URL(), "/")
+	project := strings.TrimSuffix(repoSplit[len(repoSplit)-1], ".git")
+	changes := map[string]string{
+		supported_branches.SUPPORTED_BRANCHES_FILE: buf.String(),
+	}
+	ci, err := gerrit.CreateCLWithChanges(ctx, g, project, cq.CQ_CFG_REF, commitMsg, baseCommit, changes, !dryRun)
+	if ci != nil {
+		fmt.Println(fmt.Sprintf("Uploaded change %s", g.Url(ci.Issue)))
+	}
+	return skerr.Wrap(err)
+}
+
+func updateTryjobs(ctx context.Context, g gerrit.GerritInterface, repo *gitiles.Repo, newBranch string, dryRun bool) error {
+	// Setup.
+	newRef := git.FullyQualifiedBranchName(newBranch)
+	baseCommitInfo, err := repo.Details(ctx, newRef)
+	if err != nil {
+		return skerr.Wrap(err)
+	}
+	baseCommit := baseCommitInfo.Hash
+	tmp, err := ioutil.TempDir("", "")
+	if err != nil {
+		return skerr.Wrap(err)
+	}
+	defer util.RemoveAll(tmp)
+	co, err := git.NewCheckout(ctx, repo.URL(), tmp)
+	if err != nil {
+		return skerr.Wrap(err)
+	}
+	if err := co.CleanupBranch(ctx, newBranch); err != nil {
+		return skerr.Wrap(err)
+	}
+
+	// Download and modify the jobs.json file.
+	oldJobsContents, err := repo.ReadFileAtRef(ctx, jobsJSONFile, baseCommit)
+	if err != nil {
+		return skerr.Wrap(err)
+	}
+	var jobs []struct {
+		Name     string                      `json:"name"`
+		CqConfig *specs.CommitQueueJobConfig `json:"cq_config"`
+	}
+	if err := json.Unmarshal(oldJobsContents, &jobs); err != nil {
+		return skerr.Wrapf(err, "failed to decode jobs.json")
+	}
+	for _, job := range jobs {
+		for _, re := range excludeTrybotsOnReleaseBranches {
+			if re.MatchString(job.Name) {
+				job.CqConfig = nil
+				break
+			}
+		}
+	}
+	newJobsContents, err := json.MarshalIndent(jobs, "", "  ")
+	if err != nil {
+		return skerr.Wrapf(err, "failed to encode jobs.json")
+	}
+	// Replace instances of `"cq_config": null`; these are cluttery and
+	// unnecessary, but we can't use omitempty because that causes the Marshaler
+	// to also omit `"cq_config": {}` which indicates that a job *should* be on
+	// the CQ. Also attempt to match the whitespace of the original files, to
+	// help prevent conflicts during cherry-picks.
+	newJobsContents = jobsJSONReplaceRegex.ReplaceAll(newJobsContents, jobsJSONReplaceContents)
+	if err := ioutil.WriteFile(filepath.Join(co.Dir(), jobsJSONFile), newJobsContents, os.ModePerm); err != nil {
+		return skerr.Wrapf(err, "failed to write %s", jobsJSONFile)
+	}
+
+	// Regenerate tasks.json.
+	if _, err := exec.RunCwd(ctx, co.Dir(), "go", "run", "./infra/bots/gen_tasks.go"); err != nil {
+		return skerr.Wrapf(err, "failed to regenerate tasks.json")
+	}
+
+	// Create the Gerrit CL.
+	commitMsg := fmt.Sprintf("Filter unsupported CQ try jobs on %s", newBranch)
+	repoSplit := strings.Split(repo.URL(), "/")
+	project := strings.TrimSuffix(repoSplit[len(repoSplit)-1], ".git")
+	ci, err := gerrit.CreateCLFromLocalDiffs(ctx, g, project, newBranch, commitMsg, !dryRun, co)
+	if err != nil {
+		return skerr.Wrap(err)
+	}
+	fmt.Println(fmt.Sprintf("Uploaded change %s", g.Url(ci.Issue)))
+	return nil
+}
+
+func removeCQ(ctx context.Context, g gerrit.GerritInterface, repo *gitiles.Repo, oldBranch string, dryRun bool) error {
+	// Setup.
+	oldRef := git.FullyQualifiedBranchName(oldBranch)
+	baseCommitInfo, err := repo.Details(ctx, oldRef)
+	if err != nil {
+		return skerr.Wrap(err)
+	}
+	baseCommit := baseCommitInfo.Hash
+
+	// Create the Gerrit CL.
+	if _, err := repo.ReadFileAtRef(ctx, cqJSONFile, oldRef); err != nil {
+		_, _ = fmt.Fprintf(os.Stderr, "Attempt to read %s on %s got error: %s; not removing CQ.\n", oldBranch, cqJSONFile, err)
+		return nil
+	}
+	commitMsg := fmt.Sprintf("Remove CQ for unsupported branch %s", oldBranch)
+	repoSplit := strings.Split(repo.URL(), "/")
+	project := strings.TrimSuffix(repoSplit[len(repoSplit)-1], ".git")
+	changes := map[string]string{
+		cqJSONFile: "",
+	}
+	ci, err := gerrit.CreateCLWithChanges(ctx, g, project, oldRef, commitMsg, baseCommit, changes, !dryRun)
+	if ci != nil {
+		fmt.Println(fmt.Sprintf("Uploaded change %s", g.Url(ci.Issue)))
+	}
+	return skerr.Wrap(err)
 }