X-Git-Url: https://git.arvados.org/arvados.git/blobdiff_plain/267d3c40bf1c5503e6487db2ab1f6a4339ac5f83..2e727c5d2d000faa6f1d9a566dc59568f1b276fe:/lib/cloud/ec2/ec2.go

diff --git a/lib/cloud/ec2/ec2.go b/lib/cloud/ec2/ec2.go
index 071c95006c..52b73f781c 100644
--- a/lib/cloud/ec2/ec2.go
+++ b/lib/cloud/ec2/ec2.go
@@ -20,6 +20,7 @@ import (
 	"git.arvados.org/arvados.git/lib/cloud"
 	"git.arvados.org/arvados.git/sdk/go/arvados"
 	"github.com/aws/aws-sdk-go/aws"
+	"github.com/aws/aws-sdk-go/aws/awserr"
 	"github.com/aws/aws-sdk-go/aws/credentials"
 	"github.com/aws/aws-sdk-go/aws/credentials/ec2rolecreds"
 	"github.com/aws/aws-sdk-go/aws/ec2metadata"
@@ -39,13 +40,14 @@ const (
 )
 
 type ec2InstanceSetConfig struct {
-	AccessKeyID      string
-	SecretAccessKey  string
-	Region           string
-	SecurityGroupIDs arvados.StringSet
-	SubnetID         string
-	AdminUsername    string
-	EBSVolumeType    string
+	AccessKeyID        string
+	SecretAccessKey    string
+	Region             string
+	SecurityGroupIDs   arvados.StringSet
+	SubnetID           string
+	AdminUsername      string
+	EBSVolumeType      string
+	IAMInstanceProfile string
 }
 
 type ec2Interface interface {
@@ -229,6 +231,12 @@ func (instanceSet *ec2InstanceSet) Create(
 			}}
 	}
 
+	if instanceSet.ec2config.IAMInstanceProfile != "" {
+		rii.IamInstanceProfile = &ec2.IamInstanceProfileSpecification{
+			Name: aws.String(instanceSet.ec2config.IAMInstanceProfile),
+		}
+	}
+
 	rsv, err := instanceSet.client.RunInstances(&rii)
 	err = wrapError(err, &instanceSet.throttleDelayCreate)
 	if err != nil {
@@ -349,6 +357,31 @@ func (err rateLimitError) EarliestRetry() time.Time {
 	return err.earliestRetry
 }
 
+var isCodeCapacity = map[string]bool{
+	"InsufficientInstanceCapacity": true,
+	"VcpuLimitExceeded":            true,
+	"MaxSpotInstanceCountExceeded": true,
+}
+
+// isErrorCapacity returns whether the error is to be throttled based on its code.
+// Returns false if error is nil.
+func isErrorCapacity(err error) bool {
+	if aerr, ok := err.(awserr.Error); ok && aerr != nil {
+		if _, ok := isCodeCapacity[aerr.Code()]; ok {
+			return true
+		}
+	}
+	return false
+}
+
+type ec2QuotaError struct {
+	error
+}
+
+func (er *ec2QuotaError) IsQuotaError() bool {
+	return true
+}
+
 func wrapError(err error, throttleValue *atomic.Value) error {
 	if request.IsErrorThrottle(err) {
 		// Back off exponentially until an upstream call
@@ -362,6 +395,8 @@ func wrapError(err error, throttleValue *atomic.Value) error {
 		}
 		throttleValue.Store(d)
 		return rateLimitError{error: err, earliestRetry: time.Now().Add(d)}
+	} else if isErrorCapacity(err) {
+		return &ec2QuotaError{err}
 	} else if err != nil {
 		throttleValue.Store(time.Duration(0))
 		return err