[lkg] Add tool to compute last-known-good build/revision

This tool is a skeleton which is intended to consolidate all
"last-known-good" retrieval.

This change currently implements "build" and "revision". We probably
don't need snapshot anymore after "revision" is integrated with recipes.

The algorithm to find the LKGR for n builders is:

1) For 1:n builders' latest SUCCESS builds, create a set of revisions
   which are common to all those builders.
2) Iterate through 0th builder's latest SUCCESS builds and return the
   first revision which is in the set.

We only use one (batched) buildbucket query, so this is more efficient
than the current LKGS implementation. (Overall ~10x faster than the
current application of LKGS). Also, a single buildbucket query makes
this code much easier to mock and unit test, which is partially why I
decided a rewrite would be a better use of time than to retrofit LKGS
and LKGB with unit tests.

Bug: 53486
Change-Id: Ic5e5f9069a837ef58176be00f8d9449e08f0aaca
Reviewed-on: https://fuchsia-review.googlesource.com/c/infra/infra/+/398653
Commit-Queue: Anthony Fandrianto <atyfto@google.com>
Reviewed-by: Nathan Mulcahey <nmulcahey@google.com>
diff --git a/cmd/lkg/buildbucket.go b/cmd/lkg/buildbucket.go
new file mode 100644
index 0000000..eeffe3d
--- /dev/null
+++ b/cmd/lkg/buildbucket.go
@@ -0,0 +1,111 @@
+// Copyright 2020 The Fuchsia Authors. All rights reserved.
+// 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"
+
+	"github.com/golang-collections/collections/set"
+	"go.chromium.org/luci/auth"
+	buildbucketpb "go.chromium.org/luci/buildbucket/proto"
+	"go.chromium.org/luci/grpc/prpc"
+	"google.golang.org/genproto/protobuf/field_mask"
+)
+
+// buildbucketClientWrapper provides utilities for interacting with Buildbucket.
+type buildbucketClientWrapper struct {
+	client buildbucketpb.BuildsClient
+}
+
+// newBuildbucketClient returns an authenticated buildbucketClientWrapper.
+func newBuildbucketClient(ctx context.Context, authOpts auth.Options, host string) (*buildbucketClientWrapper, error) {
+	authClient, err := auth.NewAuthenticator(ctx, auth.OptionalLogin, authOpts).Client()
+	if err != nil {
+		return nil, err
+	}
+	client := buildbucketpb.NewBuildsPRPCClient(&prpc.Client{
+		C:    authClient,
+		Host: host,
+	})
+	return &buildbucketClientWrapper{client: client}, nil
+}
+
+// GetBuilds gets the latest n successful builds for each input builder, from most to least recent.
+// The response is in the same order as the order of the input builders.
+func (c *buildbucketClientWrapper) GetBuilds(ctx context.Context, builders []*buildbucketpb.BuilderID, mask *field_mask.FieldMask, n int) ([][]*buildbucketpb.Build, error) {
+	// Always grab input by default.
+	mask.Paths = append(mask.Paths, "builds.*.input")
+	// Construct requests to execute in batch.
+	reqs := make([]*buildbucketpb.BatchRequest_Request, len(builders))
+	for i, builder := range builders {
+		reqs[i] = &buildbucketpb.BatchRequest_Request{
+			Request: &buildbucketpb.BatchRequest_Request_SearchBuilds{
+				SearchBuilds: &buildbucketpb.SearchBuildsRequest{
+					Predicate: &buildbucketpb.BuildPredicate{
+						Builder: builder,
+						Status:  buildbucketpb.Status_SUCCESS,
+					},
+					Fields:   mask,
+					PageSize: int32(n),
+				},
+			},
+		}
+	}
+	batchResp, err := c.client.Batch(ctx, &buildbucketpb.BatchRequest{Requests: reqs})
+	if err != nil {
+		return nil, fmt.Errorf("failed to retrieve latest %d successful builds for %s: %v", n, builders, err)
+	}
+	builds := make([][]*buildbucketpb.Build, len(builders))
+	for i, resp := range batchResp.Responses {
+		builds[i] = resp.Response.(*buildbucketpb.BatchResponse_Response_SearchBuilds).SearchBuilds.Builds
+	}
+	return builds, nil
+}
+
+// accessorFunc accesses a value from a Build.
+type accessorFunc func(*buildbucketpb.Build) interface{}
+
+// getLastKnownGood gets the latest common attribute from the input slices of builds.
+func getLastKnownGood(builds [][]*buildbucketpb.Build, gitilesRef string, f accessorFunc) (interface{}, error) {
+	if len(builds) == 0 {
+		return nil, errors.New("input builds is of length 0")
+	}
+	// Compute common attributes to all slices but the first.
+	var commonAttrSet *set.Set
+	for _, buildSlice := range builds[1:] {
+		attrSet := set.New()
+		for _, build := range buildSlice {
+			buildAttr := f(build)
+			// Skip nil attributes.
+			if buildAttr != nil {
+				attrSet.Insert(buildAttr)
+			}
+		}
+		// Initialize set on first iteration.
+		if commonAttrSet == nil {
+			commonAttrSet = attrSet
+			// Otherwise, take the intersection between this set and the common set.
+		} else {
+			commonAttrSet = attrSet.Intersection(commonAttrSet)
+		}
+	}
+
+	// Iterate through the first slice, returning the first attribute match.
+	for _, build := range builds[0] {
+		// Skip non-matching refs.
+		if gitilesRef != "" && gitilesRef != build.Input.GitilesCommit.Ref {
+			continue
+		}
+		attr := f(build)
+		// For length-1 inputs, the set is nil, so return the first build.
+		if commonAttrSet == nil || commonAttrSet.Has(attr) {
+			return attr, nil
+		}
+	}
+
+	return nil, errors.New("no common last-known-good attribute found")
+}
diff --git a/cmd/lkg/buildbucket_test.go b/cmd/lkg/buildbucket_test.go
new file mode 100644
index 0000000..3980ccd
--- /dev/null
+++ b/cmd/lkg/buildbucket_test.go
@@ -0,0 +1,134 @@
+// Copyright 2020 The Fuchsia Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+package main
+
+import (
+	"testing"
+
+	buildbucketpb "go.chromium.org/luci/buildbucket/proto"
+)
+
+func mockBuild(revision string) *buildbucketpb.Build {
+	return &buildbucketpb.Build{
+		Input: &buildbucketpb.Build_Input{
+			GitilesCommit: &buildbucketpb.GitilesCommit{
+				Id: revision,
+			},
+		},
+	}
+}
+
+func revisionAccessor(build *buildbucketpb.Build) interface{} {
+	return build.Input.GitilesCommit.Id
+}
+
+func TestGetLastKnownGoodRevision(t *testing.T) {
+	t.Parallel()
+	var tests = []struct {
+		builds      [][]*buildbucketpb.Build
+		gitilesRef  string
+		f           accessorFunc
+		expected    interface{}
+		expectedErr bool
+	}{
+		{
+			builds:      [][]*buildbucketpb.Build{},
+			gitilesRef:  "",
+			f:           revisionAccessor,
+			expected:    nil,
+			expectedErr: true,
+		},
+		{
+			builds: [][]*buildbucketpb.Build{
+				{
+					mockBuild("b"),
+					mockBuild("c"),
+					mockBuild("d"),
+					mockBuild("e"),
+				},
+			},
+			gitilesRef:  "",
+			f:           revisionAccessor,
+			expected:    "b",
+			expectedErr: false,
+		},
+		{
+			builds: [][]*buildbucketpb.Build{
+				{
+					mockBuild("b"),
+					mockBuild("c"),
+					mockBuild("d"),
+					mockBuild("e"),
+					mockBuild("f"),
+					mockBuild("h"),
+				},
+				{
+					mockBuild("a"),
+					mockBuild("b"),
+					mockBuild("c"),
+					mockBuild("e"),
+					mockBuild("f"),
+				},
+				{
+					mockBuild("c"),
+					mockBuild("e"),
+					mockBuild("f"),
+					mockBuild("i"),
+				},
+				{
+					mockBuild("b"),
+					mockBuild("e"),
+					mockBuild("f"),
+					mockBuild("g"),
+					mockBuild("i"),
+				},
+			},
+			gitilesRef:  "",
+			f:           revisionAccessor,
+			expected:    "e",
+			expectedErr: false,
+		},
+		{
+			builds: [][]*buildbucketpb.Build{
+				{
+					mockBuild("b"),
+					mockBuild("c"),
+					mockBuild("d"),
+				},
+				{
+					mockBuild("a"),
+					mockBuild("b"),
+					mockBuild("c"),
+				},
+				{
+					mockBuild("c"),
+					mockBuild("i"),
+				},
+				{
+					mockBuild("e"),
+					mockBuild("f"),
+					mockBuild("h"),
+				},
+			},
+			gitilesRef:  "",
+			f:           revisionAccessor,
+			expected:    nil,
+			expectedErr: true,
+		},
+	}
+	for _, test := range tests {
+		resp, err := getLastKnownGood(test.builds, test.gitilesRef, test.f)
+		if err == nil {
+			if test.expectedErr {
+				t.Errorf("expected error, got nil")
+			}
+		} else if !test.expectedErr {
+			t.Errorf("got unexpected err %v", err)
+		}
+		if resp != test.expected {
+			t.Fatalf("got %s, expected %s", resp, test.expected)
+		}
+	}
+}
diff --git a/cmd/lkg/cmd_build.go b/cmd/lkg/cmd_build.go
new file mode 100644
index 0000000..4917eb8
--- /dev/null
+++ b/cmd/lkg/cmd_build.go
@@ -0,0 +1,90 @@
+// Copyright 2020 The Fuchsia Authors. All rights reserved.
+// 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"
+
+	"github.com/maruel/subcommands"
+	"go.chromium.org/luci/auth"
+	buildbucketpb "go.chromium.org/luci/buildbucket/proto"
+	"google.golang.org/genproto/protobuf/field_mask"
+)
+
+func cmdBuild(authOpts auth.Options) *subcommands.Command {
+	return &subcommands.Command{
+		UsageLine: "build -builder project1/bucket1/builder1, ... [common options]",
+		ShortDesc: "Get the last known good build IDs for a set of builders.",
+		LongDesc:  "Get the last known good build IDs for a set of builders.",
+		CommandRun: func() subcommands.CommandRun {
+			c := &buildRun{}
+			c.Init(authOpts)
+			return c
+		},
+	}
+}
+
+type buildRun struct {
+	commonFlags
+	builders builders
+}
+
+func (c *buildRun) Init(defaultAuthOpts auth.Options) {
+	c.commonFlags.Init(defaultAuthOpts)
+	c.Flags.Var(&c.builders, "builder", "Builder ID as project/bucket/builder")
+}
+
+func (c *buildRun) checkFlags() error {
+	if err := c.commonFlags.checkFlags(); err != nil {
+		return err
+	}
+	if len(c.builders) != 1 {
+		return errors.New("exactly one -builder may be specified")
+	}
+	return nil
+}
+
+func (c *buildRun) main(a subcommands.Application) error {
+	ctx := context.Background()
+	client, err := newBuildbucketClient(ctx, c.parsedAuthOpts, c.buildbucketHost)
+	if err != nil {
+		return fmt.Errorf("failed to initialize client: %v", err)
+	}
+	// Retrieve build ID.
+	mask := &field_mask.FieldMask{
+		Paths: []string{"builds.*.id"},
+	}
+	// If not filtering by gitiles ref, decrease the search range to 1.
+	if c.gitilesRef == "" {
+		c.searchRange = 1
+	}
+	builds, err := client.GetBuilds(ctx, c.builders, mask, c.searchRange)
+	if err != nil {
+		return fmt.Errorf("failed to query builds: %v", err)
+	}
+	bid, err := getLastKnownGood(builds, c.gitilesRef, func(build *buildbucketpb.Build) interface{} {
+		return build.Id
+	})
+	if err != nil {
+		return fmt.Errorf("failed to get last known good build: %v", err)
+	}
+	fmt.Println(bid)
+	return nil
+}
+
+func (c *buildRun) Run(a subcommands.Application, args []string, env subcommands.Env) int {
+	if err := c.checkFlags(); err != nil {
+		fmt.Fprintf(a.GetErr(), "%s: %s\n", a.GetName(), err)
+		return 1
+	}
+
+	if err := c.main(a); err != nil {
+		fmt.Fprintf(a.GetErr(), "%s: %s\n", a.GetName(), err)
+		return 1
+	}
+	return 0
+}
diff --git a/cmd/lkg/cmd_revision.go b/cmd/lkg/cmd_revision.go
new file mode 100644
index 0000000..bdf0f61
--- /dev/null
+++ b/cmd/lkg/cmd_revision.go
@@ -0,0 +1,91 @@
+// Copyright 2020 The Fuchsia Authors. All rights reserved.
+// 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"
+
+	"github.com/maruel/subcommands"
+	"go.chromium.org/luci/auth"
+	buildbucketpb "go.chromium.org/luci/buildbucket/proto"
+	"google.golang.org/genproto/protobuf/field_mask"
+)
+
+func cmdRevision(authOpts auth.Options) *subcommands.Command {
+	return &subcommands.Command{
+		UsageLine: "revision -builder project1/bucket1/builder1, ... [common options]",
+		ShortDesc: "Get the last known good revision across a set of builders.",
+		LongDesc:  "Get the last known good revision across a set of builders.",
+		CommandRun: func() subcommands.CommandRun {
+			c := &revisionRun{}
+			c.Init(authOpts)
+			return c
+		},
+	}
+}
+
+type revisionRun struct {
+	commonFlags
+	builders builders
+}
+
+func (c *revisionRun) Init(defaultAuthOpts auth.Options) {
+	c.commonFlags.Init(defaultAuthOpts)
+	c.Flags.Var(&c.builders, "builder", "Builder ID as project/bucket/builder. Repeatable")
+}
+
+func (c *revisionRun) checkFlags() error {
+	if err := c.commonFlags.checkFlags(); err != nil {
+		return err
+	}
+	if len(c.builders) == 0 {
+		return errors.New("at least one -builder is required")
+	}
+	return nil
+}
+
+func (c *revisionRun) main(a subcommands.Application) error {
+	ctx := context.Background()
+	client, err := newBuildbucketClient(ctx, c.parsedAuthOpts, c.buildbucketHost)
+	if err != nil {
+		return fmt.Errorf("failed to initialize client: %v", err)
+	}
+	// Do not retrieve any fields additional to default input, as we only need revision.
+	mask := &field_mask.FieldMask{
+		Paths: []string{},
+	}
+	builds, err := client.GetBuilds(ctx, c.builders, mask, c.searchRange)
+	if err != nil {
+		return fmt.Errorf("failed to query builds: %v", err)
+	}
+	// Pass an accessorFunc which returns a build's revision.
+	revision, err := getLastKnownGood(builds, c.gitilesRef, func(build *buildbucketpb.Build) interface{} {
+		// Skip manually triggered builds, which do not have a gitiles commit.
+		if build.Input.GitilesCommit == nil {
+			return nil
+		}
+		return build.Input.GitilesCommit.Id
+	})
+	if err != nil {
+		return fmt.Errorf("failed to get last known good revision: %v", err)
+	}
+	fmt.Println(revision)
+	return nil
+}
+
+func (c *revisionRun) Run(a subcommands.Application, args []string, env subcommands.Env) int {
+	if err := c.checkFlags(); err != nil {
+		fmt.Fprintf(a.GetErr(), "%s: %s\n", a.GetName(), err)
+		return 1
+	}
+
+	if err := c.main(a); err != nil {
+		fmt.Fprintf(a.GetErr(), "%s: %s\n", a.GetName(), err)
+		return 1
+	}
+	return 0
+}
diff --git a/cmd/lkg/common.go b/cmd/lkg/common.go
new file mode 100644
index 0000000..360353d
--- /dev/null
+++ b/cmd/lkg/common.go
@@ -0,0 +1,86 @@
+// Copyright 2020 The Fuchsia Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+package main
+
+import (
+	"fmt"
+	"strings"
+
+	"github.com/maruel/subcommands"
+
+	"go.chromium.org/luci/auth"
+	"go.chromium.org/luci/auth/client/authcli"
+	buildbucketpb "go.chromium.org/luci/buildbucket/proto"
+)
+
+type commonFlags struct {
+	subcommands.CommandRunBase
+	buildbucketHost string
+	gitilesRef      string
+	searchRange     int
+	authFlags       authcli.Flags
+
+	parsedAuthOpts auth.Options
+}
+
+func (c *commonFlags) Init(authOpts auth.Options) {
+	c.authFlags = authcli.Flags{}
+	c.authFlags.Register(&c.Flags, authOpts)
+	c.Flags.StringVar(&c.buildbucketHost, "host", "cr-buildbucket.appspot.com", "Buildbucket hostname")
+	c.Flags.StringVar(&c.gitilesRef, "gitiles-ref", "", "Optionally filter for matching gitiles ref e.g. refs/heads/master")
+	c.Flags.IntVar(&c.searchRange, "search-range", 25, "Number of builds to search")
+}
+
+func (c *commonFlags) checkFlags() error {
+	var err error
+	c.parsedAuthOpts, err = c.authFlags.Options()
+	if err != nil {
+		return err
+	}
+	if c.buildbucketHost == "" {
+		return fmt.Errorf("-host is required")
+	}
+	return nil
+}
+
+// newBuilderID returns a *buildbucketpb.BuilderID for a project/bucket/builder string.
+func newBuilderID(builder string) (*buildbucketpb.BuilderID, error) {
+	components := strings.Split(builder, "/")
+	if len(components) != 3 {
+		return nil, fmt.Errorf("failed to parse builder %s", builder)
+	}
+	return &buildbucketpb.BuilderID{
+		Project: components[0],
+		Bucket:  components[1],
+		Builder: components[2],
+	}, nil
+}
+
+// builders is a flag.Getter implementation representing a []*buildbucketpb.BuilderID.
+type builders []*buildbucketpb.BuilderID
+
+// String returns a comma-separated string representation of the builder IDs.
+func (b builders) String() string {
+	strs := make([]string, len(b))
+	for i, bid := range b {
+		strs[i] = bid.String()
+	}
+	return strings.Join(strs, ", ")
+}
+
+// Set records seeing a flag value.
+func (b *builders) Set(val string) error {
+	bid, err := newBuilderID(val)
+	if err != nil {
+		return err
+	}
+	*b = append(*b, bid)
+	return nil
+}
+
+// Get retrieves the flag value.
+func (b builders) Get() interface{} {
+	return []*buildbucketpb.BuilderID(b)
+}
diff --git a/cmd/lkg/common_test.go b/cmd/lkg/common_test.go
new file mode 100644
index 0000000..996af46
--- /dev/null
+++ b/cmd/lkg/common_test.go
@@ -0,0 +1,54 @@
+// Copyright 2020 The Fuchsia Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+package main
+
+import (
+	"reflect"
+	"testing"
+
+	buildbucketpb "go.chromium.org/luci/buildbucket/proto"
+)
+
+func TestNewBuilderID(t *testing.T) {
+	t.Parallel()
+	var tests = []struct {
+		input       string
+		expected    *buildbucketpb.BuilderID
+		expectedErr bool
+	}{
+		{
+			input: "foo-project/bar-bucket/baz-builder",
+			expected: &buildbucketpb.BuilderID{
+				Project: "foo-project",
+				Bucket:  "bar-bucket",
+				Builder: "baz-builder",
+			},
+			expectedErr: false,
+		},
+		{
+			input:       "invalid-input",
+			expected:    nil,
+			expectedErr: true,
+		},
+		{
+			input:       "foo-project/bar-bucket/baz-builder/extra",
+			expected:    nil,
+			expectedErr: true,
+		},
+	}
+	for _, test := range tests {
+		builderID, err := newBuilderID(test.input)
+		if err == nil {
+			if test.expectedErr {
+				t.Fatalf("expected error, got nil")
+			}
+		} else if !test.expectedErr {
+			t.Fatalf("got unexpected err %v", err)
+		}
+		if !reflect.DeepEqual(test.expected, builderID) {
+			t.Fatalf("expected %s, got %s", test.expected, builderID)
+		}
+	}
+}
diff --git a/cmd/lkg/main.go b/cmd/lkg/main.go
new file mode 100644
index 0000000..6d7efbd
--- /dev/null
+++ b/cmd/lkg/main.go
@@ -0,0 +1,42 @@
+// Copyright 2020 The Fuchsia Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+package main
+
+import (
+	"log"
+	"os"
+
+	"github.com/maruel/subcommands"
+
+	"go.chromium.org/luci/auth"
+	"go.chromium.org/luci/auth/client/authcli"
+	"go.chromium.org/luci/client/versioncli"
+	"go.chromium.org/luci/hardcoded/chromeinfra"
+)
+
+// Version must be updated on functional change (behavior, arguments, supported commands).
+const version = "0.0.1"
+
+func getApplication(defaultAuthOpts auth.Options) *subcommands.DefaultApplication {
+	return &subcommands.DefaultApplication{
+		Name:  "lkg",
+		Title: "Tool to query a variety of last-known-good entities.",
+		Commands: []*subcommands.Command{
+			cmdBuild(defaultAuthOpts),
+			cmdRevision(defaultAuthOpts),
+			authcli.SubcommandInfo(defaultAuthOpts, "whoami", false),
+			authcli.SubcommandLogin(defaultAuthOpts, "login", false),
+			authcli.SubcommandLogout(defaultAuthOpts, "logout", false),
+			versioncli.CmdVersion(version),
+			subcommands.CmdHelp,
+		},
+	}
+}
+
+func main() {
+	log.SetFlags(log.Lmicroseconds)
+	app := getApplication(chromeinfra.DefaultAuthOptions())
+	os.Exit(subcommands.Run(app, nil))
+}
diff --git a/go.mod b/go.mod
index b5fd797..559d0b1 100644
--- a/go.mod
+++ b/go.mod
@@ -9,6 +9,7 @@
 	github.com/docker/go-connections v0.3.0 // indirect
 	github.com/docker/go-units v0.3.3 // indirect
 	github.com/gogo/protobuf v1.1.1 // indirect
