aws/ec2metadata`: modifies IMDS client to use shorter request timeout

The PR modifies EC2Metadata client to use request context within its operations. Reduces the dialer timeout and response header timeout on the EC2Metadata client to help reduce latency for known issues with EC2Metadata client running inside a container.

Author: skotambkar

CHANGELOG_PENDING.md

@@ -21,6 +21,8 @@ SDK Enhancements
   * Adds an implementation of RingBuffer data structure which acts as a revolving buffer of a predefined length. The RingBuffer implements io.ReadWriter interface.
   * Adds unit tests to test the behavior of the ring buffer. 
 * `aws/ec2metadata`: Adds support for EC2Metadata client to use secure tokens provided by the IMDS ([#453](https://github.com/aws/aws-sdk-go-v2/pull/453)) 
+  * Modifies EC2Metadata client to use request context within its operations ([#462](https://github.com/aws/aws-sdk-go-v2/pull/462))
+  * Reduces the default dialer timeout and response header timeout to help reduce latency for known issues with EC2Metadata client running inside a container
   * Modifies and adds tests to verify the behavior of the EC2Metadata client.
 * `service/dynamodb/dynamodbattribute`: Adds clarifying docs on dynamodbattribute.UnixTime ([#464](https://github.com/aws/aws-sdk-go-v2/pull/464))
 * `example/service/sts/assumeRole`: added sts assume role example ([#224](https://github.com/aws/aws-sdk-go-v2/pull/224))

aws/ec2metadata/api_client.go

@@ -9,6 +9,7 @@ package ec2metadata
 
 import (
 	"bytes"
+	"context"
 	"errors"
 	"io"
 	"net"
@@ -33,11 +34,18 @@ const (
 	tokenHeader = "x-aws-ec2-metadata-token"
 
 	// Named Handler constants
+	contextWithTimeoutHandlerName  = "ContextWithTimeoutHandler"
+	cancelContextHandlerName       = "CancelContextHandler"
 	fetchTokenHandlerName          = "FetchTokenHandler"
 	unmarshalMetadataHandlerName   = "unmarshalMetadataHandler"
 	unmarshalTokenHandlerName      = "unmarshalTokenHandler"
 	enableTokenProviderHandlerName = "enableTokenProviderHandler"
 
+	// client constants
+	defaultClientContextTimeout  = 5 * time.Second
+	defaultDialerTimeout         = 250 * time.Millisecond
+	defaultResponseHeaderTimeout = 500 * time.Millisecond
+
 	// TTL constants
 	defaultTTL          = 21600 * time.Second
 	ttlExpirationWindow = 30 * time.Second
@@ -64,7 +72,15 @@ func New(config aws.Config) *Client {
 		// environment with the service present. The client should fail fast in
 		// this case.
 		config.HTTPClient = c.WithDialerOptions(func(d *net.Dialer) {
-			d.Timeout = 5 * time.Second
+			d.Timeout = defaultDialerTimeout
+		})
+
+		// Use a custom Transport timeout for the EC2 Metadata service to account
+		// for the possibility that the application might be running in a container,
+		// and EC2Metadata service drops the connection after a single IP Hop. The client
+		// should fail fast in this case.
+		config.HTTPClient = c.WithTransportOptions(func(tr *http.Transport) {
+			tr.ResponseHeaderTimeout = defaultResponseHeaderTimeout
 		})
 	}
 
@@ -87,6 +103,23 @@ func New(config aws.Config) *Client {
 		Name: fetchTokenHandlerName,
 		Fn:   tp.fetchTokenHandler,
 	})
+
+	// The context With timeout handler function wraps a context with timeout and sets it on a request.
+	// It also sets a handler on complete handler stack that cancels the context
+	svc.Handlers.Send.PushFrontNamed(aws.NamedHandler{
+		Name: contextWithTimeoutHandlerName,
+		Fn: func(r *aws.Request) {
+			ctx, cancelFn := context.WithTimeout(r.Context(), defaultClientContextTimeout)
+			r.SetContext(ctx)
+			r.Handlers.Complete.PushBackNamed(aws.NamedHandler{
+				Name: cancelContextHandlerName,
+				Fn: func(r *aws.Request) {
+					cancelFn()
+				},
+			})
+		},
+	})
+
 	// NamedHandler for enabling token provider
 	svc.Handlers.Complete.PushBackNamed(aws.NamedHandler{
 		Name: enableTokenProviderHandlerName,

aws/ec2metadata/api_ops_test.go

@@ -761,6 +761,10 @@ func TestEC2MetadataRetryFailure(t *testing.T) {
 	cfg := unit.Config()
 	cfg.EndpointResolver = aws.EndpointResolverFunc(myCustomResolver)
 	c := ec2metadata.New(cfg)
+	// mock retryer with minimum throttle delay set to 1 ms
+	c.Retryer = aws.NewDefaultRetryer(func(d *aws.DefaultRetryer) {
+		d.MinThrottleDelay = 1 * time.Millisecond
+	})
 	// Handler on client that logs if retried
 	c.Handlers.AfterRetry.PushBack(func(i *aws.Request) {
 		t.Logf("%v received, retrying operation %v", i.HTTPResponse.StatusCode, i.Operation.Name)
@@ -794,7 +798,7 @@ func TestEC2MetadataRetryOnce(t *testing.T) {
 	mux.HandleFunc("/latest/api/token", func(w http.ResponseWriter, r *http.Request) {
 		if r.Method == "PUT" && r.Header.Get(ttlHeader) != "" {
 			w.Header().Set(ttlHeader, "200")
-			for retry {
+			if retry {
 				retry = false
 				http.Error(w, "service unavailable", http.StatusServiceUnavailable)
 				return
@@ -824,6 +828,10 @@ func TestEC2MetadataRetryOnce(t *testing.T) {
 	cfg := unit.Config()
 	cfg.EndpointResolver = aws.EndpointResolverFunc(myCustomResolver)
 	c := ec2metadata.New(cfg)
+	// mock retryer with minimum throttle delay set to 1 ms
+	c.Retryer = aws.NewDefaultRetryer(func(d *aws.DefaultRetryer) {
+		d.MinThrottleDelay = 1 * time.Millisecond
+	})
 	// Handler on client that logs if retried
 	c.Handlers.AfterRetry.PushBack(func(i *aws.Request) {
 		t.Logf("%v received, retrying operation %v", i.HTTPResponse.StatusCode, i.Operation.Name)
@@ -1088,9 +1096,13 @@ func TestRequestTimeOut(t *testing.T) {
 	}
 
 	c.Handlers.Complete.PushBack(op.addToOperationPerformedList)
-
+	start := time.Now()
 	resp, err := c.GetMetadata(context.Background(), "/some/path")
 
+	if e, a := 1*time.Second, time.Since(start); e < a {
+		t.Fatalf("expected duration of test to be less than %v, got %v", e, a)
+	}
+
 	expectedOperationsPerformed := []string{"GetToken", "GetMetadata"}
 
 	if e, a := "IMDSProfileForSDKGo", resp; e != a {
@@ -1105,7 +1117,11 @@ func TestRequestTimeOut(t *testing.T) {
 		t.Fatalf("Found diff in operations performed: \n %v \n", diff)
 	}
 
+	start = time.Now()
 	resp, err = c.GetMetadata(context.Background(), "/some/path")
+	if e, a := 1*time.Second, time.Since(start); e < a {
+		t.Fatalf("expected duration of test to be less than %v, got %v", e, a)
+	}
 
 	expectedOperationsPerformed = []string{"GetToken", "GetMetadata", "GetMetadata"}
 

aws/signer/v2/v2_test.go

@@ -63,7 +63,6 @@ func TestSignRequestWithAndWithoutSession(t *testing.T) {
 	}
 
 	// create request without a SecurityToken (session) in the credentials
-
 	query := newQuery()
 	timestamp := time.Date(2015, 7, 16, 7, 56, 16, 0, time.UTC)
 	builder := signerBuilder{

aws/signer/v4/v4.go

@@ -563,6 +563,7 @@ func buildQuery(r rule, header http.Header) (url.Values, http.Header) {
 
 	return query, unsignedHeaders
 }
+
 func (ctx *signingCtx) buildCanonicalHeaders(r rule, header http.Header) {
 	var headers []string
 	headers = append(headers, "host")