internal/cloud: add quota and instance type functions

This change adds a quota function to the AWS client. It enables
the caller to call the AWS API in order to retrieve quota (soon
to be called limits) on resource limits for caller. This will be used
to set the resource limits on how many vCPU's will be allowed to
reserved for use in the buildlet pool.

The InstanceTypesArm function has been added which will call the AWS
API and retrieve all instance types which support the arm64
architecture. Adding this allows us to store the instance types which
could possibly be called by the buildlet pool and know how many
vCPU's would be reseved for each instance that has been requested.

Updates golang/go#36841

Change-Id: Ib280a41c72f9859876fe03ee2a0d8d5eaf12cc9b
Reviewed-on: https://go-review.googlesource.com/c/build/+/243198
Run-TryBot: Carlos Amedee <carlos@golang.org>
TryBot-Result: Gobot Gobot <gobot@golang.org>
Reviewed-by: Alexander Rakoczy <alex@golang.org>
Reviewed-by: Dmitri Shuralyov <dmitshur@golang.org>
Reviewed-by: Andrew Bonventre <andybons@golang.org>

internal/cloud/aws.go

@@ -18,6 +18,7 @@ import (
 	"github.com/aws/aws-sdk-go/aws/request"
 	"github.com/aws/aws-sdk-go/aws/session"
 	"github.com/aws/aws-sdk-go/service/ec2"
+	"github.com/aws/aws-sdk-go/service/servicequotas"
 )
 
 const (
@@ -27,14 +28,29 @@ const (
 	tagDescription = "Description"
 )
 
+const (
+	// QuotaCodeCPUOnDemand is the quota code for on-demand CPUs.
+	QuotaCodeCPUOnDemand = "L-1216C47A"
+	// QuotaServiceEC2 is the service code for the EC2 service.
+	QuotaServiceEC2 = "ec2"
+)
+
 // vmClient defines the interface used to call the backing EC2 service. This is a partial interface
-// based on the EC2 package defined at `github.com/aws/aws-sdk-go/service/ec2`.
+// based on the EC2 package defined at github.com/aws/aws-sdk-go/service/ec2.
 type vmClient interface {
 	DescribeInstancesPagesWithContext(context.Context, *ec2.DescribeInstancesInput, func(*ec2.DescribeInstancesOutput, bool) bool, ...request.Option) error
 	DescribeInstancesWithContext(context.Context, *ec2.DescribeInstancesInput, ...request.Option) (*ec2.DescribeInstancesOutput, error)
 	RunInstancesWithContext(context.Context, *ec2.RunInstancesInput, ...request.Option) (*ec2.Reservation, error)
 	TerminateInstancesWithContext(context.Context, *ec2.TerminateInstancesInput, ...request.Option) (*ec2.TerminateInstancesOutput, error)
 	WaitUntilInstanceRunningWithContext(context.Context, *ec2.DescribeInstancesInput, ...request.WaiterOption) error
+	DescribeInstanceTypesPagesWithContext(context.Context, *ec2.DescribeInstanceTypesInput, func(*ec2.DescribeInstanceTypesOutput, bool) bool, ...request.Option) error
+}
+
+// quotaClient defines the interface used to call the backing service quotas service. This
+// is a partial interface based on the service quota package defined at
+// github.com/aws/aws-sdk-go/service/servicequotas.
+type quotaClient interface {
+	GetServiceQuota(*servicequotas.GetServiceQuotaInput) (*servicequotas.GetServiceQuotaOutput, error)
 }
 
 // EC2VMConfiguration is the configuration needed for an EC2 instance.
@@ -98,7 +114,8 @@ type Instance struct {
 
 // AWSClient is a client for AWS services.
 type AWSClient struct {
-	ec2Client vmClient
+	ec2Client   vmClient
+	quotaClient quotaClient
 }
 
 // NewAWSClient creates a new AWS client.
@@ -111,7 +128,8 @@ func NewAWSClient(region, keyID, accessKey string) (*AWSClient, error) {
 		return nil, fmt.Errorf("failed to create AWS session: %v", err)
 	}
 	return &AWSClient{
-		ec2Client: ec2.New(s),
+		ec2Client:   ec2.New(s),
+		quotaClient: servicequotas.New(s),
 	}, nil
 }
 
@@ -198,6 +216,58 @@ func (ac *AWSClient) WaitUntilInstanceRunning(ctx context.Context, instID string
 	return err
 }
 
+// InstanceType contains information about an EC2 vm instance type.
+type InstanceType struct {
+	// Type is the textual label used to describe an instance type.
+	Type string
+	// CPU is the Default vCPU count.
+	CPU int64
+}
+
+// InstanceTypesARM retrieves all EC2 instance types in a region which support the
+// ARM64 architecture.
+func (ac *AWSClient) InstanceTypesARM(ctx context.Context) ([]*InstanceType, error) {
+	var its []*InstanceType
+	contains := func(strs []*string, want string) bool {
+		for _, s := range strs {
+			if aws.StringValue(s) == want {
+				return true
+			}
+		}
+		return false
+	}
+	fn := func(page *ec2.DescribeInstanceTypesOutput, lastPage bool) bool {
+		for _, it := range page.InstanceTypes {
+			if !contains(it.ProcessorInfo.SupportedArchitectures, "arm64") {
+				continue
+			}
+			its = append(its, &InstanceType{
+				Type: aws.StringValue(it.InstanceType),
+				CPU:  aws.Int64Value(it.VCpuInfo.DefaultVCpus),
+			})
+		}
+		return true
+	}
+	err := ac.ec2Client.DescribeInstanceTypesPagesWithContext(ctx, &ec2.DescribeInstanceTypesInput{}, fn)
+	if err != nil {
+		return nil, fmt.Errorf("failed to retrieve arm64 instance types: %w", err)
+	}
+	return its, nil
+}
+
+// Quota retrieves the requested service quota for the service.
+func (ac *AWSClient) Quota(ctx context.Context, service, code string) (int64, error) {
+	// TODO(golang.org/issue/36841): use ctx
+	sq, err := ac.quotaClient.GetServiceQuota(&servicequotas.GetServiceQuotaInput{
+		QuotaCode:   aws.String(code),
+		ServiceCode: aws.String(service),
+	})
+	if err != nil {
+		return 0, fmt.Errorf("failed to retrieve quota: %w", err)
+	}
+	return int64(aws.Float64Value(sq.Quota.Value)), nil
+}
+
 // ec2ToInstance converts an `ec2.Instance` to an `Instance`
 func ec2ToInstance(inst *ec2.Instance) *Instance {
 	secGroup := make([]string, 0, len(inst.SecurityGroups))

internal/cloud/aws_test.go

@@ -17,20 +17,30 @@ import (
 	"github.com/aws/aws-sdk-go/aws"
 	"github.com/aws/aws-sdk-go/aws/request"
 	"github.com/aws/aws-sdk-go/service/ec2"
+	"github.com/aws/aws-sdk-go/service/servicequotas"
 	"github.com/google/go-cmp/cmp"
 )
 
-var _ vmClient = (*fakeEC2Client)(nil)
+type awsClient interface {
+	vmClient
+	quotaClient
+}
+
+var _ awsClient = (*fakeEC2Client)(nil)
 
 type fakeEC2Client struct {
 	mu sync.RWMutex
 	// instances map of instanceId -> *ec2.Instance
-	instances map[string]*ec2.Instance
+	instances     map[string]*ec2.Instance
+	instanceTypes []*ec2.InstanceTypeInfo
+	serviceQuota  map[string]float64
 }
 
-func newFakeEC2Client() *fakeEC2Client {
+func newFakeAWSClient() *fakeEC2Client {
 	return &fakeEC2Client{
-		instances: make(map[string]*ec2.Instance),
+		instances:     make(map[string]*ec2.Instance),
+		instanceTypes: []*ec2.InstanceTypeInfo{},
+		serviceQuota:  make(map[string]float64),
 	}
 }
 
@@ -237,14 +247,70 @@ func (f *fakeEC2Client) WaitUntilInstanceRunningWithContext(ctx context.Context,
 	return nil
 }
 
-func fakeClient() *AWSClient {
-	return &AWSClient{
-		ec2Client: newFakeEC2Client(),
+func (f *fakeEC2Client) DescribeInstanceTypesPagesWithContext(ctx context.Context, input *ec2.DescribeInstanceTypesInput, fn func(*ec2.DescribeInstanceTypesOutput, bool) bool, opt ...request.Option) error {
+	if ctx == nil || input == nil || fn == nil {
+		return errors.New("invalid input")
+	}
+	f.mu.RLock()
+	defer f.mu.RUnlock()
+	for it, its := range f.instanceTypes {
+		fn(&ec2.DescribeInstanceTypesOutput{
+			InstanceTypes: []*ec2.InstanceTypeInfo{its},
+		}, it == len(f.instanceTypes)-1)
+	}
+	return nil
+}
+
+func (f *fakeEC2Client) GetServiceQuota(input *servicequotas.GetServiceQuotaInput) (*servicequotas.GetServiceQuotaOutput, error) {
+	if input == nil || input.QuotaCode == nil || input.ServiceCode == nil {
+		return nil, request.ErrInvalidParams{}
+	}
+	v, ok := f.serviceQuota[aws.StringValue(input.ServiceCode)+"-"+aws.StringValue(input.QuotaCode)]
+	if !ok {
+		return nil, errors.New("quota not found")
+	}
+	return &servicequotas.GetServiceQuotaOutput{
+		Quota: &servicequotas.ServiceQuota{
+			Value: aws.Float64(v),
+		},
+	}, nil
+}
+
+type option func(*fakeEC2Client)
+
+func WithServiceQuota(service, quota string, value float64) option {
+	return func(c *fakeEC2Client) {
+		c.serviceQuota[service+"-"+quota] = value
 	}
 }
 
-func fakeClientWithInstances(t *testing.T, count int) (*AWSClient, []*Instance) {
-	c := fakeClient()
+func WithInstanceType(name, arch string, numCPU int64) option {
+	return func(c *fakeEC2Client) {
+		c.instanceTypes = append(c.instanceTypes, &ec2.InstanceTypeInfo{
+			InstanceType: aws.String(name),
+			ProcessorInfo: &ec2.ProcessorInfo{
+				SupportedArchitectures: []*string{aws.String(arch)},
+			},
+			VCpuInfo: &ec2.VCpuInfo{
+				DefaultVCpus: aws.Int64(numCPU),
+			},
+		})
+	}
+}
+
+func fakeClient(opts ...option) *AWSClient {
+	fc := newFakeAWSClient()
+	for _, opt := range opts {
+		opt(fc)
+	}
+	return &AWSClient{
+		ec2Client:   fc,
+		quotaClient: fc,
+	}
+}
+
+func fakeClientWithInstances(t *testing.T, count int, opts ...option) (*AWSClient, []*Instance) {
+	c := fakeClient(opts...)
 	ctx := context.Background()
 	insts := make([]*Instance, 0, count)
 	for i := 0; i < count; i++ {
@@ -301,6 +367,50 @@ func TestRunningInstances(t *testing.T) {
 	})
 }
 
+func TestInstanceTypesARM(t *testing.T) {
+	opts := []option{
+		WithInstanceType("zz.large", "x86_64", 10),
+		WithInstanceType("aa.xlarge", "arm64", 20),
+	}
+
+	t.Run("query-arm64-instances", func(t *testing.T) {
+		c := fakeClient(opts...)
+		gotInstTypes, gotErr := c.InstanceTypesARM(context.Background())
+		if gotErr != nil {
+			t.Fatalf("InstanceTypesArm(ctx) = %+v, %s; want nil, nil", gotInstTypes, gotErr)
+		}
+		if len(gotInstTypes) != 1 {
+			t.Errorf("got instance type count %d: want %d", len(gotInstTypes), 1)
+		}
+	})
+	t.Run("nil-request", func(t *testing.T) {
+		c := fakeClient(opts...)
+		gotInstTypes, gotErr := c.InstanceTypesARM(nil)
+		if gotErr == nil {
+			t.Fatalf("InstanceTypesArm(nil) = %+v, %s; want nil, error", gotInstTypes, gotErr)
+		}
+	})
+}
+
+func TestQuota(t *testing.T) {
+	t.Run("on-demand-vcpu", func(t *testing.T) {
+		wantQuota := int64(384)
+		c := fakeClient(WithServiceQuota(QuotaServiceEC2, QuotaCodeCPUOnDemand, float64(wantQuota)))
+		gotQuota, gotErr := c.Quota(context.Background(), QuotaServiceEC2, QuotaCodeCPUOnDemand)
+		if gotErr != nil || wantQuota != gotQuota {
+			t.Fatalf("Quota(ctx, %s, %s) = %+v, %s; want %d, nil", QuotaServiceEC2, QuotaCodeCPUOnDemand, gotQuota, gotErr, wantQuota)
+		}
+	})
+	t.Run("nil-request", func(t *testing.T) {
+		wantQuota := int64(384)
+		c := fakeClient(WithServiceQuota(QuotaServiceEC2, QuotaCodeCPUOnDemand, float64(wantQuota)))
+		gotQuota, gotErr := c.Quota(context.Background(), "", "")
+		if gotErr == nil || gotQuota != 0 {
+			t.Fatalf("Quota(ctx, %s, %s) = %+v, %s; want 0, error", QuotaServiceEC2, QuotaCodeCPUOnDemand, gotQuota, gotErr)
+		}
+	})
+}
+
 func TestInstance(t *testing.T) {
 	t.Run("query-instance", func(t *testing.T) {
 		c, wantInsts := fakeClientWithInstances(t, 1)

internal/cloud/fake_aws.go

@@ -21,14 +21,30 @@ func init() { mrand.Seed(time.Now().UnixNano()) }
 // FakeAWSClient provides a fake AWS Client used to test the AWS client
 // functionality.
 type FakeAWSClient struct {
-	mu        sync.RWMutex
-	instances map[string]*Instance
+	mu            sync.RWMutex
+	instances     map[string]*Instance
+	instanceTypes []*InstanceType
+	serviceQuotas map[serviceQuotaKey]int64
+}
+
+// serviceQuotaKey should be used as the key in the serviceQuotas map.
+type serviceQuotaKey struct {
+	code    string
+	service string
 }
 
 // NewFakeAWSClient crates a fake AWS client.
 func NewFakeAWSClient() *FakeAWSClient {
 	return &FakeAWSClient{
 		instances: make(map[string]*Instance),
+		instanceTypes: []*InstanceType{
+			&InstanceType{"ab.large", 10},
+			&InstanceType{"ab.xlarge", 20},
+			&InstanceType{"ab.small", 30},
+		},
+		serviceQuotas: map[serviceQuotaKey]int64{
+			serviceQuotaKey{QuotaCodeCPUOnDemand, QuotaServiceEC2}: 384,
+		},
 	}
 }
 
@@ -67,6 +83,36 @@ func (f *FakeAWSClient) RunningInstances(ctx context.Context) ([]*Instance, erro
 	return instances, nil
 }
 
+// InstanceTypesArm retrieves all EC2 instance types in a region which support the ARM64 architecture.
+func (f *FakeAWSClient) InstanceTypesARM(ctx context.Context) ([]*InstanceType, error) {
+	if ctx == nil {
+		return nil, errors.New("invalid params")
+	}
+	f.mu.RLock()
+	defer f.mu.RUnlock()
+
+	instanceTypes := make([]*InstanceType, 0, len(f.instanceTypes))
+	for _, it := range f.instanceTypes {
+		instanceTypes = append(instanceTypes, &InstanceType{it.Type, it.CPU})
+	}
+	return instanceTypes, nil
+}
+
+// Quota retrieves the requested service quota for the service.
+func (f *FakeAWSClient) Quota(ctx context.Context, service, code string) (int64, error) {
+	if ctx == nil || service == "" || code == "" {
+		return 0, errors.New("invalid params")
+	}
+	f.mu.RLock()
+	defer f.mu.RUnlock()
+
+	v, ok := f.serviceQuotas[serviceQuotaKey{code, service}]
+	if !ok {
+		return 0, errors.New("service quota not found")
+	}
+	return v, nil
+}
+
 // CreateInstance creates an EC2 VM instance.
 func (f *FakeAWSClient) CreateInstance(ctx context.Context, config *EC2VMConfiguration) (*Instance, error) {
 	if ctx == nil || config == nil {

internal/cloud/fake_aws_test.go

@@ -165,6 +165,62 @@ func TestFakeAWSClientRunningInstances(t *testing.T) {
 	})
 }
 
+func TestFakeAWSClientInstanceTypesARM(t *testing.T) {
+	t.Run("invalid-params", func(t *testing.T) {
+		f := NewFakeAWSClient()
+		if gotITs, gotErr := f.InstanceTypesARM(nil); gotErr == nil {
+			t.Errorf("InstanceTypesARM(nil) = %+v, nil, want error", gotITs)
+		}
+	})
+	t.Run("no-instances", func(t *testing.T) {
+		ctx := context.Background()
+		f := NewFakeAWSClient()
+		gotITs, gotErr := f.InstanceTypesARM(ctx)
+		if gotErr != nil {
+			t.Errorf("InstanceTypesARM(ctx) error = %v, no error", gotErr)
+		}
+		if !cmp.Equal(gotITs, f.instanceTypes) {
+			t.Errorf("InstanceTypesARM(ctx) = %+v, %s; want %+v", gotITs, gotErr, f.instanceTypes)
+		}
+	})
+}
+
+func TestFakeAWSClientQuota(t *testing.T) {
+	t.Run("invalid-context", func(t *testing.T) {
+		f := NewFakeAWSClient()
+		gotQuota, gotErr := f.Quota(nil, QuotaServiceEC2, QuotaCodeCPUOnDemand)
+		if gotErr == nil || gotQuota != 0 {
+			t.Errorf("Quota(nil, %s, %s) = %d, %s, want error", QuotaServiceEC2, QuotaCodeCPUOnDemand, gotQuota, gotErr)
+		}
+	})
+	t.Run("invalid-service", func(t *testing.T) {
+		f := NewFakeAWSClient()
+		gotQuota, gotErr := f.Quota(context.Background(), "", QuotaCodeCPUOnDemand)
+		if gotErr == nil || gotQuota != 0 {
+			t.Errorf("Quota(ctx, \"\", %s) = %d, %s, want error", QuotaCodeCPUOnDemand, gotQuota, gotErr)
+		}
+	})
+	t.Run("invalid-quota-code", func(t *testing.T) {
+		f := NewFakeAWSClient()
+		gotQuota, gotErr := f.Quota(context.Background(), QuotaServiceEC2, "")
+		if gotErr == nil || gotQuota != 0 {
+			t.Errorf("Quota(ctx, %s, \"\") = %d, %s, want error", QuotaServiceEC2, gotQuota, gotErr)
+		}
+	})
+	t.Run("valid-request", func(t *testing.T) {
+		f := NewFakeAWSClient()
+		wantQuota, ok := f.serviceQuotas[serviceQuotaKey{QuotaCodeCPUOnDemand, QuotaServiceEC2}]
+		if !ok {
+			t.Fatal("unable to retrieve quota value")
+		}
+		gotQuota, gotErr := f.Quota(context.Background(), QuotaServiceEC2, QuotaCodeCPUOnDemand)
+		if gotErr != nil || gotQuota != wantQuota {
+			t.Errorf("Quota(ctx, %s, %s) = %d, %s, want %d, nil", QuotaServiceEC2,
+				QuotaCodeCPUOnDemand, gotQuota, gotErr, wantQuota)
+		}
+	})
+}
+
 func TestFakeAWSClientCreateInstance(t *testing.T) {
 	t.Run("create-instance", func(t *testing.T) {
 		ctx := context.Background()