autoroll/go/autoroll-meta-config-generator/main.go

package main

import (
	"bytes"
	"context"
	"flag"
	"fmt"
	"path"
	"path/filepath"
	"strings"
	"text/template"
	"time"

	"go.skia.org/infra/autoroll/go/config"
	"go.skia.org/infra/autoroll/go/config_vars"
	"go.skia.org/infra/go/chrome_branch"
	"go.skia.org/infra/go/common"
	"go.skia.org/infra/go/gerrit"
	"go.skia.org/infra/go/gerrit/rubberstamper"
	"go.skia.org/infra/go/gitiles"
	"go.skia.org/infra/go/httputils"
	"go.skia.org/infra/go/metrics2"
	"go.skia.org/infra/go/skerr"
	"go.skia.org/infra/go/sklog"
	"go.skia.org/infra/go/util"
	"golang.org/x/oauth2/google"
	"google.golang.org/protobuf/encoding/prototext"
)

const (
	generatedFileHeaderPrefix = "# This file was generated by autoroll-meta-config-generator. DO NOT EDIT."
	generatedFileHeaderTmpl   = generatedFileHeaderPrefix + `
# Generated from: %s

`
	livenessMetric = "autoroll_meta_config_generator"
)

// flags
var (
	local         = flag.Bool("local", false, "Running locally if true. As opposed to in production.")
	port          = flag.String("port", ":8000", "HTTP service port.")
	promPort      = flag.String("prom_port", ":20000", "Metrics service address (e.g., ':10110')")
	srcsDsts      = common.NewMultiStringFlag("src-dst", nil, "Source and destination directory pairs in 'path/to/src:path/to/dst' format.  May be specified multiple times.")
	reviewers     = flag.String("reviewers", "", "Comma-separated list of reviewer email addresses.")
	repoUrl       = flag.String("repo", "", "Git repo URL containing the config files.")
	ref           = flag.String("ref", "main", "Git ref to read and write.")
	gerritUrl     = flag.String("gerrit-url", "", "Gerrit host URL.")
	gerritProject = flag.String("gerrit-project", "", "Gerrit project name.")
	interval      = flag.Duration("interval", 30*time.Minute, "How often to regenerate configs.")
)

func main() {
	common.InitWithMust(
		"autoroll-meta-config-generator",
		common.PrometheusOpt(promPort),
		common.MetricsLoggingOpt(),
	)
	if len(*srcsDsts) == 0 {
		sklog.Fatal("--src-dst is required.")
	}
	if *reviewers == "" {
		sklog.Fatal("--reviewers is required.")
	}
	if *repoUrl == "" {
		sklog.Fatal("--repo is required.")
	}
	if *gerritUrl == "" {
		sklog.Fatal("--gerrit-url is required")
	}
	if *gerritProject == "" {
		sklog.Fatal("--gerrit-project is required")
	}

	reviewersList := strings.Split(*reviewers, ",")
	srcDstMap := make(map[string]string, len(*srcsDsts))
	for _, srcDst := range *srcsDsts {
		split := strings.Split(srcDst, ":")
		if len(split) != 2 {
			sklog.Fatalf("--src-dst must be in 'path/to/src:path/to/dst' format, not %q", srcDst)
		}
		srcDstMap[split[0]] = split[1]
	}

	ctx := context.Background()
	ts, err := google.DefaultTokenSource(ctx, gerrit.AuthScope)
	if err != nil {
		sklog.Fatal(err)
	}
	client := httputils.DefaultClientConfig().WithTokenSource(ts).With2xxOnly().Client()
	reg, err := config_vars.NewRegistry(ctx, chrome_branch.NewClient(client))
	if err != nil {
		sklog.Fatal(err)
	}
	repo := gitiles.NewRepo(*repoUrl, client)
	g, err := gerrit.NewGerrit(*gerritUrl, client)
	if err != nil {
		sklog.Fatal(err)
	}

	liveness := metrics2.NewLiveness(livenessMetric)
	var ci *gerrit.ChangeInfo
	go util.RepeatCtx(ctx, *interval, func(ctx context.Context) {
		ci, err = tick(ctx, srcDstMap, reg, reviewersList, repo, *ref, g, *gerritProject, ci)
		if err != nil {
			sklog.Error(err)
		} else {
			liveness.Reset()
		}
	})
	httputils.RunHealthCheckServer(*port)
}

