support AWS_METADATA_SERVICE envs on EC2

go.mod

@@ -5,4 +5,5 @@ go 1.18
 require (
 	github.com/google/go-cmp v0.6.0
 	github.com/shogo82148/forwarded-header v0.1.0
+	github.com/shogo82148/go-retry v1.2.0
 )

go.sum

@@ -2,3 +2,5 @@ github.com/google/go-cmp v0.6.0 h1:ofyhxvXcZhMsU5ulbFiLKl/XBFqE1GSq7atu8tAmTRI=
 github.com/google/go-cmp v0.6.0/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY=
 github.com/shogo82148/forwarded-header v0.1.0 h1:OZX131i0B8GGZR+s5QEZdSEYDTUPkUX8ISIIQE+swKM=
 github.com/shogo82148/forwarded-header v0.1.0/go.mod h1:Ge73zEgn9jrdyf65vWTgDmRZQFQPHoFU0TaILJcvs54=
+github.com/shogo82148/go-retry v1.2.0 h1:A/LFdbZKJ+tsT1gF4OrzM4P10FGK7VUExpb07/U03aE=
+github.com/shogo82148/go-retry v1.2.0/go.mod h1:wttfgfwCMQvNqv4kOpqIvDDJeSmwU+AEIpUyG+5Ca6M=

internal/envconfig/metadata.go

@@ -0,0 +1,44 @@
+package envconfig
+
+import (
+	"os"
+	"strconv"
+	"time"
+)
+
+// MetadataServiceTimeout returns the timeout for the metadata service.
+func MetadataServiceTimeout() time.Duration {
+	const defaultTimeout = time.Second
+
+	timeout := os.Getenv("AWS_METADATA_SERVICE_TIMEOUT")
+	if timeout == "" {
+		return defaultTimeout
+	}
+	d, err := strconv.ParseInt(timeout, 10, 64)
+	if err != nil {
+		return defaultTimeout
+	}
+	if d <= 0 {
+		return defaultTimeout
+	}
+	return time.Duration(d) * time.Second
+}
+
+// MetadataServiceNumAttempts returns the number of attempts to make when
+// querying the metadata service.
+func MetadataServiceNumAttempts() int {
+	const defaultNumAttempts = 1
+
+	attempts := os.Getenv("AWS_METADATA_SERVICE_NUM_ATTEMPTS")
+	if attempts == "" {
+		return defaultNumAttempts
+	}
+	n, err := strconv.ParseInt(attempts, 10, 0)
+	if err != nil {
+		return defaultNumAttempts
+	}
+	if n < 1 {
+		return defaultNumAttempts
+	}
+	return int(n)
+}

internal/envconfig/metadata_test.go

