Skip IMDS calls for X-Ray telemetry if in LocalMode. (#37)

aws · Jun 22, 2023 · c74a7dc · c74a7dc
1 parent e1bc65d
commit c74a7dc
Show file tree

Hide file tree

Showing 3 changed files with 45 additions and 25 deletions.
diff --git a/exporter/awsxrayexporter/awsxray.go b/exporter/awsxrayexporter/awsxray.go
@@ -59,7 +59,7 @@ func newTracesExporter(
 	xrayClient := awsxray.NewXRayClient(logger, awsConfig, set.BuildInfo, session)
 	sender := telemetry.NewNopSender()
 	if cfg.TelemetryConfig.Enabled {
-		opts := telemetry.ToOptions(cfg.TelemetryConfig, *awsSDKV2Config, &cfg.AWSSessionSettings, logger)
+		opts := telemetry.ToOptions(cfg.TelemetryConfig, awsSDKV2Config, &cfg.AWSSessionSettings, logger)
 		opts = append(opts, telemetry.WithLogger(set.Logger))
 		sender = registry.Register(set.ID, cfg.TelemetryConfig, xrayClient, opts...)
 	}

diff --git a/internal/aws/xray/telemetry/sender.go b/internal/aws/xray/telemetry/sender.go
@@ -183,28 +183,36 @@ func (p ec2MetadataProvider) get() string {
 }
 
 // ToOptions returns the metadata options if enabled by the config.
-func ToOptions(cfg Config, config aws.Config, settings *awsutil.AWSSessionSettings, logger *zap.Logger) []Option {
+func ToOptions(cfg Config, awsConfig *aws.Config, settings *awsutil.AWSSessionSettings, logger *zap.Logger) []Option {
 	if !cfg.IncludeMetadata {
 		return nil
 	}
-	clientIMDSV2Only, clientIMDSV1Fallback := awsutil.CreateIMDSV2AndFallbackClient(config)
+	hostnameProviders := []metadataProvider{
+		simpleMetadataProvider{metadata: cfg.Hostname},
+		envMetadataProvider{envKey: envAWSHostname},
+	}
+	instanceIDProviders := []metadataProvider{
+		simpleMetadataProvider{metadata: cfg.InstanceID},
+		envMetadataProvider{envKey: envAWSInstanceID},
+	}
+	if !settings.LocalMode {
+		clientIMDSV2Only, clientIMDSV1Fallback := awsutil.CreateIMDSV2AndFallbackClient(*awsConfig)
+		hostnameProviders = append(hostnameProviders, ec2MetadataProvider{
+			clientIMDSV2Only:     clientIMDSV2Only,
+			clientIMDSV1Fallback: clientIMDSV1Fallback,
+			metadataKey:          metadataHostname,
+			logger:               logger,
+		})
+		instanceIDProviders = append(instanceIDProviders, ec2MetadataProvider{
+			clientIMDSV2Only:     clientIMDSV2Only,
+			clientIMDSV1Fallback: clientIMDSV1Fallback,
+			metadataKey:          metadataInstanceID,
+			logger:               logger,
+		})
+	}
 	return []Option{
-		WithHostname(getMetadata(
-			simpleMetadataProvider{metadata: cfg.Hostname},
-			envMetadataProvider{envKey: envAWSHostname},
-			ec2MetadataProvider{clientIMDSV2Only: clientIMDSV2Only,
-				clientIMDSV1Fallback: clientIMDSV1Fallback,
-				metadataKey:          metadataHostname,
-				logger:               logger},
-		)),
-		WithInstanceID(getMetadata(
-			simpleMetadataProvider{metadata: cfg.InstanceID},
-			envMetadataProvider{envKey: envAWSInstanceID},
-			ec2MetadataProvider{clientIMDSV2Only: clientIMDSV2Only,
-				clientIMDSV1Fallback: clientIMDSV1Fallback,
-				metadataKey:          metadataInstanceID,
-				logger:               logger},
-		)),
+		WithHostname(getMetadata(hostnameProviders...)),
+		WithInstanceID(getMetadata(instanceIDProviders...)),
 		WithResourceARN(getMetadata(
 			simpleMetadataProvider{metadata: cfg.ResourceARN},
 			simpleMetadataProvider{metadata: settings.ResourceARN},

diff --git a/internal/aws/xray/telemetry/sender_test.go b/internal/aws/xray/telemetry/sender_test.go
@@ -21,7 +21,7 @@ import (
 	"testing"
 	"time"
 
-	awsConfig "github.com/aws/aws-sdk-go-v2/config"
+	awsconfig "github.com/aws/aws-sdk-go-v2/config"
 	"github.com/aws/aws-sdk-go/service/xray"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/mock"
@@ -79,33 +79,45 @@ func TestRotateRace(t *testing.T) {
 
 func TestIncludeMetadata(t *testing.T) {
 	cfg := Config{IncludeMetadata: false}
-	awsConfig, _ := awsConfig.LoadDefaultConfig(context.Background())
+	awsConfig, _ := awsconfig.LoadDefaultConfig(context.Background())
 	set := &awsutil.AWSSessionSettings{ResourceARN: "session_arn"}
-	opts := ToOptions(cfg, awsConfig, set, nil)
+	opts := ToOptions(cfg, &awsConfig, set, zap.NewNop())
 	assert.Empty(t, opts)
 	cfg.IncludeMetadata = true
-	opts = ToOptions(cfg, awsConfig, set, nil)
+	opts = ToOptions(cfg, &awsConfig, set, zap.NewNop())
 	sender := newSender(&mockClient{}, opts...)
 	assert.Equal(t, "", sender.hostname)
 	assert.Equal(t, "", sender.instanceID)
 	assert.Equal(t, "session_arn", sender.resourceARN)
 	t.Setenv(envAWSHostname, "env_hostname")
 	t.Setenv(envAWSInstanceID, "env_instance_id")
-	opts = ToOptions(cfg, awsConfig, &awsutil.AWSSessionSettings{}, nil)
+	opts = ToOptions(cfg, &awsConfig, &awsutil.AWSSessionSettings{}, zap.NewNop())
 	sender = newSender(&mockClient{}, opts...)
 	assert.Equal(t, "env_hostname", sender.hostname)
 	assert.Equal(t, "env_instance_id", sender.instanceID)
 	assert.Equal(t, "", sender.resourceARN)
 	cfg.Hostname = "cfg_hostname"
 	cfg.InstanceID = "cfg_instance_id"
 	cfg.ResourceARN = "cfg_arn"
-	opts = ToOptions(cfg, awsConfig, &awsutil.AWSSessionSettings{}, nil)
+	opts = ToOptions(cfg, &awsConfig, &awsutil.AWSSessionSettings{}, zap.NewNop())
 	sender = newSender(&mockClient{}, opts...)
 	assert.Equal(t, "cfg_hostname", sender.hostname)
 	assert.Equal(t, "cfg_instance_id", sender.instanceID)
 	assert.Equal(t, "cfg_arn", sender.resourceARN)
 }
 
+func TestIncludeMetadataLocalMode(t *testing.T) {
+	cfg := Config{IncludeMetadata: true}
+	set := &awsutil.AWSSessionSettings{ResourceARN: "session_arn", LocalMode: true}
+	assert.NotPanics(t, func() {
+		opts := ToOptions(cfg, nil, set, zap.NewNop())
+		sender := newSender(&mockClient{}, opts...)
+		assert.Equal(t, "", sender.hostname)
+		assert.Equal(t, "", sender.instanceID)
+		assert.Equal(t, "session_arn", sender.resourceARN)
+	})
+}
+
 func TestQueueOverflow(t *testing.T) {
 	obs, logs := observer.New(zap.DebugLevel)
 	client := &mockClient{count: &atomic.Int64{}}