func tick(ctx context.Context, srcDstMap map[string]string, reg *config_vars.Registry, reviewers []string, repo gitiles.GitilesRepo, ref string, g gerrit.GerritInterface, gerritProject string, oldChange *gerrit.ChangeInfo) (*gerrit.ChangeInfo, error) {
	// Is the previously-uploaded change still active?
	if oldChange != nil {
		updatedChange, err := g.GetChange(ctx, oldChange.Id)
		if err != nil {
			return oldChange, skerr.Wrap(err)
		}
		// If so, don't upload a new one.
		if !updatedChange.IsClosed() {
			return updatedChange, nil
		}
	}
	return maybeUploadCL(ctx, srcDstMap, reg, reviewers, repo, ref, g, gerritProject)
}

func maybeUploadCL(ctx context.Context, srcDstMap map[string]string, reg *config_vars.Registry, reviewers []string, repo gitiles.GitilesRepo, ref string, g gerrit.GerritInterface, gerritProject string) (*gerrit.ChangeInfo, error) {
	// Update the config vars.
	if err := reg.Update(ctx); err != nil {
		return nil, skerr.Wrapf(err, "failed to update config vars")
	}
	vars := reg.Vars()
	commit, err := repo.ResolveRef(ctx, ref)
	if err != nil {
		return nil, skerr.Wrapf(err, "failed to resolve ref %q", ref)
	}

	// Process all of the templates
	changes := map[string]string{}
	for src, dst := range srcDstMap {
		dirChanges, err := processDir(ctx, src, dst, vars, repo, commit)
		if err != nil {
			return nil, skerr.Wrapf(err, "failed to process %s", src)
		}
		for k, v := range dirChanges {
			changes[k] = v
		}
	}
	if len(changes) == 0 {
		return nil, nil
	}

	// Upload a CL.
	commitMsg := `[autoroll] Update generated config files`
	ci, err := gerrit.CreateCLWithChanges(ctx, g, gerritProject, ref, commitMsg, commit, changes, reviewers)
	if err != nil {
		return nil, skerr.Wrap(err)
	}
	labels := map[string]int{
		gerrit.LabelChromiumAutoSubmit: gerrit.LabelChromiumAutoSubmitSubmit,
	}
	if err := g.SetReview(ctx, ci, "", labels, append(reviewers, rubberstamper.RubberStamperUser), "", nil, "", 0, nil); err != nil {
		return nil, skerr.Wrap(err)
	}
	sklog.Infof("Created %s", g.Url(ci.Issue))
	return ci, nil
}

