Skip to content

Commit 8dbd142

Browse files
author
Alex Boten
authored
opentelemetry-instrumentation-aws-lambda: Adding option to disable context propagation (#1466)
* `opentelemetry-instrumentation-aws-lambda`: Adding option to disable context propagation Adding the following option to disable context propagation `disable_aws_context_propagation`. This is similar to the disableAwsContextPropagation option in the nodejs instrumentation. * update changelog * lint * more lint
1 parent 80d0b89 commit 8dbd142

File tree

3 files changed

+148
-98
lines changed

3 files changed

+148
-98
lines changed

‎CHANGELOG.md

+2
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
1111
([#685](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/685))
1212
- Add metric instrumentation for tornado
1313
([#1252](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/1252))
14+
- `opentelemetry-instrumentation-aws-lambda` Add option to disable aws context propagation
15+
([#1466](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/1466))
1416

1517
### Added
1618

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

+29-9
Original file line numberDiff line numberDiff line change
@@ -134,7 +134,9 @@ def _default_event_context_extractor(lambda_event: Any) -> Context:
134134

135135

136136
def _determine_parent_context(
137-
lambda_event: Any, event_context_extractor: Callable[[Any], Context]
137+
lambda_event: Any,
138+
event_context_extractor: Callable[[Any], Context],
139+
disable_aws_context_propagation: bool = False,
138140
) -> Context:
139141
"""Determine the parent context for the current Lambda invocation.
140142
@@ -144,17 +146,25 @@ def _determine_parent_context(
144146
Args:
145147
lambda_event: user-defined, so it could be anything, but this
146148
method counts it being a map with a 'headers' key
149+
event_context_extractor: a method which takes the Lambda
150+
Event as input and extracts an OTel Context from it. By default,
151+
the context is extracted from the HTTP headers of an API Gateway
152+
request.
153+
disable_aws_context_propagation: By default, this instrumentation
154+
will try to read the context from the `_X_AMZN_TRACE_ID` environment
155+
variable set by Lambda, set this to `True` to disable this behavior.
147156
Returns:
148157
A Context with configuration found in the carrier.
149158
"""
150159
parent_context = None
151160

152-
xray_env_var = os.environ.get(_X_AMZN_TRACE_ID)
161+
if not disable_aws_context_propagation:
162+
xray_env_var = os.environ.get(_X_AMZN_TRACE_ID)
153163

154-
if xray_env_var:
155-
parent_context = AwsXRayPropagator().extract(
156-
{TRACE_HEADER_KEY: xray_env_var}
157-
)
164+
if xray_env_var:
165+
parent_context = AwsXRayPropagator().extract(
166+
{TRACE_HEADER_KEY: xray_env_var}
167+
)
158168

159169
if (
160170
parent_context
@@ -258,6 +268,7 @@ def _instrument(
258268
flush_timeout,
259269
event_context_extractor: Callable[[Any], Context],
260270
tracer_provider: TracerProvider = None,
271+
disable_aws_context_propagation: bool = False,
261272
):
262273
def _instrumented_lambda_handler_call(
263274
call_wrapped, instance, args, kwargs
@@ -269,7 +280,9 @@ def _instrumented_lambda_handler_call(
269280
lambda_event = args[0]
270281

271282
parent_context = _determine_parent_context(
272-
lambda_event, event_context_extractor
283+
lambda_event,
284+
event_context_extractor,
285+
disable_aws_context_propagation,
273286
)
274287

275288
span_kind = None
@@ -368,6 +381,9 @@ def _instrument(self, **kwargs):
368381
Event as input and extracts an OTel Context from it. By default,
369382
the context is extracted from the HTTP headers of an API Gateway
370383
request.
384+
``disable_aws_context_propagation``: By default, this instrumentation
385+
will try to read the context from the `_X_AMZN_TRACE_ID` environment
386+
variable set by Lambda, set this to `True` to disable this behavior.
371387
"""
372388
lambda_handler = os.environ.get(ORIG_HANDLER, os.environ.get(_HANDLER))
373389
# pylint: disable=attribute-defined-outside-init
@@ -377,11 +393,12 @@ def _instrument(self, **kwargs):
377393
) = lambda_handler.rsplit(".", 1)
378394

379395
flush_timeout_env = os.environ.get(
380-
OTEL_INSTRUMENTATION_AWS_LAMBDA_FLUSH_TIMEOUT, ""
396+
OTEL_INSTRUMENTATION_AWS_LAMBDA_FLUSH_TIMEOUT, None
381397
)
382398
flush_timeout = 30000
383399
try:
384-
flush_timeout = int(flush_timeout_env)
400+
if flush_timeout_env is not None:
401+
flush_timeout = int(flush_timeout_env)
385402
except ValueError:
386403
logger.warning(
387404
"Could not convert OTEL_INSTRUMENTATION_AWS_LAMBDA_FLUSH_TIMEOUT value %s to int",
@@ -396,6 +413,9 @@ def _instrument(self, **kwargs):
396413
"event_context_extractor", _default_event_context_extractor
397414
),
398415
tracer_provider=kwargs.get("tracer_provider"),
416+
disable_aws_context_propagation=kwargs.get(
417+
"disable_aws_context_propagation", False
418+
),
399419
)
400420

401421
def _uninstrument(self, **kwargs):

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

+117-89
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,9 @@
1212
# See the License for the specific language governing permissions and
1313
# limitations under the License.
1414
import os
15+
from dataclasses import dataclass
1516
from importlib import import_module
17+
from typing import Any, Callable, Dict
1618
from unittest import mock
1719

1820
from mocks.api_gateway_http_api_event import (
@@ -155,103 +157,129 @@ def test_active_tracing(self):
155157
test_env_patch.stop()
156158

157159
def test_parent_context_from_lambda_event(self):
158-
test_env_patch = mock.patch.dict(
159-
"os.environ",
160-
{
161-
**os.environ,
162-
# NOT Active Tracing
163-
_X_AMZN_TRACE_ID: MOCK_XRAY_TRACE_CONTEXT_NOT_SAMPLED,
164-
# NOT using the X-Ray Propagator
165-
OTEL_PROPAGATORS: "tracecontext",
166-
},
167-
)
168-
test_env_patch.start()
169-
170-
AwsLambdaInstrumentor().instrument()
171-
172-
mock_execute_lambda(
173-
{
174-
"headers": {
175-
TraceContextTextMapPropagator._TRACEPARENT_HEADER_NAME: MOCK_W3C_TRACE_CONTEXT_SAMPLED,
176-
TraceContextTextMapPropagator._TRACESTATE_HEADER_NAME: f"{MOCK_W3C_TRACE_STATE_KEY}={MOCK_W3C_TRACE_STATE_VALUE},foo=1,bar=2",
177-
}
178-
}
179-
)
180-
181-
spans = self.memory_exporter.get_finished_spans()
160+
@dataclass
161+
class TestCase:
162+
name: str
163+
custom_extractor: Callable[[Any], None]
164+
context: Dict
165+
expected_traceid: int
166+
expected_parentid: int
167+
xray_traceid: str
168+
expected_state_value: str = None
169+
expected_trace_state_len: int = 0
170+
disable_aws_context_propagation: bool = False
182171

183-
assert spans
184-
185-
self.assertEqual(len(spans), 1)
186-
span = spans[0]
187-
self.assertEqual(span.get_span_context().trace_id, MOCK_W3C_TRACE_ID)
188-
189-
parent_context = span.parent
190-
self.assertEqual(
191-
parent_context.trace_id, span.get_span_context().trace_id
192-
)
193-
self.assertEqual(parent_context.span_id, MOCK_W3C_PARENT_SPAN_ID)
194-
self.assertEqual(len(parent_context.trace_state), 3)
195-
self.assertEqual(
196-
parent_context.trace_state.get(MOCK_W3C_TRACE_STATE_KEY),
197-
MOCK_W3C_TRACE_STATE_VALUE,
198-
)
199-
self.assertTrue(parent_context.is_remote)
200-
201-
test_env_patch.stop()
202-
203-
def test_using_custom_extractor(self):
204172
def custom_event_context_extractor(lambda_event):
205173
return get_global_textmap().extract(lambda_event["foo"]["headers"])
206174

207-
test_env_patch = mock.patch.dict(
208-
"os.environ",
209-
{
210-
**os.environ,
211-
# NOT Active Tracing
212-
_X_AMZN_TRACE_ID: MOCK_XRAY_TRACE_CONTEXT_NOT_SAMPLED,
213-
# NOT using the X-Ray Propagator
214-
OTEL_PROPAGATORS: "tracecontext",
215-
},
216-
)
217-
test_env_patch.start()
218-
219-
AwsLambdaInstrumentor().instrument(
220-
event_context_extractor=custom_event_context_extractor,
221-
)
222-
223-
mock_execute_lambda(
224-
{
225-
"foo": {
175+
tests = [
176+
TestCase(
177+
name="no_custom_extractor",
178+
custom_extractor=None,
179+
context={
226180
"headers": {
227181
TraceContextTextMapPropagator._TRACEPARENT_HEADER_NAME: MOCK_W3C_TRACE_CONTEXT_SAMPLED,
228182
TraceContextTextMapPropagator._TRACESTATE_HEADER_NAME: f"{MOCK_W3C_TRACE_STATE_KEY}={MOCK_W3C_TRACE_STATE_VALUE},foo=1,bar=2",
229183
}
230-
}
231-
}
232-
)
233-
234-
spans = self.memory_exporter.get_finished_spans()
235-
236-
assert spans
237-
238-
self.assertEqual(len(spans), 1)
239-
span = spans[0]
240-
self.assertEqual(span.get_span_context().trace_id, MOCK_W3C_TRACE_ID)
241-
242-
parent_context = span.parent
243-
self.assertEqual(
244-
parent_context.trace_id, span.get_span_context().trace_id
245-
)
246-
self.assertEqual(parent_context.span_id, MOCK_W3C_PARENT_SPAN_ID)
247-
self.assertEqual(len(parent_context.trace_state), 3)
248-
self.assertEqual(
249-
parent_context.trace_state.get(MOCK_W3C_TRACE_STATE_KEY),
250-
MOCK_W3C_TRACE_STATE_VALUE,
251-
)
252-
self.assertTrue(parent_context.is_remote)
253-
254-
test_env_patch.stop()
184+
},
185+
expected_traceid=MOCK_W3C_TRACE_ID,
186+
expected_parentid=MOCK_W3C_PARENT_SPAN_ID,
187+
expected_trace_state_len=3,
188+
expected_state_value=MOCK_W3C_TRACE_STATE_VALUE,
189+
xray_traceid=MOCK_XRAY_TRACE_CONTEXT_NOT_SAMPLED,
190+
),
191+
TestCase(
192+
name="custom_extractor_not_sampled_xray",
193+
custom_extractor=custom_event_context_extractor,
194+
context={
195+
"foo": {
196+
"headers": {
197+
TraceContextTextMapPropagator._TRACEPARENT_HEADER_NAME: MOCK_W3C_TRACE_CONTEXT_SAMPLED,
198+
TraceContextTextMapPropagator._TRACESTATE_HEADER_NAME: f"{MOCK_W3C_TRACE_STATE_KEY}={MOCK_W3C_TRACE_STATE_VALUE},foo=1,bar=2",
199+
}
200+
}
201+
},
202+
expected_traceid=MOCK_W3C_TRACE_ID,
203+
expected_parentid=MOCK_W3C_PARENT_SPAN_ID,
204+
expected_trace_state_len=3,
205+
expected_state_value=MOCK_W3C_TRACE_STATE_VALUE,
206+
xray_traceid=MOCK_XRAY_TRACE_CONTEXT_NOT_SAMPLED,
207+
),
208+
TestCase(
209+
name="custom_extractor_sampled_xray",
210+
custom_extractor=custom_event_context_extractor,
211+
context={
212+
"foo": {
213+
"headers": {
214+
TraceContextTextMapPropagator._TRACEPARENT_HEADER_NAME: MOCK_W3C_TRACE_CONTEXT_SAMPLED,
215+
TraceContextTextMapPropagator._TRACESTATE_HEADER_NAME: f"{MOCK_W3C_TRACE_STATE_KEY}={MOCK_W3C_TRACE_STATE_VALUE},foo=1,bar=2",
216+
}
217+
}
218+
},
219+
expected_traceid=MOCK_XRAY_TRACE_ID,
220+
expected_parentid=MOCK_XRAY_PARENT_SPAN_ID,
221+
xray_traceid=MOCK_XRAY_TRACE_CONTEXT_SAMPLED,
222+
),
223+
TestCase(
224+
name="custom_extractor_sampled_xray_disable_aws_propagation",
225+
custom_extractor=custom_event_context_extractor,
226+
context={
227+
"foo": {
228+
"headers": {
229+
TraceContextTextMapPropagator._TRACEPARENT_HEADER_NAME: MOCK_W3C_TRACE_CONTEXT_SAMPLED,
230+
TraceContextTextMapPropagator._TRACESTATE_HEADER_NAME: f"{MOCK_W3C_TRACE_STATE_KEY}={MOCK_W3C_TRACE_STATE_VALUE},foo=1,bar=2",
231+
}
232+
}
233+
},
234+
disable_aws_context_propagation=True,
235+
expected_traceid=MOCK_W3C_TRACE_ID,
236+
expected_parentid=MOCK_W3C_PARENT_SPAN_ID,
237+
expected_trace_state_len=3,
238+
expected_state_value=MOCK_W3C_TRACE_STATE_VALUE,
239+
xray_traceid=MOCK_XRAY_TRACE_CONTEXT_SAMPLED,
240+
),
241+
]
242+
for test in tests:
243+
test_env_patch = mock.patch.dict(
244+
"os.environ",
245+
{
246+
**os.environ,
247+
# NOT Active Tracing
248+
_X_AMZN_TRACE_ID: test.xray_traceid,
249+
# NOT using the X-Ray Propagator
250+
OTEL_PROPAGATORS: "tracecontext",
251+
},
252+
)
253+
test_env_patch.start()
254+
AwsLambdaInstrumentor().instrument(
255+
event_context_extractor=test.custom_extractor,
256+
disable_aws_context_propagation=test.disable_aws_context_propagation,
257+
)
258+
mock_execute_lambda(test.context)
259+
spans = self.memory_exporter.get_finished_spans()
260+
assert spans
261+
self.assertEqual(len(spans), 1)
262+
span = spans[0]
263+
self.assertEqual(
264+
span.get_span_context().trace_id, test.expected_traceid
265+
)
266+
267+
parent_context = span.parent
268+
self.assertEqual(
269+
parent_context.trace_id, span.get_span_context().trace_id
270+
)
271+
self.assertEqual(parent_context.span_id, test.expected_parentid)
272+
self.assertEqual(
273+
len(parent_context.trace_state), test.expected_trace_state_len
274+
)
275+
self.assertEqual(
276+
parent_context.trace_state.get(MOCK_W3C_TRACE_STATE_KEY),
277+
test.expected_state_value,
278+
)
279+
self.assertTrue(parent_context.is_remote)
280+
self.memory_exporter.clear()
281+
AwsLambdaInstrumentor().uninstrument()
282+
test_env_patch.stop()
255283

256284
def test_lambda_no_error_with_invalid_flush_timeout(self):
257285

0 commit comments

Comments
 (0)