From 6bb5a84a53e5810e96e56e41cc751d4ebc054580 Mon Sep 17 00:00:00 2001
From: Ward Vandewege <ward@curii.com>
Date: Tue, 15 Jun 2021 15:40:59 -0400
Subject: [PATCH] 17776: address review comments.

Arvados-DCO-1.1-Signed-off-by: Ward Vandewege <ward@curii.com>
---
 lib/cloud/ec2/ec2.go      | 29 ++++++++++++++++++-----------
 lib/cloud/ec2/ec2_test.go | 14 +++++++++++++-
 2 files changed, 31 insertions(+), 12 deletions(-)

diff --git a/lib/cloud/ec2/ec2.go b/lib/cloud/ec2/ec2.go
index fb7afdda42..269a7d8def 100644
--- a/lib/cloud/ec2/ec2.go
+++ b/lib/cloud/ec2/ec2.go
@@ -350,28 +350,33 @@ func (err rateLimitError) EarliestRetry() time.Time {
 	return err.earliestRetry
 }
 
-var capacityCodes = map[string]struct{}{
-	"InsufficientInstanceCapacity": {},
-	"VcpuLimitExceeded":            {},
-	"MaxSpotInstanceCountExceeded": {},
+var isCodeCapacity = map[string]bool{
+	"InsufficientInstanceCapacity": true,
+	"VcpuLimitExceeded":            true,
+	"MaxSpotInstanceCountExceeded": true,
 }
 
-// IsErrorCapacity returns whether the error is to be throttled based on its code.
+// isErrorCapacity returns whether the error is to be throttled based on its code.
 // Returns false if error is nil.
-func IsErrorCapacity(err error) bool {
+func isErrorCapacity(err error) bool {
 	if aerr, ok := err.(awserr.Error); ok && aerr != nil {
-		return isCodeCapacity(aerr.Code())
+		if _, ok := isCodeCapacity[aerr.Code()]; ok {
+			return true
+		}
 	}
 	return false
 }
 
-func isCodeCapacity(code string) bool {
-	_, ok := capacityCodes[code]
-	return ok
+type ec2QuotaError struct {
+	error
+}
+
+func (er *ec2QuotaError) IsQuotaError() bool {
+	return true
 }
 
 func wrapError(err error, throttleValue *atomic.Value) error {
-	if request.IsErrorThrottle(err) || IsErrorCapacity(err) {
+	if request.IsErrorThrottle(err) {
 		// Back off exponentially until an upstream call
 		// either succeeds or returns a non-throttle error.
 		d, _ := throttleValue.Load().(time.Duration)
@@ -383,6 +388,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
diff --git a/lib/cloud/ec2/ec2_test.go b/lib/cloud/ec2/ec2_test.go
index e7319a0cb6..3cd238ded5 100644
--- a/lib/cloud/ec2/ec2_test.go
+++ b/lib/cloud/ec2/ec2_test.go
@@ -25,6 +25,7 @@ package ec2
 import (
 	"encoding/json"
 	"flag"
+	"sync/atomic"
 	"testing"
 
 	"git.arvados.org/arvados.git/lib/cloud"
@@ -32,6 +33,7 @@ import (
 	"git.arvados.org/arvados.git/sdk/go/arvados"
 	"git.arvados.org/arvados.git/sdk/go/config"
 	"github.com/aws/aws-sdk-go/aws"
+	"github.com/aws/aws-sdk-go/aws/awserr"
 	"github.com/aws/aws-sdk-go/service/ec2"
 	"github.com/sirupsen/logrus"
 	check "gopkg.in/check.v1"
@@ -246,4 +248,14 @@ func (*EC2InstanceSetSuite) TestDestroyInstances(c *check.C) {
 	}
 }
 
-var TestRateLimitErrorInterface cloud.RateLimitError = rateLimitError{}
+func (*EC2InstanceSetSuite) TestWrapError(c *check.C) {
+	retryError := awserr.New("Throttling", "", nil)
+	wrapped := wrapError(retryError, &atomic.Value{})
+	_, ok := wrapped.(cloud.RateLimitError)
+	c.Check(ok, check.Equals, true)
+
+	quotaError := awserr.New("InsufficientInstanceCapacity", "", nil)
+	wrapped = wrapError(quotaError, nil)
+	_, ok = wrapped.(cloud.QuotaError)
+	c.Check(ok, check.Equals, true)
+}
-- 
2.30.2