// processDir converts a directory of templates into a directory of configs.
func processDir(ctx context.Context, srcRelPath, dstRelPath string, vars *config_vars.Vars, repo gitiles.GitilesRepo, commit string) (map[string]string, error) {
	// Remove all configs which were generated by this tool.  This prevents us
	// from leaving outdated configs around.
	changes := map[string]string{}
	dstPaths, err := repo.ListFilesRecursiveAtRef(ctx, dstRelPath, commit)
	if err != nil {
		return nil, skerr.Wrapf(err, "failed to list files in %s", dstRelPath)
	}
	for _, cfgPath := range dstPaths {
		if strings.HasSuffix(cfgPath, ".cfg") {
			cfgPath = path.Join(dstRelPath, cfgPath)
			contents, err := repo.ReadFileAtRef(ctx, cfgPath, commit)
			if err != nil {
				return nil, skerr.Wrapf(err, "failed to retrieve %q", cfgPath)
			}
			if strings.HasPrefix(string(contents), generatedFileHeaderPrefix) {
				changes[cfgPath] = ""
			}
		}
	}

	// Find the templates in the source directory.
	srcPaths, err := repo.ListFilesRecursiveAtRef(ctx, srcRelPath, commit)
	if err != nil {
		return nil, skerr.Wrapf(err, "failed to list files in %s", srcRelPath)
	}
	for _, srcPath := range srcPaths {
		srcPath := path.Join(srcRelPath, srcPath)
		if strings.HasSuffix(srcPath, ".tmpl") {
			tmplContents, err := repo.ReadFileAtRef(ctx, srcPath, commit)
			if err != nil {
				return nil, skerr.Wrapf(err, "failed to read %s", srcPath)
			}
			tmplChanges, err := process(ctx, srcPath, string(tmplContents), dstRelPath, vars)
			if err != nil {
				return nil, skerr.Wrapf(err, "failed to convert config %s", srcPath)
			}
			for path, newContents := range tmplChanges {
				// Un-delete the file.
				delete(changes, path)

				// Load the original version of the config file. Only add it to
				// the changes map if it doesn't already exist, or if it exists
				// with different contents.
				origContents, err := repo.ReadFileAtRef(ctx, path, commit)
				if err != nil || string(origContents) != newContents {
					// Assume any error means that the file doesn't exist.
					changes[path] = newContents
				}
			}
		}
	}
	return changes, nil
}

var (
	// protoMarshalOptions are used when writing configs in text proto format.
	protoMarshalOptions = prototext.MarshalOptions{
		Multiline: true,
	}
	// funcMap is used for executing templates.
	funcMap = template.FuncMap{
		"map": makeMap,
	}
)

// process converts a single template into at least one config.
func process(ctx context.Context, srcPath, tmplContents, dstDir string, vars *config_vars.Vars) (map[string]string, error) {
	// Read and execute the template.
	tmpl, err := template.New(filepath.Base(srcPath)).Funcs(funcMap).Parse(tmplContents)
	if err != nil {
		return nil, skerr.Wrapf(err, "failed to parse template file %q", srcPath)
	}
	var buf bytes.Buffer
	if err := tmpl.Execute(&buf, vars); err != nil {
		return nil, skerr.Wrapf(err, "failed to execute template file %q", srcPath)
	}

	// Parse the template as a list of Configs.
	var configs config.Configs
	if err := prototext.Unmarshal(buf.Bytes(), &configs); err != nil {
		return nil, skerr.Wrapf(err, "failed to parse config proto from template file %q", srcPath)
	}
	sklog.Infof("  Found %d configs in %s", len(configs.Config), srcPath)
	changes := make(map[string]string, len(configs.Config))
	headerBytes := []byte(fmt.Sprintf(generatedFileHeaderTmpl, srcPath))
	for _, cfg := range configs.Config {
		encBytes, err := protoMarshalOptions.Marshal(cfg)
		if err != nil {
			return nil, skerr.Wrapf(err, "failed to encode config from %q", srcPath)
		}
		dstPath := path.Join(dstDir, cfg.RollerName+".cfg")
		changes[dstPath] = string(append(headerBytes, encBytes...))
	}
	return changes, nil
}

func makeMap(elems ...interface{}) (map[string]interface{}, error) {
	if len(elems)%2 != 0 {
		return nil, skerr.Fmt("Requires an even number of elements, not %d", len(elems))
	}
	rv := make(map[string]interface{}, len(elems)/2)
	for i := 0; i < len(elems); i += 2 {
		key, ok := elems[i].(string)
		if !ok {
			return nil, skerr.Fmt("Map keys must be strings, not %v", elems[i])
		}
		rv[key] = elems[i+1]
	}
	return rv, nil
}