aws/credentials: Add GetWithContext method to Credentials

CHANGELOG_PENDING.md

@@ -1,5 +1,8 @@
 ### SDK Features
 
 ### SDK Enhancements
+* `aws/credentials`: Add support for context when getting credentials.
+ * Adds `GetWithContext` to `Credentials` that allows canceling getting the credentials if the context is canceled, or times out. This fixes an issue where API operations would ignore their provide context when waiting for credentials to refresh.
+ * Related to [#3127](https://github.com/aws/aws-sdk-go/pull/3127).
 
 ### SDK Bugs

aws/context_background_1_5.go

@@ -2,42 +2,8 @@
 
 package aws
 
-import "time"
-
-// An emptyCtx is a copy of the Go 1.7 context.emptyCtx type. This is copied to
-// provide a 1.6 and 1.5 safe version of context that is compatible with Go
-// 1.7's Context.
-//
-// An emptyCtx is never canceled, has no values, and has no deadline. It is not
-// struct{}, since vars of this type must have distinct addresses.
-type emptyCtx int
-
-func (*emptyCtx) Deadline() (deadline time.Time, ok bool) {
- return
-}
-
-func (*emptyCtx) Done() <-chan struct{} {
- return nil
-}
-
-func (*emptyCtx) Err() error {
- return nil
-}
-
-func (*emptyCtx) Value(key interface{}) interface{} {
- return nil
-}
-
-func (e *emptyCtx) String() string {
- switch e {
- case backgroundCtx:
- return "aws.BackgroundContext"
- }
- return "unknown empty Context"
-}
-
-var (
- backgroundCtx = new(emptyCtx)
+import (
+ "github.com/aws/aws-sdk-go/internal/context"
 )
 
 // BackgroundContext returns a context that will never be canceled, has no
@@ -52,5 +18,5 @@ var (
 //
 // See https://golang.org/pkg/context for more information on Contexts.
 func BackgroundContext() Context {
- return backgroundCtx
+ return context.BackgroundCtx
 }

aws/credentials/context_background_go1.5.go

@@ -0,0 +1,22 @@
+// +build !go1.7
+
+package credentials
+
+import (
+ "github.com/aws/aws-sdk-go/internal/context"
+)
+
+// backgroundContext returns a context that will never be canceled, has no
+// values, and no deadline. This context is used by the SDK to provide
+// backwards compatibility with non-context API operations and functionality.
+//
+// Go 1.6 and before:
+// This context function is equivalent to context.Background in the Go stdlib.
+//
+// Go 1.7 and later:
+// The context returned will be the value returned by context.Background()
+//
+// See https://golang.org/pkg/context for more information on Contexts.
+func backgroundContext() Context {
+ return context.BackgroundCtx
+}

aws/credentials/context_background_go1.7.go

@@ -0,0 +1,20 @@
+// +build go1.7
+
+package credentials
+
+import "context"
+
+// backgroundContext returns a context that will never be canceled, has no
+// values, and no deadline. This context is used by the SDK to provide
+// backwards compatibility with non-context API operations and functionality.
+//
+// Go 1.6 and before:
+// This context function is equivalent to context.Background in the Go stdlib.
+//
+// Go 1.7 and later:
+// The context returned will be the value returned by context.Background()
+//
+// See https://golang.org/pkg/context for more information on Contexts.
+func backgroundContext() Context {
+ return context.Background()
+}

aws/credentials/context_go1.5.go

@@ -0,0 +1,39 @@
+// +build !go1.9
+
+package credentials
+
+import "time"
+
+// Context is an copy of the Go v1.7 stdlib's context.Context interface.
+// It is represented as a SDK interface to enable you to use the "WithContext"
+// API methods with Go v1.6 and a Context type such as golang.org/x/net/context.
+//
+// This type, aws.Context, and context.Context are equivalent.
+//
+// See https://golang.org/pkg/context on how to use contexts.
+type Context interface {
+ // Deadline returns the time when work done on behalf of this context
+ // should be canceled. Deadline returns ok==false when no deadline is
+ // set. Successive calls to Deadline return the same results.
+ Deadline() (deadline time.Time, ok bool)
+
+ // Done returns a channel that's closed when work done on behalf of this
+ // context should be canceled. Done may return nil if this context can
+ // never be canceled. Successive calls to Done return the same value.
+ Done() <-chan struct{}
+
+ // Err returns a non-nil error value after Done is closed. Err returns
+ // Canceled if the context was canceled or DeadlineExceeded if the
+ // context's deadline passed. No other values for Err are defined.
+ // After Done is closed, successive calls to Err return the same value.
+ Err() error
+
+ // Value returns the value associated with this context for key, or nil
+ // if no value is associated with key. Successive calls to Value with
+ // the same key returns the same result.
+ //
+ // Use context values only for request-scoped data that transits
+ // processes and API boundaries, not for passing optional parameters to
+ // functions.
+ Value(key interface{}) interface{}
+}

aws/credentials/context_go1.9.go

@@ -0,0 +1,13 @@
+// +build go1.9
+
+package credentials
+
+import "context"
+
+// Context is an alias of the Go stdlib's context.Context interface.
+// It can be used within the SDK's API operation "WithContext" methods.
+//
+// This type, aws.Context, and context.Context are equivalent.
+//
+// See https://golang.org/pkg/context on how to use contexts.
+type Context = context.Context

aws/credentials/credentials.go

@@ -213,34 +213,60 @@ func NewCredentials(provider Provider) *Credentials {
  return c
 }
 
-// Get returns the credentials value, or error if the credentials Value failed
-// to be retrieved.
+// GetWithContext returns the credentials value, or error if the credentials
+// Value failed to be retrieved. Will return early if the passed in context is
+// canceled.
 //
 // Will return the cached credentials Value if it has not expired. If the
 // credentials Value has expired the Provider's Retrieve() will be called
 // to refresh the credentials.
 //
 // If Credentials.Expire() was called the credentials Value will be force
 // expired, and the next call to Get() will cause them to be refreshed.
-func (c *Credentials) Get() (Value, error) {
- if creds := c.creds.Load(); !c.isExpired(creds) {
- return creds.(Value), nil
+//
+// Passed in Context is equivalent to aws.Context, and context.Context.
+func (c *Credentials) GetWithContext(ctx Context) (Value, error) {
+ if curCreds := c.creds.Load(); !c.isExpired(curCreds) {
+ return curCreds.(Value), nil
+ }
+
+ // Cannot pass context down to the actual retrieve, because the first
+ // context would cancel the whole group when there is not direct
+ // association of items in the group.
+ resCh := c.sf.DoChan("", c.singleRetrieve)
+ select {
+ case res := <-resCh:
+ return res.Val.(Value), res.Err
+ case <-ctx.Done():
+ return Value{}, awserr.New("RequestCanceled",
+ "request context canceled", ctx.Err())
  }
+}
 
- creds, err, _ := c.sf.Do("", func() (interface{}, error) {
-  if creds := c.creds.Load(); !c.isExpired(creds) {
-  return creds.(Value), nil
-  }
+func (c *Credentials) singleRetrieve() (interface{}, error) {
+ if curCreds := c.creds.Load(); !c.isExpired(curCreds) {
+ return curCreds.(Value), nil
+ }
 
-  creds, err := c.provider.Retrieve()
-  if err == nil {
-  c.creds.Store(creds)
-  }
+ creds, err := c.provider.Retrieve()
+ if err == nil {
+ c.creds.Store(creds)
+ }
 
-  return creds, err
- })
+ return creds, err
+}
 
- return creds.(Value), err
+// Get returns the credentials value, or error if the credentials Value failed
+// to be retrieved.
+//
+// Will return the cached credentials Value if it has not expired. If the
+// credentials Value has expired the Provider's Retrieve() will be called
+// to refresh the credentials.
+//
+// If Credentials.Expire() was called the credentials Value will be force
+// expired, and the next call to Get() will cause them to be refreshed.
+func (c *Credentials) Get() (Value, error) {
+ return c.GetWithContext(backgroundContext())
 }
 
 // Expire expires the credentials and forces them to be retrieved on the

aws/credentials/credentials_go1.7_test.go

@@ -0,0 +1,77 @@
+// +build go1.7
+
+package credentials
+
+import (
+ "context"
+ "testing"
+)
+
+func TestCredentialsGetWithContext(t *testing.T) {
+ stub := &stubProviderConcurrent{
+ stubProvider: stubProvider{
+ creds: Value{
+ AccessKeyID: "AKIDEXAMPLE",
+ SecretAccessKey: "KEYEXAMPLE",
+ },
+ },
+ done: make(chan struct{}),
+ }
+
+ c := NewCredentials(stub)
+
+ ctx, cancel1 := context.WithCancel(context.Background())
+ ctx1 := &ContextWaiter{Context: ctx, waiting: make(chan struct{}, 1)}
+ ctx2 := &ContextWaiter{Context: context.Background(), waiting: make(chan struct{}, 1)}
+
+ var err1, err2 error
+ var creds1, creds2 Value
+
+ done1 := make(chan struct{})
+ go func() {
+ creds1, err1 = c.GetWithContext(ctx1)
+ done1 <- struct{}{}
+ }()
+
+ done2 := make(chan struct{})
+ go func() {
+ creds2, err2 = c.GetWithContext(ctx2)
+ done2 <- struct{}{}
+ }()
+
+ <-ctx1.waiting
+ <-ctx2.waiting
+
+ cancel1()
+ <-done1
+
+ stub.done <- struct{}{}
+ <-done2
+
+ if err1 == nil {
+ t.Errorf("expect first to have error")
+ }
+ if creds1.HasKeys() {
+ t.Errorf("expect first not to have keys, %v", creds1)
+ }
+
+ if err2 != nil {
+ t.Errorf("expect second not to have error, %v", err2)
+ }
+ if !creds2.HasKeys() {
+ t.Errorf("Expect second to have keys")
+ }
+}
+
+type ContextWaiter struct {
+ context.Context
+ waiting chan struct{}
+}
+
+func (c *ContextWaiter) Done() <-chan struct{} {
+ go func() {
+ c.waiting <- struct{}{}
+ }()
+
+ return c.Context.Done()
+}

aws/signer/v4/request_context_go1.5.go

@@ -0,0 +1,13 @@
+// +build !go1.7
+
+package v4
+
+import (
+ "net/http"
+
+ "github.com/aws/aws-sdk-go/aws"
+)
+
+func requestContext(r *http.Request) aws.Context {
+ return aws.BackgroundContext()
+}

aws/signer/v4/request_context_go1.7.go

@@ -0,0 +1,13 @@
+// +build go1.7
+
+package v4
+
+import (
+ "net/http"
+
+ "github.com/aws/aws-sdk-go/aws"
+)
+
+func requestContext(r *http.Request) aws.Context {
+ return r.Context()
+}

aws/signer/v4/v4.go

@@ -340,7 +340,7 @@ func (v4 Signer) signWithBody(r *http.Request, body io.ReadSeeker, service, regi
  }
 
  var err error
- ctx.credValues, err = v4.Credentials.Get()
+ ctx.credValues, err = v4.Credentials.GetWithContext(requestContext(r))
  if err != nil {
  return http.Header{}, err
  }

internal/context/background_go1.5.go

@@ -0,0 +1,40 @@
+// +build !go1.7
+
+package context
+
+import "time"
+
+// An emptyCtx is a copy of the Go 1.7 context.emptyCtx type. This is copied to
+// provide a 1.6 and 1.5 safe version of context that is compatible with Go
+// 1.7's Context.
+//
+// An emptyCtx is never canceled, has no values, and has no deadline. It is not
+// struct{}, since vars of this type must have distinct addresses.
+type emptyCtx int
+
+func (*emptyCtx) Deadline() (deadline time.Time, ok bool) {
+ return
+}
+
+func (*emptyCtx) Done() <-chan struct{} {
+ return nil
+}
+
+func (*emptyCtx) Err() error {
+ return nil
+}
+
+func (*emptyCtx) Value(key interface{}) interface{} {
+ return nil
+}
+
+func (e *emptyCtx) String() string {
+ switch e {
+ case BackgroundCtx:
+ return "aws.BackgroundContext"
+ }
+ return "unknown empty Context"
+}
+
+// BackgroundCtx is the common base context.
+var BackgroundCtx = new(emptyCtx)