Skip to content

Commit 2645dc6

Browse files
PerfectSlayerPerfectSlayer
committed
refactor(aws): Move span pointers handling into instrumentations
1 parent 517a89d commit 2645dc6

File tree

12 files changed

+404
-377
lines changed

12 files changed

+404
-377
lines changed

‎dd-java-agent/instrumentation/aws-java/aws-java-dynamodb-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/dynamodb/DynamoDbInterceptor.java‎

Lines changed: 9 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -41,12 +41,17 @@ public void afterExecution(AfterExecution context, ExecutionAttributes execution
4141
}
4242

4343
SdkRequest request = context.request();
44+
String tableName = request.getValueForField("TableName", String.class).orElse(null);
45+
Map<String, AttributeValue> keys = null;
46+
4447
if (request instanceof UpdateItemRequest) {
45-
Map<String, AttributeValue> keys = ((UpdateItemRequest) request).key();
46-
DynamoDbUtil.exportTagsWithKnownKeys(span, keys);
48+
keys = ((UpdateItemRequest) request).key();
4749
} else if (request instanceof DeleteItemRequest) {
48-
Map<String, AttributeValue> keys = ((DeleteItemRequest) request).key();
49-
DynamoDbUtil.exportTagsWithKnownKeys(span, keys);
50+
keys = ((DeleteItemRequest) request).key();
51+
}
52+
53+
if (keys != null) {
54+
DynamoDbUtil.addSpanPointer(span, tableName, keys);
5055
}
5156
}
5257
}

‎dd-java-agent/instrumentation/aws-java/aws-java-dynamodb-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/dynamodb/DynamoDbUtil.java‎

Lines changed: 20 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,7 @@
11
package datadog.trace.instrumentation.aws.v2.dynamodb;
22

3-
import static datadog.trace.bootstrap.instrumentation.api.InstrumentationTags.DYNAMO_PRIMARY_KEY_1;
4-
import static datadog.trace.bootstrap.instrumentation.api.InstrumentationTags.DYNAMO_PRIMARY_KEY_1_VALUE;
5-
import static datadog.trace.bootstrap.instrumentation.api.InstrumentationTags.DYNAMO_PRIMARY_KEY_2;
6-
import static datadog.trace.bootstrap.instrumentation.api.InstrumentationTags.DYNAMO_PRIMARY_KEY_2_VALUE;
7-
83
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
4+
import datadog.trace.bootstrap.instrumentation.api.SpanPointerUtils;
95
import java.util.ArrayList;
106
import java.util.Collections;
117
import java.util.List;
@@ -37,36 +33,42 @@ private static String extractValueAsString(AttributeValue value) {
3733
}
3834

