Skip to content

Commit 883592d

Browse files
rapphilrapphil
committed
Feat: Use the lambda runtime as parent context as opt-in
In a previous PR open-telemetry#1657 the lambda behaviour was changed so that it extracs the context from the headers of the lambda event instead of from the lambda runtime, using the _X_AMZN_TRACE_ID env var. This behaviour is undesireble as it is a breaking change to existing users. This PR will add the previous behaviour through an opt-in flag so that users can be gracefully migrated or these two "modes" of operation can exist without conflict. Signed-off-by: Raphael Silva <[email protected]>
1 parent 1beab82 commit 883592d

File tree

2 files changed

+122
-75
lines changed

2 files changed

+122
-75
lines changed

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

Lines changed: 47 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -107,6 +107,7 @@ def custom_event_context_extractor(lambda_event):
107107
OTEL_INSTRUMENTATION_AWS_LAMBDA_FLUSH_TIMEOUT = (
108108
"OTEL_INSTRUMENTATION_AWS_LAMBDA_FLUSH_TIMEOUT"
109109
)
110+
OTEL_LAMBDA_USE_AWS_CONTEXT_PROPAGATION = "OTEL_LAMBDA_USE_AWS_CONTEXT_PROPAGATION"
110111

111112

112113
def _default_event_context_extractor(lambda_event: Any) -> Context:
@@ -140,7 +141,9 @@ def _default_event_context_extractor(lambda_event: Any) -> Context:
140141

141142

142143
def _determine_parent_context(
143-
lambda_event: Any, event_context_extractor: Callable[[Any], Context]
144+
lambda_event: Any,
145+
use_aws_context_propagation: bool,
146+
event_context_extractor: Callable[[Any], Context],
144147
) -> Context:
145148
"""Determine the parent context for the current Lambda invocation.
146149
@@ -158,15 +161,26 @@ def _determine_parent_context(
158161
A Context with configuration found in the carrier.
159162
"""
160163
parent_context = None
161-
162-
if event_context_extractor:
164+
if use_aws_context_propagation:
165+
parent_context = _get_x_ray_context()
166+
elif event_context_extractor:
163167
parent_context = event_context_extractor(lambda_event)
164168
else:
165169
parent_context = _default_event_context_extractor(lambda_event)
166170

167171
return parent_context
168172

169173

174+
def _get_x_ray_context() -> Optional[Context]:
175+
"""Determine teh context propagated through the lambda runtime"""
176+
xray_env_var = os.environ.get(_X_AMZN_TRACE_ID)
177+
if xray_env_var:
178+
env_context = AwsXRayPropagator().extract({TRACE_HEADER_KEY: xray_env_var})
179+
return env_context
180+
181+
return None
182+
183+
170184
def _determine_links() -> Optional[Sequence[Link]]:
171185
"""Determine if a Link should be added to the Span based on the
172186
environment variable `_X_AMZN_TRACE_ID`.
@@ -180,31 +194,23 @@ def _determine_links() -> Optional[Sequence[Link]]:
180194
"""
181195
links = None
182196

183-
xray_env_var = os.environ.get(_X_AMZN_TRACE_ID)
197+
x_ray_context = _get_x_ray_context()
184198

185-
if xray_env_var:
186-
env_context = AwsXRayPropagator().extract(
187-
{TRACE_HEADER_KEY: xray_env_var}
188-
)
189-
190-
span_context = get_current_span(env_context).get_span_context()
199+
if x_ray_context:
200+
span_context = get_current_span(x_ray_context).get_span_context()
191201
if span_context.is_valid:
192202
links = [Link(span_context, {"source": "x-ray-env"})]
193203

194204
return links
195205

196206

197-
def _set_api_gateway_v1_proxy_attributes(
198-
lambda_event: Any, span: Span
199-
) -> Span:
207+
def _set_api_gateway_v1_proxy_attributes(lambda_event: Any, span: Span) -> Span:
200208
"""Sets HTTP attributes for REST APIs and v1 HTTP APIs
201209
202210
More info:
203211
https://docs.aws.amazon.com/apigateway/latest/developerguide/set-up-lambda-proxy-integrations.html#api-gateway-simple-proxy-for-lambda-input-format
204212
"""
205-
span.set_attribute(
206-
SpanAttributes.HTTP_METHOD, lambda_event.get("httpMethod")
207-
)
213+
span.set_attribute(SpanAttributes.HTTP_METHOD, lambda_event.get("httpMethod"))
208214

