Add ExpiryWindow and ExpiryWindowJitterFrac to CredentialsCache

aws/credential_cache.go

@@ -3,17 +3,74 @@ package aws
 import (
 	"context"
 	"sync/atomic"
+	"time"
 
+	sdkrand "github.com/aws/aws-sdk-go-v2/internal/rand"
 	"github.com/aws/aws-sdk-go-v2/internal/sync/singleflight"
 )
 
+// CredentialsCacheOptions are the options
+type CredentialsCacheOptions struct {
+	// Provider is the CredentialProvider implementation to be wrapped by the CredentialCache.
+	Provider CredentialsProvider
+
+	// ExpiryWindow will allow the credentials to trigger refreshing prior to
+	// the credentials actually expiring. This is beneficial so race conditions
+	// with expiring credentials do not cause request to fail unexpectedly
+	// due to ExpiredTokenException exceptions.
+	//
+	// An ExpiryWindow of 10s would cause calls to IsExpired() to return true
+	// 10 seconds before the credentials are actually expired.
+	//
+	// If ExpiryWindow is 0 or less it will be ignored.
+	ExpiryWindow time.Duration
+
+	// ExpiryWindowJitterFrac provides a mechanism for randomizing the expiration of credentials
+	// within the configured ExpiryWindow by a random percentage. Valid values are between 0.0 and 1.0.
+	//
+	// As an example if ExpiryWindow is 60 seconds and ExpiryWindowJitterFrac is 0.5 then credentials will be set to
+	// expire between 30 to 60 seconds prior to their actual expiration time.
+	//
+	// If ExpiryWindow is 0 or less then ExpiryWindowJitterFrac is ignored.
+	// If ExpiryWindowJitterFrac is 0 then no randomization will be applied to the window.
+	// If ExpiryWindowJitterFrac < 0 the value will be treated as 0.
+	// If ExpiryWindowJitterFrac > 1 the value will be treated as 1.
+	ExpiryWindowJitterFrac float64
+}
+
 // CredentialsCache provides caching and concurrency safe credentials retrieval
 // via the provider's retrieve method.
 type CredentialsCache struct {
-	Provider CredentialsProvider
+	options CredentialsCacheOptions
+	creds   atomic.Value
+	sf      singleflight.Group
+}
+
+// NewCredentialsCache returns a CredentialsCache that wraps provider. Provider is expected to not be nil. A variadic
+// list of one or more functions can be provided to modify the CredentialsCache configuration. This allows for
+// configuration of credential expiry window and jitter.
+func NewCredentialsCache(provider CredentialsProvider, optFns ...func(options *CredentialsCacheOptions)) *CredentialsCache {
+	options := CredentialsCacheOptions{
+		Provider: provider,
+	}
 
-	creds atomic.Value
-	sf    singleflight.Group
+	for _, fn := range optFns {
+		fn(&options)
+	}
+
+	if options.ExpiryWindow < 0 {
+		options.ExpiryWindow = 0
+	}
+
+	if options.ExpiryWindowJitterFrac < 0 {
+		options.ExpiryWindowJitterFrac = 0
+	} else if options.ExpiryWindowJitterFrac > 1 {
+		options.ExpiryWindowJitterFrac = 1
+	}
+
+	return &CredentialsCache{
+		options: options,
+	}
 }
 
 // Retrieve returns the credentials. If the credentials have already been
@@ -41,8 +98,17 @@ func (p *CredentialsCache) singleRetrieve() (interface{}, error) {
 		return *creds, nil
 	}
 
-	creds, err := p.Provider.Retrieve(context.TODO())
+	creds, err := p.options.Provider.Retrieve(context.TODO())
 	if err == nil {
+		if creds.CanExpire {
+			randFloat64, err := sdkrand.CryptoRandFloat64()
+			if err != nil {
+				return Credentials{}, err
+			}
+			jitter := time.Duration(randFloat64 * p.options.ExpiryWindowJitterFrac * float64(p.options.ExpiryWindow))
+			creds.Expires = creds.Expires.Add(-(p.options.ExpiryWindow - jitter))
+		}
+
 		p.creds.Store(&creds)
 	}
 

aws/credential_cache_bench_test.go