3935
/**
40-
* Gets primary key/values and exports them as temporary tags on the span so that
41-
* SpanPointersProcessor.java can complete the span pointer creation.
36+
* Creates a DynamoDB span pointer link from the given primary keys and adds it to the span.
4237
*
43-
* @param span The span to set the temporary tags on
38+
* @param span The span to add the pointer link to
39+
* @param tableName The DynamoDB table name
4440
* @param keys The primary key/values to extract from
4541
*/
46-
static void exportTagsWithKnownKeys(AgentSpan span, Map<String, AttributeValue> keys) {
47-
if (keys == null || keys.isEmpty()) {
42+
static void addSpanPointer(AgentSpan span, String tableName, Map<String, AttributeValue> keys) {
43+
if (keys == null || keys.isEmpty() || tableName == null) {
4844
return;
4945
}
5046

47+
String primaryKey1Name;
48+
String primaryKey1Value;
49+
String primaryKey2Name = null;
50+
String primaryKey2Value = null;
51+
5152
if (keys.size() == 1) {
5253
// Single primary key case
5354
Map.Entry<String, AttributeValue> entry = keys.entrySet().iterator().next();
54-
span.setTag(DYNAMO_PRIMARY_KEY_1, entry.getKey());
55-
span.setTag(DYNAMO_PRIMARY_KEY_1_VALUE, extractValueAsString(entry.getValue()));
55+
primaryKey1Name = entry.getKey();
56+
primaryKey1Value = extractValueAsString(entry.getValue());
5657
} else {
5758
// Sort keys alphabetically
5859
List<String> keyNames = new ArrayList<>(keys.keySet());
5960
Collections.sort(keyNames);
6061

6162
// First key (alphabetically)
62-
String primaryKey1Name = keyNames.get(0);
63-
span.setTag(DYNAMO_PRIMARY_KEY_1, primaryKey1Name);
64-
span.setTag(DYNAMO_PRIMARY_KEY_1_VALUE, extractValueAsString(keys.get(primaryKey1Name)));
63+
primaryKey1Name = keyNames.get(0);
64+
primaryKey1Value = extractValueAsString(keys.get(primaryKey1Name));
6565

6666
// Second key
67-
String primaryKey2Name = keyNames.get(1);
68-
span.setTag(DYNAMO_PRIMARY_KEY_2, primaryKey2Name);
69-
span.setTag(DYNAMO_PRIMARY_KEY_2_VALUE, extractValueAsString(keys.get(primaryKey2Name)));
67+
primaryKey2Name = keyNames.get(1);
68+
primaryKey2Value = extractValueAsString(keys.get(primaryKey2Name));
7069
}
70+
71+
SpanPointerUtils.addDynamoDbSpanPointer(
72+
span, tableName, primaryKey1Name, primaryKey1Value, primaryKey2Name, primaryKey2Value);
7173
}
7274
}

‎dd-java-agent/instrumentation/aws-java/aws-java-dynamodb-2.0/src/test/groovy/DynamoDbClientTest.groovy‎

Lines changed: 13 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@ import datadog.trace.api.DDSpanTypes
33
import datadog.trace.api.DDTraceId
44
import datadog.trace.bootstrap.instrumentation.api.SpanAttributes
55
import datadog.trace.bootstrap.instrumentation.api.SpanLink
6-
import datadog.trace.core.tagprocessor.SpanPointersProcessor
6+
import datadog.trace.bootstrap.instrumentation.api.SpanPointerUtils
77
import org.testcontainers.containers.GenericContainer
88
import org.testcontainers.utility.DockerImageName
99
import software.amazon.awssdk.auth.credentials.AwsBasicCredentials
@@ -193,11 +193,11 @@ class DynamoDbClientTest extends InstrumentationSpecification {
193193
spanType DDSpanTypes.HTTP_CLIENT
194194
links {
195195
link(DDTraceId.ZERO, 0, SpanLink.DEFAULT_FLAGS, SpanAttributes.builder()
196-
.put("ptr.kind", SpanPointersProcessor.DYNAMODB_PTR_KIND)
197-
.put("ptr.dir", SpanPointersProcessor.DOWN_DIRECTION)
196+
.put("ptr.kind", SpanPointerUtils.DYNAMODB_PTR_KIND)
197+
.put("ptr.dir", SpanPointerUtils.DOWN_DIRECTION)
198198
// First 32 chars of SHA256("dynamodb-one-key-table|id|test-id-1||")
199199
.put("ptr.hash", "ca8daaa857b00545ed5186a915cf1ab5")
200-
.put("link.kind", SpanPointersProcessor.LINK_KIND)
200+
.put("link.kind", SpanPointerUtils.LINK_KIND)
201201
.build())
202202
}
203203
tags {
@@ -295,11 +295,11 @@ class DynamoDbClientTest extends InstrumentationSpecification {
295295
spanType DDSpanTypes.HTTP_CLIENT
296296
links {
297297
link(DDTraceId.ZERO, 0, SpanLink.DEFAULT_FLAGS, SpanAttributes.builder()
298-
.put("ptr.kind", SpanPointersProcessor.DYNAMODB_PTR_KIND)
299-
.put("ptr.dir", SpanPointersProcessor.DOWN_DIRECTION)
298+
.put("ptr.kind", SpanPointerUtils.DYNAMODB_PTR_KIND)
299+
.put("ptr.dir", SpanPointerUtils.DOWN_DIRECTION)
300300
// First 32 chars of SHA256("dynamodb-two-key-table|primaryKey|customer-123|sortKey|order-456")
301301
.put("ptr.hash", "90922c7899a82ea34406fdcdfb95161e")
302-
.put("link.kind", SpanPointersProcessor.LINK_KIND)
302+
.put("link.kind", SpanPointerUtils.LINK_KIND)
303303
.build())
304304
}
305305
tags {
@@ -383,11 +383,11 @@ class DynamoDbClientTest extends InstrumentationSpecification {
383383
spanType DDSpanTypes.HTTP_CLIENT
384384
links {
385385
link(DDTraceId.ZERO, 0, SpanLink.DEFAULT_FLAGS, SpanAttributes.builder()
386-
.put("ptr.kind", SpanPointersProcessor.DYNAMODB_PTR_KIND)
387-
.put("ptr.dir", SpanPointersProcessor.DOWN_DIRECTION)
386+
.put("ptr.kind", SpanPointerUtils.DYNAMODB_PTR_KIND)
387+
.put("ptr.dir", SpanPointerUtils.DOWN_DIRECTION)
388388
// First 32 chars of SHA256("dynamodb-one-key-table|id|delete-test-id||")
389389
.put("ptr.hash", "65031164be5e929fddd274a02cba3f9f")
390-
.put("link.kind", SpanPointersProcessor.LINK_KIND)
390+
.put("link.kind", SpanPointerUtils.LINK_KIND)
391391
.build())
392392
}
393393
tags {
@@ -482,11 +482,11 @@ class DynamoDbClientTest extends InstrumentationSpecification {
482482
spanType DDSpanTypes.HTTP_CLIENT
483483
links {
484484
link(DDTraceId.ZERO, 0, SpanLink.DEFAULT_FLAGS, SpanAttributes.builder()
485-
.put("ptr.kind", SpanPointersProcessor.DYNAMODB_PTR_KIND)
486-
.put("ptr.dir", SpanPointersProcessor.DOWN_DIRECTION)
485+
.put("ptr.kind", SpanPointerUtils.DYNAMODB_PTR_KIND)
486+
.put("ptr.dir", SpanPointerUtils.DOWN_DIRECTION)
487487
// First 32 chars of SHA256("dynamodb-two-key-table|primaryKey|user-789|sortKey|profile")
488488
.put("ptr.hash", "e5ce1148208c6f88041c73ceb9bbbf3a")
489-
.put("link.kind", SpanPointersProcessor.LINK_KIND)
489+
.put("link.kind", SpanPointerUtils.LINK_KIND)
490490
.build())
491491
}
492492
tags {

‎dd-java-agent/instrumentation/aws-java/aws-java-dynamodb-2.0/src/test/groovy/DynamoDbUtilTest.groovy‎

Lines changed: 79 additions & 45 deletions
Original file line numberDiff line numberDiff line change
@@ -1,130 +1,164 @@
11
import datadog.trace.bootstrap.instrumentation.api.AgentSpan
2-
import datadog.trace.bootstrap.instrumentation.api.InstrumentationTags
2+
import datadog.trace.bootstrap.instrumentation.api.AgentSpanLink
3+
import datadog.trace.bootstrap.instrumentation.api.SpanPointerUtils
34
import datadog.trace.instrumentation.aws.v2.dynamodb.DynamoDbUtil
45
import org.junit.jupiter.api.Test
56
import software.amazon.awssdk.core.SdkBytes
67
import software.amazon.awssdk.services.dynamodb.model.AttributeValue
78

89
class DynamoDbUtilTest {
910
static createMockSpan() {
10-
def tags = [:]
11+
def links = []
1112

1213
def mockSpan = [
13-
setTag: { String key, String value ->
14-
tags[key] = value
15-
return null
14+
addLink: { AgentSpanLink link ->
15+
links.add(link)
1616
}
1717
] as AgentSpan
1818

19-
return [span: mockSpan, tags: tags]
19+
return [span: mockSpan, links: links]
2020
}
2121

2222
@Test
23-
void testExportTagsWithNullKeys() {
23+
void testAddSpanPointerWithNullKeys() {
2424
def mockData = createMockSpan()
2525
def mockSpan = mockData.span
26-
def tags = mockData.tags
26+
def links = mockData.links
2727

28-
DynamoDbUtil.exportTagsWithKnownKeys(mockSpan, null)
28+
DynamoDbUtil.addSpanPointer(mockSpan, "table", null)
2929

30-
assert tags.isEmpty()
30+
assert links.isEmpty()
3131
}
3232

3333
@Test
34-
void testExportTagsWithEmptyKeys() {
34+
void testAddSpanPointerWithEmptyKeys() {
3535
def mockData = createMockSpan()
3636
def mockSpan = mockData.span
37-
def tags = mockData.tags
37+
def links = mockData.links
3838

39-
DynamoDbUtil.exportTagsWithKnownKeys(mockSpan, [:])
39+
DynamoDbUtil.addSpanPointer(mockSpan, "table", [:])
4040

41-
assert tags.isEmpty()
41+
assert links.isEmpty()
4242
}
4343

4444
@Test
45-
void testExportTagsWithSingleStringKey() {
45+
void testAddSpanPointerWithNullTableName() {
4646
def mockData = createMockSpan()
4747
def mockSpan = mockData.span
48-
def tags = mockData.tags
48+
def links = mockData.links
4949

5050
def keys = [
5151
"id": AttributeValue.builder().s("12345").build()
5252
]
5353

54-
DynamoDbUtil.exportTagsWithKnownKeys(mockSpan, keys)
54+
DynamoDbUtil.addSpanPointer(mockSpan, null, keys)
5555

56-
assert tags[InstrumentationTags.DYNAMO_PRIMARY_KEY_1] == "id"
57-
assert tags[InstrumentationTags.DYNAMO_PRIMARY_KEY_1_VALUE] == "12345"
56+
assert links.isEmpty()
5857
}
5958

6059
@Test
61-
void testExportTagsWithSingleNumberKey() {
60+
void testAddSpanPointerWithSingleStringKey() {
6261
def mockData = createMockSpan()
6362
def mockSpan = mockData.span
64-
def tags = mockData.tags
63+
def links = mockData.links
64+
65+
def keys = [
66+
"id": AttributeValue.builder().s("12345").build()
67+
]
68+
69+
DynamoDbUtil.addSpanPointer(mockSpan, "my-table", keys)
70+
71+
assert links.size() == 1
72+
def link = links[0]
73+
assert link.attributes().asMap().get("ptr.kind") == SpanPointerUtils.DYNAMODB_PTR_KIND
74+
assert link.attributes().asMap().get("ptr.dir") == SpanPointerUtils.DOWN_DIRECTION
75+
assert link.attributes().asMap().get("link.kind") == SpanPointerUtils.LINK_KIND
76+
assert link.attributes().asMap().get("ptr.hash") != null
77+
}
78+
79+
@Test
80+
void testAddSpanPointerWithSingleNumberKey() {
81+
def mockData = createMockSpan()
82+
def mockSpan = mockData.span
83+
def links = mockData.links
6584

6685
def keys = [
6786
"count": AttributeValue.builder().n("42").build()
6887
]
6988

70-
DynamoDbUtil.exportTagsWithKnownKeys(mockSpan, keys)
89+
DynamoDbUtil.addSpanPointer(mockSpan, "my-table", keys)
7190

72-
assert tags[InstrumentationTags.DYNAMO_PRIMARY_KEY_1] == "count"
73-
assert tags[InstrumentationTags.DYNAMO_PRIMARY_KEY_1_VALUE] == "42"
91+
assert links.size() == 1
92+
def link = links[0]
93+
assert link.attributes().asMap().get("ptr.kind") == SpanPointerUtils.DYNAMODB_PTR_KIND
7494
}
7595

7696
@Test
77-
void testExportTagsWithSingleBinaryKey() {
97+
void testAddSpanPointerWithSingleBinaryKey() {
7898
def mockData = createMockSpan()
7999
def mockSpan = mockData.span
80-
def tags = mockData.tags
100+
def links = mockData.links
81101

82102
def binaryData = "binary-data".getBytes()
83103
def keys = [
84104
"data": AttributeValue.builder().b(SdkBytes.fromByteArray(binaryData)).build()
85105
]
86106

87-
DynamoDbUtil.exportTagsWithKnownKeys(mockSpan, keys)
107+
DynamoDbUtil.addSpanPointer(mockSpan, "my-table", keys)
88108

89-
assert tags[InstrumentationTags.DYNAMO_PRIMARY_KEY_1] == "data"
90-
assert tags[InstrumentationTags.DYNAMO_PRIMARY_KEY_1_VALUE] == "binary-data"
109+
assert links.size() == 1
110+
def link = links[0]
111+
assert link.attributes().asMap().get("ptr.kind") == SpanPointerUtils.DYNAMODB_PTR_KIND
91112
}
92113

93114
@Test
94-
void testExportTagsWithTwoKeys() {
115+
void testAddSpanPointerWithTwoKeys() {
95116
def mockData = createMockSpan()
96117
def mockSpan = mockData.span
97-
def tags = mockData.tags
118+
def links = mockData.links
98119

99120
def keys = [
100121
"id": AttributeValue.builder().s("12345").build(),
101122
"name": AttributeValue.builder().s("item-name").build()
102123
]
103124

104-
DynamoDbUtil.exportTagsWithKnownKeys(mockSpan, keys)
125+
DynamoDbUtil.addSpanPointer(mockSpan, "my-table", keys)
105126

106-
assert tags[InstrumentationTags.DYNAMO_PRIMARY_KEY_1] == "id"
107-
assert tags[InstrumentationTags.DYNAMO_PRIMARY_KEY_1_VALUE] == "12345"
108-
assert tags[InstrumentationTags.DYNAMO_PRIMARY_KEY_2] == "name"
109-
assert tags[InstrumentationTags.DYNAMO_PRIMARY_KEY_2_VALUE] == "item-name"
127+
assert links.size() == 1
128+
def link = links[0]
129+
assert link.attributes().asMap().get("ptr.kind") == SpanPointerUtils.DYNAMODB_PTR_KIND
130+
assert link.attributes().asMap().get("ptr.hash") != null
110131
}
111132

112133
@Test
113-
void testExportTagsWithTwoKeysSortsAlphabetically() {
134+
void testAddSpanPointerWithTwoKeysSortsAlphabetically() {
114135
def mockData = createMockSpan()
115-
def mockSpan = mockData.span
116-
def tags = mockData.tags
136+
def mockSpan1 = mockData.span
137+
def links1 = mockData.links
117138

118-
def keys = [
139+
// Keys in order bKey, aKey
140+
def keys1 = [
119141
"bKey": AttributeValue.builder().s("abc").build(),
120142
"aKey": AttributeValue.builder().s("zxy").build()
121143
]
122144

123-
DynamoDbUtil.exportTagsWithKnownKeys(mockSpan, keys)
145+
DynamoDbUtil.addSpanPointer(mockSpan1, "my-table", keys1)
146+
147+
// Reverse order: aKey, bKey — should produce the same hash
148+
def mockData2 = createMockSpan()
149+
def mockSpan2 = mockData2.span
150+
def links2 = mockData2.links
151+
152+
def keys2 = [
153+
"aKey": AttributeValue.builder().s("zxy").build(),
154+
"bKey": AttributeValue.builder().s("abc").build()
155+
]
156+
157+
DynamoDbUtil.addSpanPointer(mockSpan2, "my-table", keys2)
124158

125-
assert tags[InstrumentationTags.DYNAMO_PRIMARY_KEY_1] == "aKey"
126-
assert tags[InstrumentationTags.DYNAMO_PRIMARY_KEY_1_VALUE] == "zxy"
127-
assert tags[InstrumentationTags.DYNAMO_PRIMARY_KEY_2] == "bKey"
128-
assert tags[InstrumentationTags.DYNAMO_PRIMARY_KEY_2_VALUE] == "abc"
159+
assert links1.size() == 1
160+
assert links2.size() == 1
161+
// Both should produce the same hash regardless of input order
162+
assert links1[0].attributes().asMap().get("ptr.hash") == links2[0].attributes().asMap().get("ptr.hash")
129163
}
130164
}

‎dd-java-agent/instrumentation/aws-java/aws-java-s3-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/s3/S3Interceptor.java‎

Lines changed: 5 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,10 @@
11
package datadog.trace.instrumentation.aws.v2.s3;
22

3-
import static datadog.trace.bootstrap.instrumentation.api.InstrumentationTags.S3_ETAG;
4-
53
import datadog.context.Context;
64
import datadog.trace.api.Config;
75
import datadog.trace.bootstrap.InstanceStore;
86
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
7+
import datadog.trace.bootstrap.instrumentation.api.SpanPointerUtils;
98
import org.slf4j.Logger;
109
import org.slf4j.LoggerFactory;
1110
import software.amazon.awssdk.core.interceptor.Context.AfterExecution;
@@ -52,9 +51,9 @@ public void afterExecution(AfterExecution context, ExecutionAttributes execution
5251
return;
5352
}
5453

55-
// Store eTag as tag, then calculate hash + add span pointers in SpanPointersProcessor.
56-
// Bucket and key are already stored as tags in AwsSdkClientDecorator, so need to make redundant
57-
// tags.
58-
span.setTag(S3_ETAG, eTag);
54+
// Get bucket and key from the request to create the span pointer directly
55+
String bucket = context.request().getValueForField("Bucket", String.class).orElse(null);
56+
String key = context.request().getValueForField("Key", String.class).orElse(null);
57+
SpanPointerUtils.addS3SpanPointer(span, bucket, key, eTag);
5958
}
6059
}

0 commit comments

Comments
 (0)