Skip to content

Commit

Permalink
Use existing carrier; avoid using confusing pointers
Browse files Browse the repository at this point in the history
  • Loading branch information
nhulston committed Oct 10, 2024
1 parent f917961 commit f95440c
Show file tree
Hide file tree
Showing 6 changed files with 17 additions and 90 deletions.
8 changes: 1 addition & 7 deletions contrib/aws/internal/eventbridge/eventbridge.go
Original file line number Diff line number Diff line change
Expand Up @@ -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":
Expand Down Expand Up @@ -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)
Expand Down
22 changes: 2 additions & 20 deletions contrib/aws/internal/eventbridge/eventbridge_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down Expand Up @@ -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
}
}

Expand Down
15 changes: 4 additions & 11 deletions contrib/aws/internal/sns/sns.go
Original file line number Diff line number Diff line change
Expand Up @@ -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":
Expand Down Expand Up @@ -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)
}
}

Expand All @@ -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())
Expand Down
24 changes: 3 additions & 21 deletions contrib/aws/internal/sns/sns_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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())
Expand Down
15 changes: 4 additions & 11 deletions contrib/aws/internal/sqs/sqs.go
Original file line number Diff line number Diff line change
Expand Up @@ -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":
Expand Down Expand Up @@ -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)
}
}

Expand All @@ -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())
Expand Down
23 changes: 3 additions & 20 deletions contrib/aws/internal/sqs/sqs_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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())
Expand Down

0 comments on commit f95440c

Please sign in to comment.