+	github.com/golang-collections/collections v0.0.0-20130729185459-604e922904d3
 	github.com/golang/mock v1.4.0
 	github.com/golang/protobuf v1.4.1
 	github.com/google/go-cmp v0.4.0
diff --git a/go.sum b/go.sum
index 9f9b2b9..4fd65bd 100644
--- a/go.sum
+++ b/go.sum
@@ -53,6 +53,8 @@
 github.com/go-gl/glfw/v3.3/glfw v0.0.0-20191125211704-12ad95a8df72/go.mod h1:tQ2UAYgL5IevRw8kRxooKSPJfGvJ9fJQFa0TUsXzTg8=
 github.com/gogo/protobuf v1.1.1 h1:72R+M5VuhED/KujmZVcIquuo8mBgX4oVda//DQb3PXo=
 github.com/gogo/protobuf v1.1.1/go.mod h1:r8qH/GZQm5c6nD/R0oafs1akxWv10x8SbQlK7atdtwQ=
+github.com/golang-collections/collections v0.0.0-20130729185459-604e922904d3 h1:zN2lZNZRflqFyxVaTIU61KNKQ9C0055u9CAfpmqUvo4=
+github.com/golang-collections/collections v0.0.0-20130729185459-604e922904d3/go.mod h1:nPpo7qLxd6XL3hWJG/O60sR8ZKfMCiIoNap5GvD12KU=
 github.com/golang/glog v0.0.0-20160126235308-23def4e6c14b h1:VKtxabqXZkF25pY9ekfRL6a582T4P37/31XEstQ5p58=
 github.com/golang/glog v0.0.0-20160126235308-23def4e6c14b/go.mod h1:SBH7ygxi8pfUlaOkMMuAQtPIUF8ecWP5IEl/CR7VP2Q=
 github.com/golang/groupcache v0.0.0-20190702054246-869f871628b6/go.mod h1:cIg4eruTrX1D+g88fzRXU5OdNfaM+9IcxsU14FzY7Hc=