aws/credentials: Update Credentials cache to have less overhead

aws/credentials.go

@@ -3,6 +3,7 @@ package aws
 import (
 	"math"
 	"sync"
+	"sync/atomic"
 	"time"
 
 	"github.com/aws/aws-sdk-go-v2/internal/sdk"
@@ -85,7 +86,7 @@ type CredentialsProvider interface {
 type SafeCredentialsProvider struct {
 	RetrieveFn func() (Credentials, error)
 
-	creds Credentials
+	creds atomic.Value
 	m     sync.Mutex
 }
 
@@ -95,28 +96,43 @@ type SafeCredentialsProvider struct {
 //
 // Retruns and error if RetrieveFn returns an error.
 func (p *SafeCredentialsProvider) Retrieve() (Credentials, error) {
+	if creds := p.getCreds(); creds != nil {
+		return *creds, nil
+	}
+
 	p.m.Lock()
 	defer p.m.Unlock()
 
-	if p.creds.HasKeys() && !p.creds.Expired() {
-		return p.creds, nil
+	// Make sure another goroutine didn't already update the credentials.
+	if creds := p.getCreds(); creds != nil {
+		return *creds, nil
 	}
 
 	creds, err := p.RetrieveFn()
 	if err != nil {
 		return Credentials{}, err
 	}
+	p.creds.Store(&creds)
+
+	return creds, nil
+}
+
+func (p *SafeCredentialsProvider) getCreds() *Credentials {
+	v := p.creds.Load()
+	if v == nil {
+		return nil
+	}
 
-	p.creds = creds
+	c := v.(*Credentials)
+	if c != nil && c.HasKeys() && !c.Expired() {
+		return c
+	}
 
-	return p.creds, nil
+	return nil
 }
 
 // Invalidate will invalidate the cached credentials. The next call to Retrieve
 // will cause RetrieveFn to be called.
 func (p *SafeCredentialsProvider) Invalidate() {
-	p.m.Lock()
-	defer p.m.Unlock()
-
-	p.creds = Credentials{}
+	p.creds.Store((*Credentials)(nil))
 }

aws/credentials_bench_test.go

@@ -0,0 +1,87 @@
+package aws
+
+import (
+	"fmt"
+	"strconv"
+	"sync"
+	"testing"
+	"time"
+)
+
+func BenchmarkSafeCredentialsProvider_Retrieve(b *testing.B) {
+	retrieveFn := func() (Credentials, error) {
+		return Credentials{
+			AccessKeyID:     "key",
+			SecretAccessKey: "secret",
+			Source:          "benchmark",
+		}, nil
+	}
+
+	cases := []int{1, 10, 100, 500, 1000, 10000}
+	for _, c := range cases {
+		b.Run(strconv.Itoa(c), func(b *testing.B) {
+			p := SafeCredentialsProvider{
+				RetrieveFn: retrieveFn,
+			}
+			var wg sync.WaitGroup
+			wg.Add(c)
+			for i := 0; i < c; i++ {
+				go func() {
+					for j := 0; j < b.N; j++ {
+						v, err := p.Retrieve()
+						if err != nil {
+							b.Fatalf("expect no error %v, %v", v, err)
+						}
+					}
+					wg.Done()
+				}()
+			}
+			b.ResetTimer()
+
+			wg.Wait()
+		})
+	}
+}
+
+func BenchmarkSafeCredentialsProvider_Retrieve_Invalidate(b *testing.B) {
+	retrieveFn := func() (Credentials, error) {
+		time.Sleep(time.Millisecond)
+		return Credentials{
+			AccessKeyID:     "key",
+			SecretAccessKey: "secret",
+			Source:          "benchmark",
+		}, nil
+	}
+
+	expRates := []int{10000, 1000, 100}
+	cases := []int{1, 10, 100, 500, 1000, 10000}
+	for _, expRate := range expRates {
+		for _, c := range cases {
+			b.Run(fmt.Sprintf("%d-%d", expRate, c), func(b *testing.B) {
+				p := SafeCredentialsProvider{
+					RetrieveFn: retrieveFn,
+				}
+				var wg sync.WaitGroup
+				wg.Add(c)
+				for i := 0; i < c; i++ {
+					go func(id int) {
+						for j := 0; j < b.N; j++ {
+							v, err := p.Retrieve()
+							if err != nil {
+								b.Fatalf("expect no error %v, %v", v, err)
+							}
+							// periodically expire creds to cause rwlock
+							if id == 0 && j%expRate == 0 {
+								p.Invalidate()
+							}
+						}
+						wg.Done()
+					}(i)
+				}
+				b.ResetTimer()
+
+				wg.Wait()
+			})
+		}
+	}
+}