dashboard, buildlet: add support for ARM64 COS images

This adds the ability to query for COS images which are ARM64 based.
The HostConfig has been updated to provide the ability to note that a
host requires an ARM64 COS image.

Updates golang/go#53851

Change-Id: Ib57bf6359ca365078528dbb6912dca2402939df9
Reviewed-on: https://go-review.googlesource.com/c/build/+/447259
TryBot-Result: Gopher Robot <[email protected]>
Run-TryBot: Carlos Amedee <[email protected]>
Reviewed-by: Carlos Amedee <[email protected]>
Reviewed-by: Heschi Kreinick <[email protected]>

buildlet/gce.go

@@ -121,7 +121,7 @@ func StartNewVM(creds *google.Credentials, buildEnv *buildenv.Environment, instN
 			srcImage = "https://www.googleapis.com/compute/v1/projects/" + projectID + "/global/images/" + vm
 		} else {
 			var err error
-			srcImage, err = cosImage(ctx, computeService)
+			srcImage, err = cosImage(ctx, computeService, hconf.CosArchitecture())
 			if err != nil {
 				return nil, fmt.Errorf("error find Container-Optimized OS image: %v", err)
 			}
@@ -417,41 +417,63 @@ func instanceIPs(inst *compute.Instance) (intIP, extIP string) {
 }
 
 var (
-	cosListMu      sync.Mutex
-	cosCachedTime  time.Time
-	cosCachedImage string
+	cosListMu     sync.Mutex
+	cosCachedTime time.Time
+	cosCache      = map[dashboard.CosArch]*cosCacheEntry{}
 )
 
+type cosCacheEntry struct {
+	cachedTime  time.Time
+	cachedImage string
+}
+
 // cosImage returns the GCP VM image name of the latest stable
 // Container-Optimized OS image. It caches results for 15 minutes.
-func cosImage(ctx context.Context, svc *compute.Service) (string, error) {
+func cosImage(ctx context.Context, svc *compute.Service, arch dashboard.CosArch) (string, error) {
 	const cacheDuration = 15 * time.Minute
 	cosListMu.Lock()
 	defer cosListMu.Unlock()
-	if cosCachedImage != "" && cosCachedTime.After(time.Now().Add(-cacheDuration)) {
-		return cosCachedImage, nil
-	}
 
-	imList, err := svc.Images.List("cos-cloud").Filter(`(family eq "cos-stable")`).Context(ctx).Do()
-	if err != nil {
-		return "", err
+	cosQuery := func(a dashboard.CosArch) (string, error) {
+		imList, err := svc.Images.List("cos-cloud").Filter(fmt.Sprintf("(family eq %q)", string(arch))).Context(ctx).Do()
+		if err != nil {
+			return "", err
+		}
+		if imList.NextPageToken != "" {
+			return "", fmt.Errorf("too many images; pagination not supported")
+		}
+		ims := imList.Items
+		if len(ims) == 0 {
+			return "", errors.New("no image found")
+		}
+		sort.Slice(ims, func(i, j int) bool {
+			if ims[i].Deprecated == nil && ims[j].Deprecated != nil {
+				return true
+			}
+			return ims[i].CreationTimestamp > ims[j].CreationTimestamp
+		})
+		return ims[0].SelfLink, nil
 	}
-	if imList.NextPageToken != "" {
-		return "", fmt.Errorf("too many images; pagination not supported")
+	c, ok := cosCache[arch]
+	if !ok {
+		image, err := cosQuery(arch)
+		if err != nil {
+			return "", err
+		}
+		cosCache[arch] = &cosCacheEntry{
+			cachedTime:  time.Now(),
+			cachedImage: image,
+		}
+		return image, nil
 	}
-	ims := imList.Items
-	if len(ims) == 0 {
-		return "", errors.New("no image found")
+	if c.cachedImage != "" && c.cachedTime.After(time.Now().Add(-cacheDuration)) {
+		return c.cachedImage, nil
 	}
-	sort.Slice(ims, func(i, j int) bool {
-		if ims[i].Deprecated == nil && ims[j].Deprecated != nil {
-			return true
-		}
-		return ims[i].CreationTimestamp > ims[j].CreationTimestamp
-	})
-
-	im := ims[0].SelfLink
-	cosCachedImage = im
-	cosCachedTime = time.Now()
-	return im, nil
+	image, err := cosQuery(arch)
+	if err != nil {
+		return "", err
+	}
+	c.cachedImage = image
+	c.cachedTime = time.Now()
+	return image, nil
 }

dashboard/builders.go

@@ -361,6 +361,13 @@ var Hosts = map[string]*HostConfig{
 		isEC2:          true,
 		SSHUsername:    "root",
 	},
+	"host-linux-arm64-bullseye": {
+		Notes:           "Debian Bullseye",
+		ContainerImage:  "linux-arm64-bullseye:latest",
+		machineType:     "t2a",
+		SSHUsername:     "root",
+		cosArchitecture: CosArchARM64,
+	},
 	"host-linux-loong64-3a5000": {
 		Notes:       "Loongson 3A5000 Box hosted by Loongson; loong64 is the short name of LoongArch 64 bit version",
 		Owners:      []*gophers.Person{gh("XiaodongLoong")},
@@ -635,6 +642,14 @@ func init() {
 	}
 }
 
+// CosArch defines the diffrent COS images types used.
+type CosArch string
+
+const (
+	CosArchAMD64 CosArch = "cos-stable"       // COS image for AMD64 architecture
+	CosArchARM64         = "cos-arm64-stable" // COS image for ARM64 architecture
+)
+
 // A HostConfig describes the available ways to obtain buildlets of
 // different types. Some host configs can serve multiple
 // builders. For example, a host config of "host-linux-amd64-bullseye" can
@@ -678,9 +693,10 @@ type HostConfig struct {
 	IsReverse      bool   // if true, only use the reverse buildlet pool
 
 	// GCE options, if VMImage != "" || ContainerImage != ""
-	machineType    string // optional GCE instance type ("n2-standard-4") or instance class ("n2")
-	RegularDisk    bool   // if true, use spinning disk instead of SSD
-	MinCPUPlatform string // optional. e2 instances are not supported. see https://cloud.google.com/compute/docs/instances/specify-min-cpu-platform.
+	machineType     string  // optional GCE instance type ("n2-standard-4") or instance class ("n2")
+	RegularDisk     bool    // if true, use spinning disk instead of SSD
+	MinCPUPlatform  string  // optional. e2 instances are not supported. see https://cloud.google.com/compute/docs/instances/specify-min-cpu-platform.
+	cosArchitecture CosArch // optional. GCE instances which use COS need the architecture set. Default: CosArchAMD64
 
 	// EC2 options
 	isEC2 bool // if true, the instance is configured to run on EC2
@@ -710,6 +726,14 @@ type HostConfig struct {
 	SSHUsername string // username to ssh as, empty means not supported
 }
 
+// CosArchitecture returns the COS architecture to use with the host. The default is CosArchAMD64.
+func (hc *HostConfig) CosArchitecture() CosArch {
+	if hc.cosArchitecture == CosArch("") {
+		return CosArchAMD64
+	}
+	return hc.cosArchitecture
+}
+
 // A BuildConfig describes how to run a builder.
 type BuildConfig struct {
 	// Name is the unique name of the builder, in the form of
@@ -2489,6 +2513,11 @@ func init() {
 		tryBot:            defaultTrySet(),
 		numTryTestHelpers: 1,
 	})
+	addBuilder(BuildConfig{
+		Name:        "linux-arm64",
+		HostType:    "host-linux-arm64-bullseye",
+		KnownIssues: []int{53851},
+	})
 	addBuilder(BuildConfig{
 		Name:     "linux-arm64-boringcrypto",
 		HostType: "host-linux-arm64-aws",

dashboard/builders_test.go

@@ -1113,3 +1113,22 @@ func TestHostsSort(t *testing.T) {
 		last = key
 	}
 }
+
+func TestHostConfigCosArchitecture(t *testing.T) {
+	testCases := []struct {
+		desc       string
+		hostConfig HostConfig
+		want       CosArch
+	}{
+		{"default", HostConfig{}, CosArchAMD64},
+		{"amd64", HostConfig{cosArchitecture: CosArchAMD64}, CosArchAMD64},
+		{"arm64", HostConfig{cosArchitecture: CosArchARM64}, CosArchARM64},
+	}
+	for _, tc := range testCases {
+		t.Run(tc.desc, func(t *testing.T) {
+			if got := tc.hostConfig.CosArchitecture(); got != tc.want {
+				t.Errorf("HostConfig.CosArchitecture() = %+v; want %+v", got, tc.want)
+			}
+		})
+	}
+}