Skip to content

Commit 706d1b9

Browse files
committed
update awslambda to use _X_AMZN_TRACE_ID as a Span Link
1 parent df32e8c commit 706d1b9

File tree

2 files changed

+103
-83
lines changed

2 files changed

+103
-83
lines changed

instrumentation/opentelemetry-instrumentation-aws-lambda/src/opentelemetry/instrumentation/aws_lambda/__init__.py

Lines changed: 38 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -71,7 +71,7 @@ def custom_event_context_extractor(lambda_event):
7171
import os
7272
import time
7373
from importlib import import_module
74-
from typing import Any, Callable, Collection
74+
from typing import Any, Callable, Collection, Optional, Sequence
7575
from urllib.parse import urlencode
7676

7777
from wrapt import wrap_function_wrapper
@@ -90,9 +90,11 @@ def custom_event_context_extractor(lambda_event):
9090
from opentelemetry.semconv.resource import ResourceAttributes
9191
from opentelemetry.semconv.trace import SpanAttributes
9292
from opentelemetry.trace import (
93+
Link,
9394
Span,
9495
SpanKind,
9596
TracerProvider,
97+
get_current_span,
9698
get_tracer,
9799
get_tracer_provider,
98100
)
@@ -149,7 +151,7 @@ def _determine_parent_context(
149151
"""Determine the parent context for the current Lambda invocation.
150152
151153
See more:
152-
https://github.com/open-telemetry/opentelemetry-specification/blob/main/specification/trace/semantic_conventions/instrumentation/aws-lambda.md#determining-the-parent-of-a-span
154+
https://github.com/open-telemetry/opentelemetry-specification/blob/main/specification/trace/semantic_conventions/instrumentation/aws-lambda.md
153155
154156
Args:
155157
lambda_event: user-defined, so it could be anything, but this
@@ -166,28 +168,45 @@ def _determine_parent_context(
166168
"""
167169
parent_context = None
168170

171+
if event_context_extractor:
172+
parent_context = event_context_extractor(lambda_event)
173+
else:
174+
parent_context = _default_event_context_extractor(lambda_event)
175+
176+
return parent_context
177+
178+
179+
def _determine_links(
180+
disable_aws_context_propagation: bool = False,
181+
) -> Optional[Sequence[Link]]:
182+
"""Determine if a Link should be added to the Span based on the
183+
environment variable `_X_AMZN_TRACE_ID`.
184+
185+
186+
See more:
187+
https://github.com/open-telemetry/opentelemetry-specification/blob/main/specification/trace/semantic_conventions/instrumentation/aws-lambda.md
188+
189+
Args:
190+
disable_aws_context_propagation: By default, this instrumentation
191+
will try to read the context from the `_X_AMZN_TRACE_ID` environment
192+
variable set by Lambda, set this to `True` to disable this behavior.
193+
Returns:
194+
A Link or None
195+
"""
196+
env_context = None
197+
links = None
198+
169199
if not disable_aws_context_propagation:
170200
xray_env_var = os.environ.get(_X_AMZN_TRACE_ID)
171201

172202
if xray_env_var:
173-
parent_context = AwsXRayPropagator().extract(
203+
env_context = AwsXRayPropagator().extract(
174204
{TRACE_HEADER_KEY: xray_env_var}
175205
)
176206

177-
if (
178-
parent_context
179-
and get_current_span(parent_context)
180-
.get_span_context()
181-
.trace_flags.sampled
182-
):
183-
return parent_context
184-
185-
if event_context_extractor:
186-
parent_context = event_context_extractor(lambda_event)
187-
else:
188-
parent_context = _default_event_context_extractor(lambda_event)
207+
links = [Link(get_current_span(env_context).get_span_context())]
189208

190-
return parent_context
209+
return links
191210

192211

193212
def _set_api_gateway_v1_proxy_attributes(
@@ -294,6 +313,8 @@ def _instrumented_lambda_handler_call( # noqa pylint: disable=too-many-branches
294313
disable_aws_context_propagation,
295314
)
296315

316+
links = _determine_links(disable_aws_context_propagation)
317+
297318
span_kind = None
298319
try:
299320
if lambda_event["Records"][0]["eventSource"] in {
@@ -319,6 +340,7 @@ def _instrumented_lambda_handler_call( # noqa pylint: disable=too-many-branches
319340
name=orig_handler_name,
320341
context=parent_context,
321342
kind=span_kind,
343+
links=links,
322344
) as span:
323345
if span.is_recording():
324346
lambda_context = args[1]

instrumentation/opentelemetry-instrumentation-aws-lambda/tests/test_aws_lambda_instrumentation_manual.py

Lines changed: 65 additions & 67 deletions
Original file line numberDiff line numberDiff line change
@@ -138,7 +138,9 @@ def test_active_tracing(self):
138138
self.assertEqual(len(spans), 1)
139139
span = spans[0]
140140
self.assertEqual(span.name, os.environ[_HANDLER])
141-
self.assertEqual(span.get_span_context().trace_id, MOCK_XRAY_TRACE_ID)
141+
self.assertNotEqual(
142+
span.get_span_context().trace_id, MOCK_XRAY_TRACE_ID
143+
)
142144
self.assertEqual(span.kind, SpanKind.SERVER)
143145
self.assertSpanHasAttributes(
144146
span,
@@ -149,11 +151,7 @@ def test_active_tracing(self):
149151
)
150152

151153
parent_context = span.parent
152-
self.assertEqual(
153-
parent_context.trace_id, span.get_span_context().trace_id
154-
)
155-
self.assertEqual(parent_context.span_id, MOCK_XRAY_PARENT_SPAN_ID)
156-
self.assertTrue(parent_context.is_remote)
154+
self.assertEqual(None, parent_context)
157155

158156
test_env_patch.stop()
159157

@@ -165,11 +163,8 @@ class TestCase:
165163
context: Dict
166164
expected_traceid: int
167165
expected_parentid: int
168-
xray_traceid: str
169166
expected_state_value: str = None
170167
expected_trace_state_len: int = 0
171-
disable_aws_context_propagation: bool = False
172-
disable_aws_context_propagation_envvar: str = ""
173168

174169
def custom_event_context_extractor(lambda_event):
175170
return get_global_textmap().extract(lambda_event["foo"]["headers"])
@@ -188,42 +183,9 @@ def custom_event_context_extractor(lambda_event):
188183
expected_parentid=MOCK_W3C_PARENT_SPAN_ID,
189184
expected_trace_state_len=3,
190185
expected_state_value=MOCK_W3C_TRACE_STATE_VALUE,
191-
xray_traceid=MOCK_XRAY_TRACE_CONTEXT_NOT_SAMPLED,
192-
),
193-
TestCase(
194-
name="custom_extractor_not_sampled_xray",
195-
custom_extractor=custom_event_context_extractor,
196-
context={
197-
"foo": {
198-
"headers": {
199-
TraceContextTextMapPropagator._TRACEPARENT_HEADER_NAME: MOCK_W3C_TRACE_CONTEXT_SAMPLED,
200-
TraceContextTextMapPropagator._TRACESTATE_HEADER_NAME: f"{MOCK_W3C_TRACE_STATE_KEY}={MOCK_W3C_TRACE_STATE_VALUE},foo=1,bar=2",
201-
}
202-
}
203-
},
204-
expected_traceid=MOCK_W3C_TRACE_ID,
205-
expected_parentid=MOCK_W3C_PARENT_SPAN_ID,
206-
expected_trace_state_len=3,
207-
expected_state_value=MOCK_W3C_TRACE_STATE_VALUE,
208-
xray_traceid=MOCK_XRAY_TRACE_CONTEXT_NOT_SAMPLED,
209-
),
210-
TestCase(
211-
name="custom_extractor_sampled_xray",
212-
custom_extractor=custom_event_context_extractor,
213-
context={
214-
"foo": {
215-
"headers": {
216-
TraceContextTextMapPropagator._TRACEPARENT_HEADER_NAME: MOCK_W3C_TRACE_CONTEXT_SAMPLED,
217-
TraceContextTextMapPropagator._TRACESTATE_HEADER_NAME: f"{MOCK_W3C_TRACE_STATE_KEY}={MOCK_W3C_TRACE_STATE_VALUE},foo=1,bar=2",
218-
}
219-
}
220-
},
221-
expected_traceid=MOCK_XRAY_TRACE_ID,
222-
expected_parentid=MOCK_XRAY_PARENT_SPAN_ID,
223-
xray_traceid=MOCK_XRAY_TRACE_CONTEXT_SAMPLED,
224186
),
225187
TestCase(
226-
name="custom_extractor_sampled_xray_disable_aws_propagation",
188+
name="custom_extractor",
227189
custom_extractor=custom_event_context_extractor,
228190
context={
229191
"foo": {
@@ -233,47 +195,24 @@ def custom_event_context_extractor(lambda_event):
233195
}
234196
}
235197
},
236-
disable_aws_context_propagation=True,
237-
expected_traceid=MOCK_W3C_TRACE_ID,
238-
expected_parentid=MOCK_W3C_PARENT_SPAN_ID,
239-
expected_trace_state_len=3,
240-
expected_state_value=MOCK_W3C_TRACE_STATE_VALUE,
241-
xray_traceid=MOCK_XRAY_TRACE_CONTEXT_SAMPLED,
242-
),
243-
TestCase(
244-
name="no_custom_extractor_xray_disable_aws_propagation_via_env_var",
245-
custom_extractor=None,
246-
context={
247-
"headers": {
248-
TraceContextTextMapPropagator._TRACEPARENT_HEADER_NAME: MOCK_W3C_TRACE_CONTEXT_SAMPLED,
249-
TraceContextTextMapPropagator._TRACESTATE_HEADER_NAME: f"{MOCK_W3C_TRACE_STATE_KEY}={MOCK_W3C_TRACE_STATE_VALUE},foo=1,bar=2",
250-
}
251-
},
252-
disable_aws_context_propagation=False,
253-
disable_aws_context_propagation_envvar="true",
254198
expected_traceid=MOCK_W3C_TRACE_ID,
255199
expected_parentid=MOCK_W3C_PARENT_SPAN_ID,
256200
expected_trace_state_len=3,
257201
expected_state_value=MOCK_W3C_TRACE_STATE_VALUE,
258-
xray_traceid=MOCK_XRAY_TRACE_CONTEXT_SAMPLED,
259202
),
260203
]
261204
for test in tests:
262205
test_env_patch = mock.patch.dict(
263206
"os.environ",
264207
{
265208
**os.environ,
266-
# NOT Active Tracing
267-
_X_AMZN_TRACE_ID: test.xray_traceid,
268-
OTEL_LAMBDA_DISABLE_AWS_CONTEXT_PROPAGATION: test.disable_aws_context_propagation_envvar,
269209
# NOT using the X-Ray Propagator
270210
OTEL_PROPAGATORS: "tracecontext",
271211
},
272212
)
273213
test_env_patch.start()
274214
AwsLambdaInstrumentor().instrument(
275-
event_context_extractor=test.custom_extractor,
276-
disable_aws_context_propagation=test.disable_aws_context_propagation,
215+
event_context_extractor=test.custom_extractor
277216
)
278217
mock_execute_lambda(test.context)
279218
spans = self.memory_exporter.get_finished_spans()
@@ -301,6 +240,65 @@ def custom_event_context_extractor(lambda_event):
301240
AwsLambdaInstrumentor().uninstrument()
302241
test_env_patch.stop()
303242

243+
def test_links_from_lambda_event(self):
244+
@dataclass
245+
class TestCase:
246+
name: str
247+
context: Dict
248+
expected_link_trace_id: int
249+
xray_traceid: str
250+
disable_aws_context_propagation: bool = False
251+
disable_aws_context_propagation_envvar: str = ""
252+
253+
tests = [
254+
TestCase(
255+
name="enabled_aws_context_propagation",
256+
context={},
257+
expected_link_trace_id=MOCK_XRAY_TRACE_ID,
258+
xray_traceid=MOCK_XRAY_TRACE_CONTEXT_NOT_SAMPLED,
259+
),
260+
TestCase(
261+
name="disable_aws_context_propagation",
262+
context={},
263+
expected_link_trace_id=None,
264+
xray_traceid=MOCK_XRAY_TRACE_CONTEXT_NOT_SAMPLED,
265+
disable_aws_context_propagation=True,
266+
),
267+
]
268+
for test in tests:
269+
test_env_patch = mock.patch.dict(
270+
"os.environ",
271+
{
272+
**os.environ,
273+
# NOT Active Tracing
274+
_X_AMZN_TRACE_ID: test.xray_traceid,
275+
OTEL_LAMBDA_DISABLE_AWS_CONTEXT_PROPAGATION: test.disable_aws_context_propagation_envvar,
276+
# NOT using the X-Ray Propagator
277+
OTEL_PROPAGATORS: "tracecontext",
278+
},
279+
)
280+
test_env_patch.start()
281+
AwsLambdaInstrumentor().instrument(
282+
disable_aws_context_propagation=test.disable_aws_context_propagation,
283+
)
284+
mock_execute_lambda(test.context)
285+
spans = self.memory_exporter.get_finished_spans()
286+
assert spans
287+
self.assertEqual(len(spans), 1)
288+
span = spans[0]
289+
290+
if test.expected_link_trace_id == None:
291+
self.assertEqual(0, len(span.links))
292+
else:
293+
link = span.links[0]
294+
self.assertEqual(
295+
link.context.trace_id, test.expected_link_trace_id
296+
)
297+
298+
self.memory_exporter.clear()
299+
AwsLambdaInstrumentor().uninstrument()
300+
test_env_patch.stop()
301+
304302
def test_lambda_no_error_with_invalid_flush_timeout(self):
305303
test_env_patch = mock.patch.dict(
306304
"os.environ",

0 commit comments

Comments
 (0)