go.fuchsia.dev/infra@v0.0.0-20240507153436-9b593402251b/cmd/size_diff/common.go

// Copyright 2021 The Fuchsia Authors.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

package main

import (
	"context"
	"errors"
	"fmt"
	"log"
	"strings"
	"time"

	"github.com/maruel/subcommands"
	"go.chromium.org/luci/auth"
	"go.chromium.org/luci/auth/client/authcli"
	buildbucketpb "go.chromium.org/luci/buildbucket/proto"
	"go.chromium.org/luci/common/proto/structmask"
	"go.chromium.org/luci/common/retry"
	"go.chromium.org/luci/common/retry/transient"
	"go.chromium.org/luci/grpc/prpc"
	"go.fuchsia.dev/infra/buildbucket"
	"google.golang.org/protobuf/types/known/fieldmaskpb"
)

// buildFieldMask identifies required fields and output properties that should
// be requested from the Buildbucket API.
func buildFieldMask(outputProps []string) *buildbucketpb.BuildMask {
	return &buildbucketpb.BuildMask{
		Fields: &fieldmaskpb.FieldMask{Paths: []string{
			"id",
			"status",
		}},
		OutputProperties: []*structmask.StructMask{{Path: outputProps}},
	}
}

type commonFlags struct {
	subcommands.CommandRunBase
	authFlags authcli.Flags

	parsedAuthOpts auth.Options
	bbHost         string
	gitilesRemote  string
	baseCommit     string
	builder        buildbucket.BuilderIDFlag
	jsonOutput     string
}

func (c *commonFlags) Init(authOpts auth.Options) {
	c.authFlags = authcli.Flags{}
	c.authFlags.Register(&c.Flags, authOpts)
	c.Flags.StringVar(&c.bbHost, "bb-host", "cr-buildbucket.appspot.com", "Buildbucket host to use.")
	c.Flags.StringVar(&c.gitilesRemote, "gitiles-remote", "", "Gitiles remote for base commit.")
	c.Flags.StringVar(&c.baseCommit, "base-commit", "", "Base commit as sha1.")
	c.Flags.Var(&c.builder, "builder", "Fully-qualified Buildbucket CI builder name to inspect.")
	c.Flags.StringVar(&c.jsonOutput, "json-output", "-", "Path to output results to.")
}

func (c *commonFlags) Parse() error {
	var err error
	c.parsedAuthOpts, err = c.authFlags.Options()
	if err != nil {
		return err
	}
	if c.gitilesRemote == "" {
		return errors.New("-gitiles-remote is required")
	}
	c.gitilesRemote = strings.TrimPrefix(c.gitilesRemote, "https://")
	if c.baseCommit == "" {
		return errors.New("-base-commit is required")
	}
	if &c.builder == nil {
		return errors.New("-builder is required")
	}
	if c.jsonOutput == "" {
		return errors.New("-json-output is required")
	}
	return nil
}

// searchForBuild gets the CI build with gitiles commit matching the base
// commit.
func searchForBuild(ctx context.Context, buildsClient buildbucketpb.BuildsClient, builder *buildbucketpb.BuilderID, gitilesRemote, baseCommit string, outputProps []string) (*buildbucketpb.Build, error) {
	fullBuilderName := fmt.Sprintf("%s/%s/%s", builder.Project, builder.Bucket, builder.Builder)
	resp, err := buildsClient.SearchBuilds(ctx, &buildbucketpb.SearchBuildsRequest{
		Predicate: &buildbucketpb.BuildPredicate{
			Builder: builder,
			Tags: []*buildbucketpb.StringPair{
				{
					Key:   "buildset",
					Value: fmt.Sprintf("commit/gitiles/%s/+/%s", gitilesRemote, baseCommit),
				},
			},
		},
		Mask:     buildFieldMask(outputProps),
		PageSize: int32(1),
	})
	if err != nil {
		return nil, fmt.Errorf("failed to query builds for %s: %w", fullBuilderName, err)
	}
	builds := resp.Builds
	if len(builds) == 0 {
		return nil, fmt.Errorf("no builds found for %s", fullBuilderName)
	}
	if len(builds) != 1 {
		return nil, fmt.Errorf("expected exactly one build, got %d", len(builds))
	}
	return builds[0], nil
}

// waitForBuildCompletion waits for the given build to complete, as CI could be
// on the cusp of completion.
func waitForBuildCompletion(ctx context.Context, buildsClient buildbucketpb.BuildsClient, buildID int64, bbHost string, outputProps []string) (*buildbucketpb.Build, error) {
	buildLink := fmt.Sprintf("https://%s/build/%d", bbHost, buildID)
	retryPolicy := transient.Only(func() retry.Iterator {
		return &retry.ExponentialBackoff{
			Limited: retry.Limited{
				Delay:   time.Minute,
				Retries: 20,
			},
			Multiplier: 1,
		}
	})
	var build *buildbucketpb.Build
	if err := retry.Retry(ctx, retryPolicy, func() error {
		var err error
		build, err = buildsClient.GetBuild(ctx, &buildbucketpb.GetBuildRequest{
			Id:   buildID,
			Mask: buildFieldMask(outputProps),
		})
		if err != nil {
			return fmt.Errorf("failed to get build %d: %w", buildID, err)
		}
		if !hasEnded(build) {
			err := buildNotSuccessfulError{
				msg:    fmt.Sprintf("a finished build is needed to attempt the size diff but got status %s, see %s", build.Status, buildLink),
				status: build.Status,
			}
			log.Printf("%s is still status %s", buildLink, build.Status)
			return transient.Tag.Apply(err)
		}
		return nil
	}, nil); err != nil {
		return nil, err
	}
	return build, nil
}

// hasEnded returns whether a build has completed.
func hasEnded(build *buildbucketpb.Build) bool {
	bitmasked := build.Status & buildbucketpb.Status_ENDED_MASK
	return bitmasked > 0
}

func getBuild(ctx context.Context, c commonFlags, outputProps []string) (*buildbucketpb.Build, error) {
	authClient, err := auth.NewAuthenticator(ctx, auth.OptionalLogin, c.parsedAuthOpts).Client()
	if err != nil {
		return nil, fmt.Errorf("failed to initialize auth client: %w", err)
	}

	buildsClient := buildbucketpb.NewBuildsPRPCClient(&prpc.Client{
		C:    authClient,
		Host: c.bbHost,
	})
	builder, ok := c.builder.Get().(*buildbucketpb.BuilderID)
	if !ok {
		return nil, errors.New("builder input is not of the form project/bucket/builder")
	}

	build, err := searchForBuild(ctx, buildsClient, builder, c.gitilesRemote, c.baseCommit, outputProps)
	if err != nil {
		return nil, err
	}

	if hasEnded(build) {
		// The build has already finished, exit early to avoid an extra RPC
		// request.
		return build, err
	}
	return waitForBuildCompletion(ctx, buildsClient, build.Id, c.bbHost, outputProps)
}