Skip to content

Commit

Permalink
update awslambda to use _X_AMZN_TRACE_ID as a Span Link (#1657)
Browse files Browse the repository at this point in the history 
Co-authored-by: Ron Yishai <ronyish@cisco.com>
Co-authored-by: Srikanth Chekuri <srikanth.chekuri92@gmail.com>
Co-authored-by: Shalev Roda <65566801+shalevr@users.noreply.github.com>
Co-authored-by: Diego Hurtado <ocelotl@users.noreply.github.com>
  • Loading branch information
@tsloughter @ronyis
@srikanthccv @shalevr @ocelotl
5 people authored Jul 18, 2023
1 parent 74fcbf4 commit 7603a1f
Show file tree
Hide file tree
Showing 4 changed files with 104 additions and 114 deletions.
3 changes: 3 additions & 0 deletions CHANGELOG.md
View file
Original file line number Diff line number Diff line change
Expand Up @@ -150,6 +150,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
([#1592](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/1592))
- `opentelemetry-instrumentation-django` Allow explicit `excluded_urls` configuration through `instrument()`
([#1618](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/1618))
- `opentelemetry-instrumentation-aws-lambda` Use env var `_X_AMZN_TRACE_ID` as a
Span Link instead of parent
([#1657](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/1657))

### Fixed

Expand Down
80 changes: 35 additions & 45 deletions ...metry-instrumentation-aws-lambda/src/opentelemetry/instrumentation/aws_lambda/__init__.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -71,7 +71,7 @@ def custom_event_context_extractor(lambda_event):
import os
import time
from importlib import import_module
from typing import Any, Callable, Collection
from typing import Any, Callable, Collection, Optional, Sequence
from urllib.parse import urlencode

from wrapt import wrap_function_wrapper
Expand All @@ -90,6 +90,7 @@ def custom_event_context_extractor(lambda_event):
from opentelemetry.semconv.resource import ResourceAttributes
from opentelemetry.semconv.trace import SpanAttributes
from opentelemetry.trace import (
Link,
Span,
SpanKind,
TracerProvider,
Expand All @@ -106,9 +107,6 @@ def custom_event_context_extractor(lambda_event):
OTEL_INSTRUMENTATION_AWS_LAMBDA_FLUSH_TIMEOUT = (
"OTEL_INSTRUMENTATION_AWS_LAMBDA_FLUSH_TIMEOUT"
)
OTEL_LAMBDA_DISABLE_AWS_CONTEXT_PROPAGATION = (
"OTEL_LAMBDA_DISABLE_AWS_CONTEXT_PROPAGATION"
)


def _default_event_context_extractor(lambda_event: Any) -> Context:
Expand Down Expand Up @@ -142,14 +140,12 @@ def _default_event_context_extractor(lambda_event: Any) -> Context:


def _determine_parent_context(
lambda_event: Any,
event_context_extractor: Callable[[Any], Context],
disable_aws_context_propagation: bool = False,
lambda_event: Any, event_context_extractor: Callable[[Any], Context]
) -> Context:
"""Determine the parent context for the current Lambda invocation.
See more:
https://github.com/open-telemetry/opentelemetry-specification/blob/main/specification/trace/semantic_conventions/instrumentation/aws-lambda.md#determining-the-parent-of-a-span
https://github.com/open-telemetry/opentelemetry-specification/blob/main/specification/trace/semantic_conventions/instrumentation/aws-lambda.md
Args:
lambda_event: user-defined, so it could be anything, but this
Expand All @@ -158,30 +154,11 @@ def _determine_parent_context(
Event as input and extracts an OTel Context from it. By default,
the context is extracted from the HTTP headers of an API Gateway
request.
disable_aws_context_propagation: By default, this instrumentation
will try to read the context from the `_X_AMZN_TRACE_ID` environment
variable set by Lambda, set this to `True` to disable this behavior.
Returns:
A Context with configuration found in the carrier.
"""
parent_context = None

if not disable_aws_context_propagation:
xray_env_var = os.environ.get(_X_AMZN_TRACE_ID)

if xray_env_var:
parent_context = AwsXRayPropagator().extract(
{TRACE_HEADER_KEY: xray_env_var}
)

if (
parent_context
and get_current_span(parent_context)
.get_span_context()
.trace_flags.sampled
):
return parent_context

if event_context_extractor:
parent_context = event_context_extractor(lambda_event)
else:
Expand All @@ -190,6 +167,33 @@ def _determine_parent_context(
return parent_context


def _determine_links() -> Optional[Sequence[Link]]:
"""Determine if a Link should be added to the Span based on the
environment variable `_X_AMZN_TRACE_ID`.
See more:
https://github.com/open-telemetry/opentelemetry-specification/blob/main/specification/trace/semantic_conventions/instrumentation/aws-lambda.md#aws-x-ray-environment-span-link
Returns:
A Link or None
"""
links = None

xray_env_var = os.environ.get(_X_AMZN_TRACE_ID)

if xray_env_var:
env_context = AwsXRayPropagator().extract(
{TRACE_HEADER_KEY: xray_env_var}
)

span_context = get_current_span(env_context).get_span_context()
if span_context.is_valid:
links = [Link(span_context, {"source": "x-ray-env"})]

return links


def _set_api_gateway_v1_proxy_attributes(
lambda_event: Any, span: Span
) -> Span:
Expand Down Expand Up @@ -284,7 +288,6 @@ def _instrument(
flush_timeout,
event_context_extractor: Callable[[Any], Context],
tracer_provider: TracerProvider = None,
disable_aws_context_propagation: bool = False,
meter_provider: MeterProvider = None,
):
def _instrumented_lambda_handler_call( # noqa pylint: disable=too-many-branches
Expand All @@ -297,11 +300,11 @@ def _instrumented_lambda_handler_call( # noqa pylint: disable=too-many-branches
lambda_event = args[0]

parent_context = _determine_parent_context(
lambda_event,
event_context_extractor,
disable_aws_context_propagation,
lambda_event, event_context_extractor
)

links = _determine_links()

span_kind = None
try:
if lambda_event["Records"][0]["eventSource"] in {
Expand All @@ -327,6 +330,7 @@ def _instrumented_lambda_handler_call( # noqa pylint: disable=too-many-branches
name=orig_handler_name,
context=parent_context,
kind=span_kind,
links=links,
) as span:
if span.is_recording():
lambda_context = args[1]
Expand Down Expand Up @@ -420,9 +424,6 @@ def _instrument(self, **kwargs):
Event as input and extracts an OTel Context from it. By default,
the context is extracted from the HTTP headers of an API Gateway
request.
``disable_aws_context_propagation``: By default, this instrumentation
will try to read the context from the `_X_AMZN_TRACE_ID` environment
variable set by Lambda, set this to `True` to disable this behavior.
"""
lambda_handler = os.environ.get(ORIG_HANDLER, os.environ.get(_HANDLER))
# pylint: disable=attribute-defined-outside-init
Expand All @@ -444,16 +445,6 @@ def _instrument(self, **kwargs):
flush_timeout_env,
)

disable_aws_context_propagation = kwargs.get(
"disable_aws_context_propagation", False
) or os.getenv(
OTEL_LAMBDA_DISABLE_AWS_CONTEXT_PROPAGATION, "False"
).strip().lower() in (
"true",
"1",
"t",
)

_instrument(
self._wrapped_module_name,
self._wrapped_function_name,
Expand All @@ -462,7 +453,6 @@ def _instrument(self, **kwargs):
"event_context_extractor", _default_event_context_extractor
),
tracer_provider=kwargs.get("tracer_provider"),
disable_aws_context_propagation=disable_aws_context_propagation,
meter_provider=kwargs.get("meter_provider"),
)

Expand Down
133 changes: 65 additions & 68 deletions .../opentelemetry-instrumentation-aws-lambda/tests/test_aws_lambda_instrumentation_manual.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -27,7 +27,6 @@
_HANDLER,
_X_AMZN_TRACE_ID,
OTEL_INSTRUMENTATION_AWS_LAMBDA_FLUSH_TIMEOUT,
OTEL_LAMBDA_DISABLE_AWS_CONTEXT_PROPAGATION,
AwsLambdaInstrumentor,
)
from opentelemetry.propagate import get_global_textmap
Expand Down Expand Up @@ -138,7 +137,9 @@ def test_active_tracing(self):
self.assertEqual(len(spans), 1)
span = spans[0]
self.assertEqual(span.name, os.environ[_HANDLER])
self.assertEqual(span.get_span_context().trace_id, MOCK_XRAY_TRACE_ID)
self.assertNotEqual(
span.get_span_context().trace_id, MOCK_XRAY_TRACE_ID
)
self.assertEqual(span.kind, SpanKind.SERVER)
self.assertSpanHasAttributes(
span,
Expand All @@ -149,11 +150,7 @@ def test_active_tracing(self):
)

parent_context = span.parent
self.assertEqual(
parent_context.trace_id, span.get_span_context().trace_id
)
self.assertEqual(parent_context.span_id, MOCK_XRAY_PARENT_SPAN_ID)
self.assertTrue(parent_context.is_remote)
self.assertEqual(None, parent_context)

test_env_patch.stop()

Expand All @@ -165,11 +162,8 @@ class TestCase:
context: Dict
expected_traceid: int
expected_parentid: int
xray_traceid: str
expected_state_value: str = None
expected_trace_state_len: int = 0
disable_aws_context_propagation: bool = False
disable_aws_context_propagation_envvar: str = ""

def custom_event_context_extractor(lambda_event):
return get_global_textmap().extract(lambda_event["foo"]["headers"])
Expand All @@ -188,42 +182,9 @@ def custom_event_context_extractor(lambda_event):
expected_parentid=MOCK_W3C_PARENT_SPAN_ID,
expected_trace_state_len=3,
expected_state_value=MOCK_W3C_TRACE_STATE_VALUE,
xray_traceid=MOCK_XRAY_TRACE_CONTEXT_NOT_SAMPLED,
),
TestCase(
name="custom_extractor_not_sampled_xray",
custom_extractor=custom_event_context_extractor,
context={
"foo": {
"headers": {
TraceContextTextMapPropagator._TRACEPARENT_HEADER_NAME: MOCK_W3C_TRACE_CONTEXT_SAMPLED,
TraceContextTextMapPropagator._TRACESTATE_HEADER_NAME: f"{MOCK_W3C_TRACE_STATE_KEY}={MOCK_W3C_TRACE_STATE_VALUE},foo=1,bar=2",
}
}
},
expected_traceid=MOCK_W3C_TRACE_ID,
expected_parentid=MOCK_W3C_PARENT_SPAN_ID,
expected_trace_state_len=3,
expected_state_value=MOCK_W3C_TRACE_STATE_VALUE,
xray_traceid=MOCK_XRAY_TRACE_CONTEXT_NOT_SAMPLED,
),
TestCase(
name="custom_extractor_sampled_xray",
custom_extractor=custom_event_context_extractor,
context={
"foo": {
"headers": {
TraceContextTextMapPropagator._TRACEPARENT_HEADER_NAME: MOCK_W3C_TRACE_CONTEXT_SAMPLED,
TraceContextTextMapPropagator._TRACESTATE_HEADER_NAME: f"{MOCK_W3C_TRACE_STATE_KEY}={MOCK_W3C_TRACE_STATE_VALUE},foo=1,bar=2",
}
}
},
expected_traceid=MOCK_XRAY_TRACE_ID,
expected_parentid=MOCK_XRAY_PARENT_SPAN_ID,
xray_traceid=MOCK_XRAY_TRACE_CONTEXT_SAMPLED,
),
TestCase(
name="custom_extractor_sampled_xray_disable_aws_propagation",
name="custom_extractor",
custom_extractor=custom_event_context_extractor,
context={
"foo": {
Expand All @@ -233,47 +194,24 @@ def custom_event_context_extractor(lambda_event):
}
}
},
disable_aws_context_propagation=True,
expected_traceid=MOCK_W3C_TRACE_ID,
expected_parentid=MOCK_W3C_PARENT_SPAN_ID,
expected_trace_state_len=3,
expected_state_value=MOCK_W3C_TRACE_STATE_VALUE,
xray_traceid=MOCK_XRAY_TRACE_CONTEXT_SAMPLED,
),
TestCase(
name="no_custom_extractor_xray_disable_aws_propagation_via_env_var",
custom_extractor=None,
context={
"headers": {
TraceContextTextMapPropagator._TRACEPARENT_HEADER_NAME: MOCK_W3C_TRACE_CONTEXT_SAMPLED,
TraceContextTextMapPropagator._TRACESTATE_HEADER_NAME: f"{MOCK_W3C_TRACE_STATE_KEY}={MOCK_W3C_TRACE_STATE_VALUE},foo=1,bar=2",
}
},
disable_aws_context_propagation=False,
disable_aws_context_propagation_envvar="true",
expected_traceid=MOCK_W3C_TRACE_ID,
expected_parentid=MOCK_W3C_PARENT_SPAN_ID,
expected_trace_state_len=3,
expected_state_value=MOCK_W3C_TRACE_STATE_VALUE,
xray_traceid=MOCK_XRAY_TRACE_CONTEXT_SAMPLED,
),
]
for test in tests:
test_env_patch = mock.patch.dict(
"os.environ",
{
**os.environ,
# NOT Active Tracing
_X_AMZN_TRACE_ID: test.xray_traceid,
OTEL_LAMBDA_DISABLE_AWS_CONTEXT_PROPAGATION: test.disable_aws_context_propagation_envvar,
# NOT using the X-Ray Propagator
OTEL_PROPAGATORS: "tracecontext",
},
)
test_env_patch.start()
AwsLambdaInstrumentor().instrument(
event_context_extractor=test.custom_extractor,
disable_aws_context_propagation=test.disable_aws_context_propagation,
event_context_extractor=test.custom_extractor
)
mock_execute_lambda(test.context)
spans = self.memory_exporter.get_finished_spans()
Expand Down Expand Up @@ -301,6 +239,65 @@ def custom_event_context_extractor(lambda_event):
AwsLambdaInstrumentor().uninstrument()
test_env_patch.stop()

def test_links_from_lambda_event(self):
@dataclass
class TestCase:
name: str
context: Dict
expected_link_trace_id: int
expected_link_attributes: dict
xray_traceid: str

tests = [
TestCase(
name="valid_xray_trace",
context={},
expected_link_trace_id=MOCK_XRAY_TRACE_ID,
expected_link_attributes={"source": "x-ray-env"},
xray_traceid=MOCK_XRAY_TRACE_CONTEXT_SAMPLED,
),
TestCase(
name="invalid_xray_trace",
context={},
expected_link_trace_id=None,
expected_link_attributes={},
xray_traceid="0",
),
]
for test in tests:
test_env_patch = mock.patch.dict(
"os.environ",
{
**os.environ,
# NOT Active Tracing
_X_AMZN_TRACE_ID: test.xray_traceid,
# NOT using the X-Ray Propagator
OTEL_PROPAGATORS: "tracecontext",
},
)
test_env_patch.start()
AwsLambdaInstrumentor().instrument()
mock_execute_lambda(test.context)
spans = self.memory_exporter.get_finished_spans()
assert spans
self.assertEqual(len(spans), 1)
span = spans[0]

if test.expected_link_trace_id is None:
self.assertEqual(0, len(span.links))
else:
link = span.links[0]
self.assertEqual(
link.context.trace_id, test.expected_link_trace_id
)
self.assertEqual(
link.attributes, test.expected_link_attributes
)

self.memory_exporter.clear()
AwsLambdaInstrumentor().uninstrument()
test_env_patch.stop()

def test_lambda_no_error_with_invalid_flush_timeout(self):
test_env_patch = mock.patch.dict(
"os.environ",
Expand Down
2 changes: 1 addition & 1 deletion tox.ini
View file
Original file line number Diff line number Diff line change
Expand Up @@ -34,7 +34,7 @@ envlist =
; instrumentation-aiopg intentionally excluded from pypy3

; opentelemetry-instrumentation-aws-lambda
py3{7,8,9}-test-instrumentation-aws-lambda
py3{7,8,9,10,11}-test-instrumentation-aws-lambda

; opentelemetry-instrumentation-botocore
py3{7,8,9,10,11}-test-instrumentation-botocore
Expand Down

0 comments on commit 7603a1f

Please sign in to comment.