@@ -21,9 +21,7 @@ func BenchmarkCredentialsCache_Retrieve(b *testing.B) {
 	cases := []int{1, 10, 100, 500, 1000, 10000}
 	for _, c := range cases {
 		b.Run(strconv.Itoa(c), func(b *testing.B) {
-			p := CredentialsCache{
-				Provider: provider,
-			}
+			p := NewCredentialsCache(provider)
 			var wg sync.WaitGroup
 			wg.Add(c)
 			for i := 0; i < c; i++ {
@@ -59,9 +57,7 @@ func BenchmarkCredentialsCache_Retrieve_Invalidate(b *testing.B) {
 	for _, expRate := range expRates {
 		for _, c := range cases {
 			b.Run(fmt.Sprintf("%d-%d", expRate, c), func(b *testing.B) {
-				p := CredentialsCache{
-					Provider: provider,
-				}
+				p := NewCredentialsCache(provider)
 				var wg sync.WaitGroup
 				wg.Add(c)
 				for i := 0; i < c; i++ {

aws/credential_cache_test.go

@@ -42,15 +42,13 @@ func TestCredentialsCache_Cache(t *testing.T) {
 	}
 
 	var called bool
-	p := &CredentialsCache{
-		Provider: CredentialsProviderFunc(func(ctx context.Context) (Credentials, error) {
-			if called {
-				t.Fatalf("expect provider.Retrieve to only be called once")
-			}
-			called = true
-			return expect, nil
-		}),
-	}
+	p := NewCredentialsCache(CredentialsProviderFunc(func(ctx context.Context) (Credentials, error) {
+		if called {
+			t.Fatalf("expect provider.Retrieve to only be called once")
+		}
+		called = true
+		return expect, nil
+	}))
 
 	for i := 0; i < 2; i++ {
 		creds, err := p.Retrieve(context.Background())
@@ -108,12 +106,10 @@ func TestCredentialsCache_Expires(t *testing.T) {
 
 	for _, c := range cases {
 		var called int
-		p := &CredentialsCache{
-			Provider: CredentialsProviderFunc(func(ctx context.Context) (Credentials, error) {
-				called++
-				return c.Creds(), nil
-			}),
-		}
+		p := NewCredentialsCache(CredentialsProviderFunc(func(ctx context.Context) (Credentials, error) {
+			called++
+			return c.Creds(), nil
+		}))
 
 		p.Retrieve(context.Background())
 		p.Retrieve(context.Background())
@@ -131,13 +127,92 @@ func TestCredentialsCache_Expires(t *testing.T) {
 	}
 }
 
-func TestCredentialsCache_Error(t *testing.T) {
-	p := &CredentialsCache{
-		Provider: CredentialsProviderFunc(func(ctx context.Context) (Credentials, error) {
-			return Credentials{}, fmt.Errorf("failed")
-		}),
+func TestCredentialsCache_ExpireTime(t *testing.T) {
+	orig := sdk.NowTime
+	defer func() { sdk.NowTime = orig }()
+	var mockTime time.Time
+	sdk.NowTime = func() time.Time { return mockTime }
+
+	cases := map[string]struct {
+		ExpireTime   time.Time
+		ExpiryWindow time.Duration
+		JitterFrac   float64
+		Validate     func(t *testing.T, v time.Time)
+	}{
+		"no expire window": {
+			Validate: func(t *testing.T, v time.Time) {
+				t.Helper()
+				if e, a := mockTime, v; !e.Equal(a) {
+					t.Errorf("expect %v, got %v", e, a)
+				}
+			},
+		},
+		"expire window": {
+			ExpireTime:   mockTime.Add(100),
+			ExpiryWindow: 50,
+			Validate: func(t *testing.T, v time.Time) {
+				t.Helper()
+				if e, a := mockTime.Add(50), v; !e.Equal(a) {
+					t.Errorf("expect %v, got %v", e, a)
+				}
+			},
+		},
+		"expire window with jitter": {
+			ExpireTime:   mockTime.Add(100),
+			JitterFrac:   0.5,
+			ExpiryWindow: 50,
+			Validate: func(t *testing.T, v time.Time) {
+				t.Helper()
+				max := mockTime.Add(75)
+				min := mockTime.Add(50)
+				if v.Before(min) {
+					t.Errorf("expect %v to be before %s", v, min)
+				}
+				if v.After(max) {
+					t.Errorf("expect %v to be after %s", v, max)
+				}
+			},
+		},
+		"no expire window with jitter": {
+			ExpireTime:   mockTime,
+			JitterFrac:   0.5,
+			Validate: func(t *testing.T, v time.Time) {
+				t.Helper()
+				if e, a := mockTime, v; !e.Equal(a) {
+					t.Errorf("expect %v, got %v", e, a)
+				}
+			},
+		},
 	}
 
+	for name, tt := range cases {
+		t.Run(name, func(t *testing.T) {
+			p := NewCredentialsCache(CredentialsProviderFunc(func(ctx context.Context) (Credentials, error) {
+				return Credentials{
+					AccessKeyID:     "accessKey",
+					SecretAccessKey: "secretKey",
+					CanExpire:       true,
+					Expires:         tt.ExpireTime,
+				}, nil
+			}), func(options *CredentialsCacheOptions) {
+				options.ExpiryWindow = tt.ExpiryWindow
+				options.ExpiryWindowJitterFrac = tt.JitterFrac
+			})
+
+			credentials, err := p.Retrieve(context.Background())
+			if err != nil {
+				t.Fatalf("expect no error, got %v", err)
+			}
+			tt.Validate(t, credentials.Expires)
+		})
+	}
+}
+
+func TestCredentialsCache_Error(t *testing.T) {
+	p := NewCredentialsCache(CredentialsProviderFunc(func(ctx context.Context) (Credentials, error) {
+		return Credentials{}, fmt.Errorf("failed")
+	}))
+
 	creds, err := p.Retrieve(context.Background())
 	if err == nil {
 		t.Fatalf("expect error, not none")
@@ -156,16 +231,14 @@ func TestCredentialsCache_Race(t *testing.T) {
 		SecretAccessKey: "secret",
 	}
 	var called bool
-	p := &CredentialsCache{
-		Provider: CredentialsProviderFunc(func(ctx context.Context) (Credentials, error) {
-			time.Sleep(time.Duration(rand.Intn(10)) * time.Millisecond)
-			if called {
-				t.Fatalf("expect provider.Retrieve only called once")
-			}
-			called = true
-			return expect, nil
-		}),
-	}
+	p := NewCredentialsCache(CredentialsProviderFunc(func(ctx context.Context) (Credentials, error) {
+		time.Sleep(time.Duration(rand.Intn(10)) * time.Millisecond)
+		if called {
+			t.Fatalf("expect provider.Retrieve only called once")
+		}
+		called = true
+		return expect, nil
+	}))
 
 	var wg sync.WaitGroup
 	wg.Add(100)
@@ -206,9 +279,7 @@ func TestCredentialsCache_RetrieveConcurrent(t *testing.T) {
 	stub := &stubConcurrentProvider{
 		done: make(chan struct{}),
 	}
-	provider := CredentialsCache{
-		Provider: stub,
-	}
+	provider := NewCredentialsCache(stub)
 
 	var wg sync.WaitGroup
 	wg.Add(2)

credentials/ec2rolecreds/provider.go

@@ -46,17 +46,6 @@ type Options struct {
 	//
 	// If nil, the provider will default to the ec2imds client.
 	Client GetMetadataAPIClient
-
-	// ExpiryWindow will allow the credentials to trigger refreshing prior to
-	// the credentials actually expiring. This is beneficial so race conditions
-	// with expiring credentials do not cause request to fail unexpectedly
-	// due to ExpiredTokenException exceptions.
-	//
-	// So a ExpiryWindow of 10s would cause calls to IsExpired() to return true
-	// 10 seconds before the credentials are actually expired.
-	//
-	// If ExpiryWindow is 0 or less it will be ignored.
-	ExpiryWindow time.Duration
 }
 
 // New returns an initialized Provider value configured to retrieve
@@ -102,7 +91,7 @@ func (p *Provider) Retrieve(ctx context.Context) (aws.Credentials, error) {
 		Source:          ProviderName,
 
 		CanExpire: true,
-		Expires:   roleCreds.Expiration.Add(-p.options.ExpiryWindow),
+		Expires:   roleCreds.Expiration,
 	}
 
 	return creds, nil

credentials/endpointcreds/provider.go

@@ -33,7 +33,6 @@ import (
 	"context"
 	"fmt"
 	"net/http"
-	"time"
 
 	"github.com/aws/aws-sdk-go-v2/aws"
 	"github.com/aws/aws-sdk-go-v2/credentials/endpointcreds/internal/client"
@@ -65,17 +64,6 @@ type HTTPClient interface {
 
 // Options is structure of configurable options for Provider
 type Options struct {
-	// ExpiryWindow will allow the credentials to trigger refreshing prior to
-	// the credentials actually expiring. This is beneficial so race conditions
-	// with expiring credentials do not cause request to fail unexpectedly
-	// due to ExpiredTokenException exceptions.
-	//
-	// So a ExpiryWindow of 10s would cause calls to IsExpired() to return true
-	// 10 seconds before the credentials are actually expired.
-	//
-	// If ExpiryWindow is 0 or less it will be ignored.
-	ExpiryWindow time.Duration
-
 	// Endpoint to retrieve credentials from. Required
 	Endpoint string
 
@@ -134,7 +122,7 @@ func (p *Provider) Retrieve(ctx context.Context) (aws.Credentials, error) {
 
 	if resp.Expiration != nil {
 		creds.CanExpire = true
-		creds.Expires = resp.Expiration.Add(-p.options.ExpiryWindow)
+		creds.Expires = *resp.Expiration
 	}
 
 	return creds, nil

credentials/processcreds/provider.go

@@ -55,18 +55,6 @@ type Provider struct {
 
 // Options is the configuration options for configuring the Provider.
 type Options struct {
-	// ExpiryWindow will allow the credentials to trigger refreshing prior to
-	// the credentials actually expiring. This is beneficial so race conditions
-	// with expiring credentials do not cause request to fail unexpectedly
-	// due to ExpiredTokenException exceptions.
-	//
-	// For example, an ExpiryWindow of 10s would cause calls to the
-	// Credentials.IsExpired() method to return true 10 seconds before the
-	// credentials would of actually expired.
-	//
-	// If ExpiryWindow is 0 or less, it will be ignored.
-	ExpiryWindow time.Duration
-
 	// Timeout limits the time a process can run.
 	Timeout time.Duration
 }
@@ -213,7 +201,7 @@ func (p *Provider) Retrieve(ctx context.Context) (aws.Credentials, error) {
 	// Handle expiration
 	if resp.Expiration != nil {
 		creds.CanExpire = true
-		creds.Expires = (*resp.Expiration).Add(-p.options.ExpiryWindow)
+		creds.Expires = *resp.Expiration
 	}
 
 	return creds, nil