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>
diff --git a/internal/cloud/aws.go b/internal/cloud/aws.go
index 2cdca87..008660a 100644
--- a/internal/cloud/aws.go
+++ b/internal/cloud/aws.go
@@ -18,6 +18,7 @@
"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 @@
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 @@
// 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 @@
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 @@
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))
diff --git a/internal/cloud/aws_test.go b/internal/cloud/aws_test.go
index a90a73e..4861ddf 100644
--- a/internal/cloud/aws_test.go
+++ b/internal/cloud/aws_test.go
@@ -17,20 +17,30 @@
"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 @@
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 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)
diff --git a/internal/cloud/fake_aws.go b/internal/cloud/fake_aws.go
index 2a47b2c..9f7aa9e 100644
--- a/internal/cloud/fake_aws.go
+++ b/internal/cloud/fake_aws.go
@@ -21,14 +21,30 @@
// 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 @@
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 {
diff --git a/internal/cloud/fake_aws_test.go b/internal/cloud/fake_aws_test.go
index 5d78399..38f791d 100644
--- a/internal/cloud/fake_aws_test.go
+++ b/internal/cloud/fake_aws_test.go
@@ -165,6 +165,62 @@
})
}
+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()
