Use existing carrier; avoid using confusing pointers

DataDog · Oct 10, 2024 · f95440c · f95440c
1 parent f917961
commit f95440c
Show file tree

Hide file tree

Showing 6 changed files with 17 additions and 90 deletions.
diff --git a/contrib/aws/internal/eventbridge/eventbridge.go b/contrib/aws/internal/eventbridge/eventbridge.go
@@ -25,12 +25,6 @@ const (
 	maxSizeBytes    = 256 * 1024 // 256 KB
 )
 
-type messageCarrier map[string]string
-
-func (carrier messageCarrier) Set(key, val string) {
-	carrier[key] = val
-}
-
 func EnrichOperation(ctx context.Context, in middleware.InitializeInput, operation string) {
 	switch operation {
 	case "PutEvents":
@@ -61,7 +55,7 @@ func injectTraceContext(ctx context.Context, entryPtr *types.PutEventsRequestEnt
 		return
 	}
 
-	carrier := make(messageCarrier)
+	carrier := tracer.TextMapCarrier{}
 	err := tracer.Inject(span.Context(), carrier)
 	if err != nil {
 		log.Debug("Unable to inject trace context: %s", err)

diff --git a/contrib/aws/internal/eventbridge/eventbridge_test.go b/contrib/aws/internal/eventbridge/eventbridge_test.go
@@ -22,23 +22,6 @@ import (
 	"gopkg.in/DataDog/dd-trace-go.v1/ddtrace/tracer"
 )
 
-type testCarrier struct {
-	m map[string]string
-}
-
-func (c *testCarrier) Set(key, val string) {
-	c.m[key] = val
-}
-
-func (c *testCarrier) ForeachKey(handler func(key, val string) error) error {
-	for k, v := range c.m {
-		if err := handler(k, v); err != nil {
-			return err
-		}
-	}
-	return nil
-}
-
 func TestEnrichOperation(t *testing.T) {
 	mt := mocktracer.Start()
 	defer mt.Stop()
@@ -145,11 +128,10 @@ func TestInjectTraceContext(t *testing.T) {
 			assert.NoError(t, err)
 			assert.Greater(t, startTime, int64(0))
 
-			var carrier testCarrier
-			carrier.m = make(map[string]string)
+			carrier := tracer.TextMapCarrier{}
 			for k, v := range ddData {
 				if s, ok := v.(string); ok {
-					carrier.m[k] = s
+					carrier[k] = s
 				}
 			}
 

diff --git a/contrib/aws/internal/sns/sns.go b/contrib/aws/internal/sns/sns.go
@@ -21,12 +21,6 @@ const (
 	maxMessageAttributes = 10
 )
 
-type messageCarrier map[string]string
-
-func (carrier messageCarrier) Set(key, val string) {
-	carrier[key] = val
-}
-
 func EnrichOperation(ctx context.Context, in middleware.InitializeInput, operation string) {
 	switch operation {
 	case "Publish":
@@ -58,11 +52,10 @@ func handlePublishBatch(ctx context.Context, in middleware.InitializeInput) {
 	}
 
 	for i := range params.PublishBatchRequestEntries {
-		entryPtr := &params.PublishBatchRequestEntries[i]
-		if entryPtr.MessageAttributes == nil {
-			entryPtr.MessageAttributes = make(map[string]types.MessageAttributeValue)
+		if params.PublishBatchRequestEntries[i].MessageAttributes == nil {
+			params.PublishBatchRequestEntries[i].MessageAttributes = make(map[string]types.MessageAttributeValue)
 		}
-		injectTraceContext(ctx, entryPtr.MessageAttributes)
+		injectTraceContext(ctx, params.PublishBatchRequestEntries[i].MessageAttributes)
 	}
 }
 
@@ -81,7 +74,7 @@ func injectTraceContext(ctx context.Context, messageAttributes map[string]types.
 		return
 	}
 
-	carrier := make(messageCarrier)
+	carrier := tracer.TextMapCarrier{}
 	err := tracer.Inject(span.Context(), carrier)
 	if err != nil {
 		log.Debug("Unable to inject trace context: %s", err.Error())

diff --git a/contrib/aws/internal/sns/sns_test.go b/contrib/aws/internal/sns/sns_test.go
@@ -21,23 +21,6 @@ import (
 	"gopkg.in/DataDog/dd-trace-go.v1/ddtrace/tracer"
 )
 
-type testCarrier struct {
-	m map[string]string
-}
-
-func (c *testCarrier) Set(key, val string) {
-	c.m[key] = val
-}
-
-func (c *testCarrier) ForeachKey(handler func(key, val string) error) error {
-	for k, v := range c.m {
-		if err := handler(k, v); err != nil {
-			return err
-		}
-	}
-	return nil
-}
-
 func TestEnrichOperation(t *testing.T) {
 	tests := []struct {
 		name      string
@@ -179,12 +162,11 @@ func TestInjectTraceContext(t *testing.T) {
 				assert.NotNil(t, messageAttributes[datadogKey].StringValue)
 				assert.NotEmpty(t, *messageAttributes[datadogKey].StringValue)
 
-				var carrier testCarrier
-				carrier.m = make(map[string]string)
-				err := json.Unmarshal([]byte(*messageAttributes[datadogKey].StringValue), &carrier.m)
+				carrier := tracer.TextMapCarrier{}
+				err := json.Unmarshal([]byte(*messageAttributes[datadogKey].StringValue), &carrier)
 				assert.NoError(t, err)
 
-				extractedSpanContext, err := tracer.Extract(&carrier)
+				extractedSpanContext, err := tracer.Extract(carrier)
 				assert.NoError(t, err)
 				assert.Equal(t, span.Context().TraceID(), extractedSpanContext.TraceID())
 				assert.Equal(t, span.Context().SpanID(), extractedSpanContext.SpanID())

diff --git a/contrib/aws/internal/sqs/sqs.go b/contrib/aws/internal/sqs/sqs.go
@@ -21,12 +21,6 @@ const (
 	maxMessageAttributes = 10
 )
 
-type messageCarrier map[string]string
-
-func (carrier messageCarrier) Set(key, val string) {
-	carrier[key] = val
-}
-
 func EnrichOperation(ctx context.Context, in middleware.InitializeInput, operation string) {
 	switch operation {
 	case "SendMessage":
@@ -58,11 +52,10 @@ func handleSendMessageBatch(ctx context.Context, in middleware.InitializeInput)
 	}
 
 	for i := range params.Entries {
-		entryPtr := &params.Entries[i]
-		if entryPtr.MessageAttributes == nil {
-			entryPtr.MessageAttributes = make(map[string]types.MessageAttributeValue)
+		if params.Entries[i].MessageAttributes == nil {
+			params.Entries[i].MessageAttributes = make(map[string]types.MessageAttributeValue)
 		}
-		injectTraceContext(ctx, entryPtr.MessageAttributes)
+		injectTraceContext(ctx, params.Entries[i].MessageAttributes)
 	}
 }
 
@@ -81,7 +74,7 @@ func injectTraceContext(ctx context.Context, messageAttributes map[string]types.
 		return
 	}
 
-	carrier := make(messageCarrier)
+	carrier := tracer.TextMapCarrier{}
 	err := tracer.Inject(span.Context(), carrier)
 	if err != nil {
 		log.Debug("Unable to inject trace context: %s", err.Error())

diff --git a/contrib/aws/internal/sqs/sqs_test.go b/contrib/aws/internal/sqs/sqs_test.go
@@ -21,23 +21,6 @@ import (
 	"gopkg.in/DataDog/dd-trace-go.v1/ddtrace/tracer"
 )
 
-type testCarrier struct {
-	m map[string]string
-}
-
-func (c *testCarrier) Set(key, val string) {
-	c.m[key] = val
-}
-
-func (c *testCarrier) ForeachKey(handler func(key, val string) error) error {
-	for k, v := range c.m {
-		if err := handler(k, v); err != nil {
-			return err
-		}
-	}
-	return nil
-}
-
 func TestEnrichOperation(t *testing.T) {
 	tests := []struct {
 		name      string
@@ -183,11 +166,11 @@ func TestInjectTraceContext(t *testing.T) {
 				assert.NotNil(t, messageAttributes[datadogKey].StringValue)
 				assert.NotEmpty(t, *messageAttributes[datadogKey].StringValue)
 
-				var carrier testCarrier
-				err := json.Unmarshal([]byte(*messageAttributes[datadogKey].StringValue), &carrier.m)
+				carrier := tracer.TextMapCarrier{}
+				err := json.Unmarshal([]byte(*messageAttributes[datadogKey].StringValue), &carrier)
 				assert.NoError(t, err)
 
-				extractedSpanContext, err := tracer.Extract(&carrier)
+				extractedSpanContext, err := tracer.Extract(carrier)
 				assert.NoError(t, err)
 				assert.Equal(t, span.Context().TraceID(), extractedSpanContext.TraceID())
 				assert.Equal(t, span.Context().SpanID(), extractedSpanContext.SpanID())