209215
if lambda_event.get("headers"):
210216
if "User-Agent" in lambda_event["headers"]:
@@ -231,16 +237,12 @@ def _set_api_gateway_v1_proxy_attributes(
231237
f"{lambda_event['resource']}?{urlencode(lambda_event['queryStringParameters'])}",
232238
)
233239
else:
234-
span.set_attribute(
235-
SpanAttributes.HTTP_TARGET, lambda_event["resource"]
236-
)
240+
span.set_attribute(SpanAttributes.HTTP_TARGET, lambda_event["resource"])
237241

238242
return span
239243

240244

241-
def _set_api_gateway_v2_proxy_attributes(
242-
lambda_event: Any, span: Span
243-
) -> Span:
245+
def _set_api_gateway_v2_proxy_attributes(lambda_event: Any, span: Span) -> Span:
244246
"""Sets HTTP attributes for v2 HTTP APIs
245247
246248
More info:
@@ -289,21 +291,26 @@ def _instrument(
289291
event_context_extractor: Callable[[Any], Context],
290292
tracer_provider: TracerProvider = None,
291293
meter_provider: MeterProvider = None,
294+
use_aws_context_propagation: bool = False,
292295
):
293296
def _instrumented_lambda_handler_call( # noqa pylint: disable=too-many-branches
294297
call_wrapped, instance, args, kwargs
295298
):
296-
orig_handler_name = ".".join(
297-
[wrapped_module_name, wrapped_function_name]
298-
)
299+
orig_handler_name = ".".join([wrapped_module_name, wrapped_function_name])
299300

300301
lambda_event = args[0]
301302

303+
# We are not fully complying with the specification here to be backwards
304+
# compatible with the old version of the specification.
305+
# the ``use_aws_context_propagation`` flag allow us
306+
# to opt-in into the previous behavior
302307
parent_context = _determine_parent_context(
303-
lambda_event, event_context_extractor
308+
lambda_event, use_aws_context_propagation, event_context_extractor
304309
)
305310

306-
links = _determine_links()
311+
links = None
312+
if not use_aws_context_propagation:
313+
links = _determine_links()
307314

308315
span_kind = None
309316
try:
@@ -354,9 +361,7 @@ def _instrumented_lambda_handler_call( # noqa pylint: disable=too-many-branches
354361
# If the request came from an API Gateway, extract http attributes from the event
355362
# https://github.com/open-telemetry/opentelemetry-specification/blob/main/specification/trace/semantic_conventions/instrumentation/aws-lambda.md#api-gateway
356363
# https://github.com/open-telemetry/opentelemetry-specification/blob/main/specification/trace/semantic_conventions/http.md#http-server-semantic-conventions
357-
if isinstance(lambda_event, dict) and lambda_event.get(
358-
"requestContext"
359-
):
364+
if isinstance(lambda_event, dict) and lambda_event.get("requestContext"):
360365
span.set_attribute(SpanAttributes.FAAS_TRIGGER, "http")
361366

362367
if lambda_event.get("version") == "2.0":
@@ -424,6 +429,13 @@ def _instrument(self, **kwargs):
424429
Event as input and extracts an OTel Context from it. By default,
425430
the context is extracted from the HTTP headers of an API Gateway
426431
request.
432+
``use_aws_context_propagation``: whether to use the AWS context propagation
433+
to populate the parent context. When set to true, the spans
434+
from the lambda runtime will not be added as span link to the
435+
span of the lambda invocation.
436+
Defaults to False.
437+
"""
438+
_instrument(**kwargs)
427439
"""
428440
lambda_handler = os.environ.get(ORIG_HANDLER, os.environ.get(_HANDLER))
429441
# pylint: disable=attribute-defined-outside-init
@@ -454,6 +466,11 @@ def _instrument(self, **kwargs):
454466
),
455467
tracer_provider=kwargs.get("tracer_provider"),
456468
meter_provider=kwargs.get("meter_provider"),
469+
use_aws_context_propagation=kwargs.get(
470+
"use_aws_context_propagation",
471+
os.environ.get(OTEL_LAMBDA_USE_AWS_CONTEXT_PROPAGATION, "False").lower()
472+
in ("true", "1", "t"),
473+
),
457474
)
458475

459476
def _uninstrument(self, **kwargs):

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

Lines changed: 75 additions & 45 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@
2727
_HANDLER,
2828
_X_AMZN_TRACE_ID,
2929
OTEL_INSTRUMENTATION_AWS_LAMBDA_FLUSH_TIMEOUT,
30+
OTEL_LAMBDA_USE_AWS_CONTEXT_PROPAGATION,
3031
AwsLambdaInstrumentor,
3132
)
3233
from opentelemetry.propagate import get_global_textmap
@@ -59,9 +60,7 @@ def __init__(self, aws_request_id, invoked_function_arn):
5960
MOCK_XRAY_PARENT_SPAN_ID = 0x3328B8445A6DBAD2
6061
MOCK_XRAY_TRACE_CONTEXT_COMMON = f"Root={TRACE_ID_VERSION}-{MOCK_XRAY_TRACE_ID_STR[:TRACE_ID_FIRST_PART_LENGTH]}-{MOCK_XRAY_TRACE_ID_STR[TRACE_ID_FIRST_PART_LENGTH:]};Parent={MOCK_XRAY_PARENT_SPAN_ID:x}"
6162
MOCK_XRAY_TRACE_CONTEXT_SAMPLED = f"{MOCK_XRAY_TRACE_CONTEXT_COMMON};Sampled=1"
62-
MOCK_XRAY_TRACE_CONTEXT_NOT_SAMPLED = (
63-
f"{MOCK_XRAY_TRACE_CONTEXT_COMMON};Sampled=0"
64-
)
63+
MOCK_XRAY_TRACE_CONTEXT_NOT_SAMPLED = f"{MOCK_XRAY_TRACE_CONTEXT_COMMON};Sampled=0"
6564

6665
# See more:
6766
# https://www.w3.org/TR/trace-context/#examples-of-http-traceparent-headers
@@ -116,43 +115,70 @@ def tearDown(self):
116115
AwsLambdaInstrumentor().uninstrument()
117116

118117
def test_active_tracing(self):
119-
test_env_patch = mock.patch.dict(
120-
"os.environ",
121-
{
122-
**os.environ,
123-
# Using Active tracing
124-
_X_AMZN_TRACE_ID: MOCK_XRAY_TRACE_CONTEXT_SAMPLED,
125-
},
126-
)
127-
test_env_patch.start()
118+
@dataclass
119+
class Testcase:
120+
name: str
121+
use_aws_context_propagation: str
122+
expected_trace_id: str
128123

129-
AwsLambdaInstrumentor().instrument()
124+
tests = [
125+
Testcase(
126+
name="Use aws context propgation",
127+
use_aws_context_propagation="true",
128+
expected_trace_id=MOCK_XRAY_TRACE_ID,
129+
),
130+
Testcase(
131+
name="Do not use aws context propgation",
132+
use_aws_context_propagation="false",
133+
expected_trace_id=None,
134+
),
135+
]
130136

131-
mock_execute_lambda()
137+
for test in tests:
138+
test_env_patch = mock.patch.dict(
139+
"os.environ",
140+
{
141+
**os.environ,
142+
# Using Active tracing
143+
_X_AMZN_TRACE_ID: MOCK_XRAY_TRACE_CONTEXT_SAMPLED,
144+
OTEL_LAMBDA_USE_AWS_CONTEXT_PROPAGATION: test.use_aws_context_propagation,
145+
},
146+
)
147+
test_env_patch.start()
132148

133-
spans = self.memory_exporter.get_finished_spans()
149+
AwsLambdaInstrumentor().instrument()
134150

135-
assert spans
151+
mock_execute_lambda()
136152

137-
self.assertEqual(len(spans), 1)
138-
span = spans[0]
139-
self.assertEqual(span.name, os.environ[_HANDLER])
140-
self.assertNotEqual(
141-
span.get_span_context().trace_id, MOCK_XRAY_TRACE_ID
142-
)
143-
self.assertEqual(span.kind, SpanKind.SERVER)
144-
self.assertSpanHasAttributes(
145-
span,
146-
{
147-
ResourceAttributes.FAAS_ID: MOCK_LAMBDA_CONTEXT.invoked_function_arn,
148-
SpanAttributes.FAAS_EXECUTION: MOCK_LAMBDA_CONTEXT.aws_request_id,
149-
},
150-
)
153+
spans = self.memory_exporter.get_finished_spans()
151154

152-
parent_context = span.parent
153-
self.assertEqual(None, parent_context)
155+
assert spans
154156

155-
test_env_patch.stop()
157+
self.assertEqual(len(spans), 1)
158+
span = spans[0]
159+
self.assertEqual(span.name, os.environ[_HANDLER])
160+
parent_context = span.parent
161+
if test.expected_trace_id is None:
162+
self.assertNotEqual(
163+
span.get_span_context().trace_id, MOCK_XRAY_TRACE_ID
164+
)
165+
self.assertEqual(None, parent_context)
166+
else:
167+
self.assertEqual(span.get_span_context().trace_id, MOCK_XRAY_TRACE_ID)
168+
self.assertEqual(
169+
parent_context.trace_id, span.get_span_context().trace_id
170+
)
171+
self.assertEqual(span.kind, SpanKind.SERVER)
172+
self.assertSpanHasAttributes(
173+
span,
174+
{
175+
ResourceAttributes.FAAS_ID: MOCK_LAMBDA_CONTEXT.invoked_function_arn,
176+
SpanAttributes.FAAS_EXECUTION: MOCK_LAMBDA_CONTEXT.aws_request_id,
177+
},
178+
)
179+
self.memory_exporter.clear()
180+
AwsLambdaInstrumentor().uninstrument()
181+
test_env_patch.stop()
156182

157183
def test_parent_context_from_lambda_event(self):
158184
@dataclass
@@ -218,14 +244,10 @@ def custom_event_context_extractor(lambda_event):
218244
assert spans
219245
self.assertEqual(len(spans), 1)
220246
span = spans[0]
221-
self.assertEqual(
222-
span.get_span_context().trace_id, test.expected_traceid
223-
)
247+
self.assertEqual(span.get_span_context().trace_id, test.expected_traceid)
224248

225249
parent_context = span.parent
226-
self.assertEqual(
227-
parent_context.trace_id, span.get_span_context().trace_id
228-
)
250+
self.assertEqual(parent_context.trace_id, span.get_span_context().trace_id)
229251
self.assertEqual(parent_context.span_id, test.expected_parentid)
230252
self.assertEqual(
231253
len(parent_context.trace_state), test.expected_trace_state_len
@@ -247,6 +269,7 @@ class TestCase:
247269
expected_link_trace_id: int
248270
expected_link_attributes: dict
249271
xray_traceid: str
272+
use_xray_propagator: str
250273

251274
tests = [
252275
TestCase(
@@ -255,13 +278,23 @@ class TestCase:
255278
expected_link_trace_id=MOCK_XRAY_TRACE_ID,
256279
expected_link_attributes={"source": "x-ray-env"},
257280
xray_traceid=MOCK_XRAY_TRACE_CONTEXT_SAMPLED,
281+
use_xray_propagator="false",
258282
),
259283
TestCase(
260284
name="invalid_xray_trace",
261285
context={},
262286
expected_link_trace_id=None,
263287
expected_link_attributes={},
264288
xray_traceid="0",
289+
use_xray_propagator="false",
290+
),
291+
TestCase(
292+
name="use_xra",
293+
context={},
294+
expected_link_trace_id=None,
295+
expected_link_attributes={},
296+
xray_traceid=MOCK_XRAY_TRACE_CONTEXT_SAMPLED,
297+
use_xray_propagator="true",
265298
),
266299
]
267300
for test in tests:
@@ -273,6 +306,7 @@ class TestCase:
273306
_X_AMZN_TRACE_ID: test.xray_traceid,
274307
# NOT using the X-Ray Propagator
275308
OTEL_PROPAGATORS: "tracecontext",
309+
OTEL_LAMBDA_USE_AWS_CONTEXT_PROPAGATION: test.use_xray_propagator,
276310
},
277311
)
278312
test_env_patch.start()
@@ -287,12 +321,8 @@ class TestCase:
287321
self.assertEqual(0, len(span.links))
288322
else:
289323
link = span.links[0]
290-
self.assertEqual(
291-
link.context.trace_id, test.expected_link_trace_id
292-
)
293-
self.assertEqual(
294-
link.attributes, test.expected_link_attributes
295-
)
324+
self.assertEqual(link.context.trace_id, test.expected_link_trace_id)
325+
self.assertEqual(link.attributes, test.expected_link_attributes)
296326

297327
self.memory_exporter.clear()
298328
AwsLambdaInstrumentor().uninstrument()

0 commit comments

Comments
 (0)