Skip to content

Commit 288878e

Browse files
David-DAMDavid-DAM
committed
Add SNS and SQS context extraction and enrichment for OpenTelemetry tracing
1 parent c424115 commit 288878e

8 files changed

Lines changed: 264 additions & 72 deletions

File tree

‎powertools-tracing-opentelemetry/src/main/java/software/amazon/lambda/powertools/tracing/opentelemetry/TracingOpenTelemetry.java‎

Lines changed: 12 additions & 40 deletions
Original file line numberDiff line numberDiff line change
@@ -101,18 +101,13 @@ public SpanScope addSpan(String name) {
101101
}
102102

103103

104-
public SpanScope addSpan(
105-
String name,
106-
SpanKind kind) {
104+
public SpanScope addSpan(String name, SpanKind kind) {
107105

108106
return addSpan(name, kind, Attributes.empty());
109107
}
110108

111109

112-
public SpanScope addSpan(
113-
String name,
114-
SpanKind kind,
115-
Attributes attributes) {
110+
public SpanScope addSpan(String name, SpanKind kind, Attributes attributes) {
116111

117112
Objects.requireNonNull(name, "name must not be null");
118113
Objects.requireNonNull(kind, "kind must not be null");
@@ -126,11 +121,7 @@ public SpanScope addSpan(
126121
return new SpanScope(span);
127122
}
128123

129-
public SpanScope addSpan(
130-
String name,
131-
SpanKind kind,
132-
Attributes attributes,
133-
Context parentContext) {
124+
public SpanScope addSpan(String name, SpanKind kind, Attributes attributes, Context parentContext) {
134125

135126
Objects.requireNonNull(parentContext, "parentContext must not be null");
136127

@@ -144,9 +135,7 @@ public SpanScope addSpan(
144135
}
145136

146137

147-
public <T> T withSpan(
148-
String name,
149-
SpanOperation<T> operation) throws Exception {
138+
public <T> T withSpan(String name, SpanOperation<T> operation) throws Exception {
150139

151140
Objects.requireNonNull(operation, "operation must not be null");
152141

@@ -164,7 +153,8 @@ public <T> T captureLambdaHandler(
164153
String name,
165154
com.amazonaws.services.lambda.runtime.Context lambdaContext,
166155
io.opentelemetry.context.Context parentContext,
167-
SpanOperation<T> operation) throws Exception {
156+
SpanOperation<T> operation
157+
) throws Exception {
168158

169159
Objects.requireNonNull(name, "name must not be null");
170160
Objects.requireNonNull(parentContext, "parentContext must not be null");
@@ -191,51 +181,33 @@ public <T> T captureLambdaHandler(
191181
}
192182
}
193183

194-
public <T> Context extractContext(
195-
T carrier,
196-
TextMapGetter<T> getter) {
184+
public <T> Context extractContext(T carrier, TextMapGetter<T> getter) {
197185

198186
return extractContext(Context.current(), carrier, getter);
199187
}
200188

201189

202-
public <T> Context extractContext(
203-
Context context,
204-
T carrier,
205-
TextMapGetter<T> getter) {
190+
public <T> Context extractContext(Context context, T carrier, TextMapGetter<T> getter) {
206191

207192
Objects.requireNonNull(context, "context must not be null");
208193
Objects.requireNonNull(getter, "getter must not be null");
209194

210-
return propagator.extract(
211-
context,
212-
carrier,
213-
getter
214-
);
195+
return propagator.extract(context, carrier, getter);
215196
}
216197

217198

218-
public <T> void injectContext(
219-
T carrier,
220-
TextMapSetter<T> setter) {
199+
public <T> void injectContext(T carrier, TextMapSetter<T> setter) {
221200

222201
injectContext(Context.current(), carrier, setter);
223202
}
224203

225204

226-
public <T> void injectContext(
227-
Context context,
228-
T carrier,
229-
TextMapSetter<T> setter) {
205+
public <T> void injectContext(Context context, T carrier, TextMapSetter<T> setter) {
230206

231207
Objects.requireNonNull(context, "context must not be null");
232208
Objects.requireNonNull(setter, "setter must not be null");
233209

234-
propagator.inject(
235-
context,
236-
carrier,
237-
setter
238-
);
210+
propagator.inject(context, carrier, setter);
239211
}
240212

241213
private static TextMapPropagator createDefaultPropagator() {

‎powertools-tracing-opentelemetry/src/main/java/software/amazon/lambda/powertools/tracing/opentelemetry/context/ApiGatewayTraceContextExtractor.java‎

Lines changed: 2 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -3,34 +3,12 @@
33
import com.amazonaws.services.lambda.runtime.events.APIGatewayProxyRequestEvent;
44
import io.opentelemetry.api.trace.Span;
55
import io.opentelemetry.context.Context;
6-
import io.opentelemetry.context.propagation.TextMapGetter;
76
import io.opentelemetry.context.propagation.TextMapPropagator;
8-
import java.util.Collections;
97
import java.util.Map;
108
import java.util.stream.Collectors;
9+
import software.amazon.lambda.powertools.tracing.opentelemetry.provider.OpenTelemetryProvider;
1110

1211
public final class ApiGatewayTraceContextExtractor implements LambdaEventContextExtractor {
13-
private static final TextMapGetter<Map<String, String>> HEADER_GETTER =
14-
new TextMapGetter<>() {
15-
16-
@Override
17-
public Iterable<String> keys(Map<String, String> carrier) {
18-
19-
return carrier != null
20-
? carrier.keySet()
21-
: Collections.emptyList();
22-
}
23-
24-
@Override
25-
public String get(Map<String, String> carrier, String key) {
26-
27-
if (carrier == null) {
28-
return null;
29-
}
30-
31-
return carrier.get(key);
32-
}
33-
};
3412

3513

3614
@Override
@@ -52,7 +30,7 @@ public Context extract(Object event, Context parentContext, TextMapPropagator pr
5230
return propagator.extract(
5331
parentContext,
5432
headers,
55-
HEADER_GETTER
33+
OpenTelemetryProvider.textMapGetter()
5634
);
5735
}
5836

‎powertools-tracing-opentelemetry/src/main/java/software/amazon/lambda/powertools/tracing/opentelemetry/context/LambdaEventContextExtractorResolver.java‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,9 @@ public LambdaEventContextExtractorResolver(List<LambdaEventContextExtractor> ext
1717
public static LambdaEventContextExtractorResolver create() {
1818
return new LambdaEventContextExtractorResolver(
1919
List.of(
20-
new ApiGatewayTraceContextExtractor()
20+
new ApiGatewayTraceContextExtractor(),
21+
new SqsTraceContextExtractor(),
22+
new SnsTraceContextExtractor()
2123
)
2224
);
2325
}

‎powertools-tracing-opentelemetry/src/main/java/software/amazon/lambda/powertools/tracing/opentelemetry/context/SnsTraceContextExtractor.java‎

Lines changed: 84 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,84 @@
1+
package software.amazon.lambda.powertools.tracing.opentelemetry.context;
2+
3+
import com.amazonaws.services.lambda.runtime.events.SNSEvent;
4+
import io.opentelemetry.api.trace.Span;
5+
import io.opentelemetry.context.Context;
6+
import io.opentelemetry.context.propagation.TextMapPropagator;
7+
import java.util.Map;
8+
import java.util.stream.Collectors;
9+
import software.amazon.lambda.powertools.tracing.opentelemetry.provider.OpenTelemetryProvider;
10+
11+
public final class SnsTraceContextExtractor implements LambdaEventContextExtractor {
12+
13+
@Override
14+
public boolean supports(Object event) {
15+
return event instanceof SNSEvent;
16+
}
17+
18+
@Override
19+
public Context extract(Object event, Context parentContext, TextMapPropagator propagator) {
20+
21+
SNSEvent snsEvent = (SNSEvent) event;
22+
23+
if (snsEvent.getRecords() == null || snsEvent.getRecords().isEmpty()) {
24+
return parentContext;
25+
}
26+
27+
SNSEvent.SNSRecord record = snsEvent.getRecords().get(0);
28+
29+
Map<String, SNSEvent.MessageAttribute> attributes = record.getSNS().getMessageAttributes();
30+
31+
if (attributes == null || attributes.isEmpty()) {
32+
return parentContext;
33+
}
34+
35+
Map<String, String> propagationAttributes = attributes.entrySet()
36+
.stream()
37+
.filter(entry -> entry.getValue() != null)
38+
.filter(entry -> entry.getValue().getValue() != null)
39+
.collect(Collectors.toMap(
40+
Map.Entry::getKey,
41+
entry -> entry.getValue().getValue()
42+
));
43+
44+
return propagator.extract(
45+
parentContext,
46+
propagationAttributes,
47+
OpenTelemetryProvider.textMapGetter()
48+
);
49+
}
50+
51+
@Override
52+
public void enrichSpan(Object event, Span span) {
53+
54+
SNSEvent snsEvent = (SNSEvent) event;
55+
56+
if (snsEvent.getRecords() == null || snsEvent.getRecords().isEmpty()) {
57+
return;
58+
}
59+
60+
SNSEvent.SNSRecord record = snsEvent.getRecords().get(0);
61+
62+
if (record.getSNS() == null) {
63+
return;
64+
}
65+
66+
span.setAttribute("messaging.system", "aws.sns");
67+
68+
if (record.getSNS().getMessageId() != null) {
69+
span.setAttribute("messaging.message.id", record.getSNS().getMessageId());
70+
}
71+
72+
if (record.getSNS().getTopicArn() != null) {
73+
span.setAttribute("messaging.destination.name", extractTopicName(record.getSNS().getTopicArn()));
74+
}
75+
}
76+
77+
private String extractTopicName(String topicArn) {
78+
int separator = topicArn.lastIndexOf(':');
79+
80+
return separator >= 0
81+
? topicArn.substring(separator + 1)
82+
: topicArn;
83+
}
84+
}

‎powertools-tracing-opentelemetry/src/main/java/software/amazon/lambda/powertools/tracing/opentelemetry/context/SqsTraceContextExtractor.java‎

Lines changed: 78 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,78 @@
1+
package software.amazon.lambda.powertools.tracing.opentelemetry.context;
2+
3+
import com.amazonaws.services.lambda.runtime.events.SQSEvent;
4+
import io.opentelemetry.api.trace.Span;
5+
import io.opentelemetry.context.Context;
6+
import io.opentelemetry.context.propagation.TextMapPropagator;
7+
import java.util.Map;
8+
import java.util.stream.Collectors;
9+
import software.amazon.lambda.powertools.tracing.opentelemetry.provider.OpenTelemetryProvider;
10+
11+
public final class SqsTraceContextExtractor implements LambdaEventContextExtractor {
12+
13+
@Override
14+
public boolean supports(Object event) {
15+
return event instanceof SQSEvent;
16+
}
17+
18+
@Override
19+
public Context extract(Object event, Context parentContext, TextMapPropagator propagator) {
20+
21+
SQSEvent sqsEvent = (SQSEvent) event;
22+
23+
if (sqsEvent.getRecords() == null || sqsEvent.getRecords().isEmpty()) {
24+
return parentContext;
25+
}
26+
27+
SQSEvent.SQSMessage message = sqsEvent.getRecords().get(0);
28+
29+
Map<String, SQSEvent.MessageAttribute> attributes = message.getMessageAttributes();
30+
31+
if (attributes == null || attributes.isEmpty()) {
32+
return parentContext;
33+
}
34+
35+
Map<String, String> propagationAttributes = attributes.entrySet()
36+
.stream()
37+
.filter(entry -> entry.getValue() != null)
38+
.collect(Collectors.toMap(
39+
Map.Entry::getKey,
40+
entry -> entry.getValue().getStringValue()
41+
));
42+
43+
return propagator.extract(
44+
parentContext,
45+
propagationAttributes,
46+
OpenTelemetryProvider.textMapGetter()
47+
);
48+
}
49+
50+
@Override
51+
public void enrichSpan(Object event, Span span) {
52+
SQSEvent sqsEvent = (SQSEvent) event;
53+
54+
if (sqsEvent.getRecords() == null || sqsEvent.getRecords().isEmpty()) {
55+
return;
56+
}
57+
58+
SQSEvent.SQSMessage message = sqsEvent.getRecords().get(0);
59+
60+
span.setAttribute("messaging.system", "aws.sqs");
61+
62+
if (message.getMessageId() != null) {
63+
span.setAttribute("messaging.message.id", message.getMessageId());
64+
}
65+
66+
if (message.getEventSourceArn() != null) {
67+
span.setAttribute("messaging.destination.name", extractQueueName(message.getEventSourceArn()));
68+
}
69+
}
70+
71+
private String extractQueueName(String arn) {
72+
int separator = arn.lastIndexOf(':');
73+
74+
return separator >= 0
75+
? arn.substring(separator + 1)
76+
: arn;
77+
}
78+
}

‎powertools-tracing-opentelemetry/src/main/java/software/amazon/lambda/powertools/tracing/opentelemetry/internal/TracingOpenTelemetryAspect.java‎

Lines changed: 5 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -66,9 +66,7 @@ private Object traceHandler(ProceedingJoinPoint pjp, Tracing tracing, String spa
6666

6767
Span span = scope.span();
6868

69-
Object event = pjp.getArgs()[0];
70-
71-
tracingOtel.eventContextExtractorResolver().enrichSpan(event, span);
69+
tracingOtel.eventContextExtractorResolver().enrichSpan(pjp.getArgs()[0], span);
7270

7371
addLambdaInvocationAttributes(pjp, span);
7472

@@ -161,7 +159,7 @@ private void captureResponse(
161159
Tracing tracing,
162160
Object response) throws Exception {
163161

164-
if (!captureResponse(tracing)) {
162+
if (!isCaptureResponseEnabled(tracing)) {
165163
return;
166164
}
167165

@@ -176,12 +174,12 @@ private void captureError(
176174
Tracing tracing,
177175
Throwable throwable) {
178176

179-
if (captureError(tracing)) {
177+
if (isCaptureErrorEnabled(tracing)) {
180178
scope.recordException(throwable);
181179
}
182180
}
183181

184-
private boolean captureResponse(Tracing tracing) {
182+
private boolean isCaptureResponseEnabled(Tracing tracing) {
185183
switch (tracing.captureMode()) {
186184
case ENVIRONMENT_VAR:
187185
return isEnvironmentVariableSet(
@@ -200,7 +198,7 @@ && environmentVariable(
200198
}
201199
}
202200

203-
private boolean captureError(Tracing tracing) {
201+
private boolean isCaptureErrorEnabled(Tracing tracing) {
204202
switch (tracing.captureMode()) {
205203
case ENVIRONMENT_VAR:
206204
return isEnvironmentVariableSet(

0 commit comments

Comments
 (0)