@@ -0,0 +1,79 @@
+package envconfig
+
+import (
+	"fmt"
+	"testing"
+	"time"
+)
+
+func TestMetadataServiceTimeout(t *testing.T) {
+	tests := []struct {
+		name string
+		want time.Duration
+	}{
+		{
+			name: "",
+			want: time.Second,
+		},
+		{
+			name: "1",
+			want: time.Second,
+		},
+		{
+			name: "invalid",
+			want: time.Second,
+		},
+		{
+			name: "10",
+			want: 10 * time.Second,
+		},
+		{
+			name: "0",
+			want: time.Second,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(fmt.Sprintf("%q", tt.name), func(t *testing.T) {
+			t.Setenv("AWS_METADATA_SERVICE_TIMEOUT", tt.name)
+			if got := MetadataServiceTimeout(); got != tt.want {
+				t.Errorf("MetadataServiceTimeout() = %v, want %v", got, tt.want)
+			}
+		})
+	}
+}
+
+func TestMetadataServiceNumAttempts(t *testing.T) {
+	tests := []struct {
+		name string
+		want int
+	}{
+		{
+			name: "",
+			want: 1,
+		},
+		{
+			name: "1",
+			want: 1,
+		},
+		{
+			name: "invalid",
+			want: 1,
+		},
+		{
+			name: "10",
+			want: 10,
+		},
+		{
+			name: "0",
+			want: 1,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(fmt.Sprintf("%q", tt.name), func(t *testing.T) {
+			t.Setenv("AWS_METADATA_SERVICE_NUM_ATTEMPTS", tt.name)
+			if got := MetadataServiceNumAttempts(); got != tt.want {
+				t.Errorf("MetadataServiceNumAttempts() = %v, want %v", got, tt.want)
+			}
+		})
+	}
+}

plugins/ec2/ec2.go

@@ -29,9 +29,11 @@ import (
 	"sync"
 	"time"
 
+	"github.com/shogo82148/aws-xray-yasdk-go/internal/envconfig"
 	"github.com/shogo82148/aws-xray-yasdk-go/xray"
 	"github.com/shogo82148/aws-xray-yasdk-go/xray/schema"
 	"github.com/shogo82148/aws-xray-yasdk-go/xray/xraylog"
+	"github.com/shogo82148/go-retry"
 )
 
 type ec2InstanceIdentityDocument struct {
@@ -68,6 +70,9 @@ type client struct {
 	// TTL for token
 	ttl time.Time
 
+	timeout time.Duration
+	policy  *retry.Policy
+
 	httpClient *http.Client
 }
 
@@ -92,8 +97,14 @@ func newClient(ctx context.Context) *client {
 		}
 	}
 	base = strings.TrimSuffix(base, "/")
+	timeout := envconfig.MetadataServiceTimeout()
+
 	c := &client{
-		base: base,
+		base:    base,
+		timeout: timeout,
+		policy: &retry.Policy{
+			MaxCount: envconfig.MetadataServiceNumAttempts(),
+		},
 		httpClient: &http.Client{
 			Transport: &http.Transport{
 				// ignore proxy configure from the environment values
@@ -102,16 +113,15 @@ func newClient(ctx context.Context) *client {
 				// metadata endpoint is in same network,
 				// so timeout can be shorter.
 				DialContext: (&net.Dialer{
-					Timeout:   time.Second,
-					KeepAlive: time.Second,
+					Timeout:   timeout,
+					KeepAlive: timeout,
 					DualStack: true,
 				}).DialContext,
 				MaxIdleConns:          5,
-				IdleConnTimeout:       time.Second,
-				TLSHandshakeTimeout:   time.Second,
-				ExpectContinueTimeout: time.Second,
+				IdleConnTimeout:       timeout,
+				TLSHandshakeTimeout:   timeout,
+				ExpectContinueTimeout: timeout,
 			},
-			Timeout: time.Second,
 		},
 	}
 	return c
@@ -122,66 +132,75 @@ func (c *client) Close() {
 }
 
 func (c *client) refreshToken(ctx context.Context) error {
-	now := time.Now()
-	if c.token != "" && c.ttl.After(now) {
-		// no need to refresh
-		return nil
-	}
+	return c.policy.Do(ctx, func() error {
+		now := time.Now()
+		if c.token != "" && c.ttl.After(now) {
+			// no need to refresh
+			return nil
+		}
 
-	req, err := http.NewRequestWithContext(ctx, http.MethodPut, c.base+"/latest/api/token", nil)
-	if err != nil {
-		return err
-	}
-	req.Header.Set("x-aws-ec2-metadata-token-ttl-seconds", "10")
-	resp, err := c.httpClient.Do(req)
-	if err != nil {
-		return err
-	}
-	defer resp.Body.Close()
+		ctx, cancel := context.WithTimeout(ctx, c.timeout)
+		defer cancel()
 
-	if resp.StatusCode != http.StatusOK {
-		// IMDSv2 may be disabled; fallback to IMDSv1.
-		return nil
-	}
+		req, err := http.NewRequestWithContext(ctx, http.MethodPut, c.base+"/latest/api/token", nil)
+		if err != nil {
+			return err
+		}
+		req.Header.Set("x-aws-ec2-metadata-token-ttl-seconds", "10")
+		resp, err := c.httpClient.Do(req)
+		if err != nil {
+			return err
+		}
+		defer resp.Body.Close()
 
-	token, err := io.ReadAll(resp.Body)
-	if err != nil {
-		return err
-	}
-	c.token = string(token)
-	c.ttl = now.Add(5 * time.Second)
+		if resp.StatusCode != http.StatusOK {
+			// IMDSv2 may be disabled; fallback to IMDSv1.
+			return nil
+		}
 
-	return nil
+		token, err := io.ReadAll(resp.Body)
+		if err != nil {
+			return err
+		}
+		c.token = string(token)
+		c.ttl = now.Add(5 * time.Second)
+		return nil
+	})
 }
 
 func (c *client) getInstanceIdentityDocument(ctx context.Context) (*ec2InstanceIdentityDocument, error) {
 	if err := c.refreshToken(ctx); err != nil {
 		return nil, err
 	}
 
-	req, err := http.NewRequestWithContext(ctx, http.MethodGet, c.base+"/latest/dynamic/instance-identity/document", nil)
-	if err != nil {
-		return nil, err
-	}
-	if c.token != "" {
-		req.Header.Set("x-aws-ec2-metadata-token", c.token)
-	}
-	resp, err := c.httpClient.Do(req)
-	if err != nil {
-		return nil, err
-	}
-	defer resp.Body.Close()
+	return retry.DoValue(ctx, c.policy, func() (*ec2InstanceIdentityDocument, error) {
+		ctx, cancel := context.WithTimeout(ctx, c.timeout)
+		defer cancel()
 
-	if resp.StatusCode != http.StatusOK {
-		return nil, fmt.Errorf("unexpected status code: %d", resp.StatusCode)
-	}
+		req, err := http.NewRequestWithContext(ctx, http.MethodGet, c.base+"/latest/dynamic/instance-identity/document", nil)
+		if err != nil {
+			return nil, err
+		}
+		if c.token != "" {
+			req.Header.Set("x-aws-ec2-metadata-token", c.token)
+		}
+		resp, err := c.httpClient.Do(req)
+		if err != nil {
+			return nil, err
+		}
+		defer resp.Body.Close()
 
-	var doc ec2InstanceIdentityDocument
-	dec := json.NewDecoder(resp.Body)
-	if err := dec.Decode(&doc); err != nil {
-		return nil, err
-	}
-	return &doc, nil
+		if resp.StatusCode != http.StatusOK {
+			return nil, fmt.Errorf("unexpected status code: %d", resp.StatusCode)
+		}
+
+		var doc ec2InstanceIdentityDocument
+		dec := json.NewDecoder(resp.Body)
+		if err := dec.Decode(&doc); err != nil {
+			return nil, err
+		}
+		return &doc, nil
+	})
 }
 
 // Find logging configure of CloudWatch Agent.

plugins/ec2/ec2_test.go

@@ -13,6 +13,7 @@ import (
 
 	"github.com/google/go-cmp/cmp"
 	"github.com/shogo82148/aws-xray-yasdk-go/xray/schema"
+	"github.com/shogo82148/go-retry"
 )
 
 func TestGetInstanceIdentityDocument_IMDSv1(t *testing.T) {
@@ -53,6 +54,10 @@ func TestGetInstanceIdentityDocument_IMDSv1(t *testing.T) {
 	c := &client{
 		base:       ts.URL,
 		httpClient: &http.Client{},
+		timeout:    time.Second,
+		policy: &retry.Policy{
+			MaxCount: 1,
+		},
 	}
 	got, err := c.getInstanceIdentityDocument(context.Background())
 	if err != nil {
@@ -126,6 +131,10 @@ func TestGetInstanceIdentityDocument_IMDSv2(t *testing.T) {
 	c := &client{
 		base:       ts.URL,
 		httpClient: &http.Client{},
+		timeout:    time.Second,
+		policy: &retry.Policy{
+			MaxCount: 1,
+		},
 	}
 	got, err := c.getInstanceIdentityDocument(context.Background())
 	if err != nil {
@@ -221,8 +230,9 @@ func TestGetInstanceIdentityDocument_IMDSEndpointFromEnv(t *testing.T) {
 	ts := httptest.NewServer(mux)
 	defer ts.Close()
 
-	os.Setenv("AWS_EC2_METADATA_SERVICE_ENDPOINT", ts.URL)
-	defer os.Unsetenv("AWS_EC2_METADATA_SERVICE_ENDPOINT")
+	t.Setenv("AWS_EC2_METADATA_SERVICE_ENDPOINT", ts.URL)
+	t.Setenv("AWS_METADATA_SERVICE_TIMEOUT", "1")
+	t.Setenv("AWS_METADATA_SERVICE_NUM_ATTEMPTS", "1")
 
 	Init()