Respect trace context from Lambda environment variable _X_AMZ_TRACE_ID

lambda-layer/src/Dockerfile

@@ -6,6 +6,10 @@ ADD . /workspace
 
 WORKDIR /workspace
 
+RUN sed -i "/opentelemetry-exporter-otlp-proto-grpc/d" ./aws-opentelemetry-distro/pyproject.toml
+
+RUN sed -i "/opentelemetry-instrumentation-system-metrics/d" ./aws-opentelemetry-distro/pyproject.toml
+
 RUN mkdir -p /build && \
   python3 -m pip install aws-opentelemetry-distro/ -t /build/python && \
   mv otel_wrapper.py /build/python && \

lambda-layer/src/otel-instrument

@@ -126,7 +126,7 @@ fi
 # - Enable botocore instrumentation by default
 
 if [ -z ${OTEL_PYTHON_DISABLED_INSTRUMENTATIONS} ]; then
-    export OTEL_PYTHON_DISABLED_INSTRUMENTATIONS="aio-pika,aiohttp-client,aiopg,asgi,asyncpg,boto3sqs,boto,cassandra,celery,confluent-kafka,dbapi,django,elasticsearch,fastapi,falcon,flask,grpc,httpx,jinja2,kafka-python,logging,mysql,mysqlclient,pika,psycopg2,pymemcache,pymongo,pymysql,pyramid,redis,remoulade,requests,sklearn,sqlalchemy,sqlite3,starlette,system-metrics,tornado,tortoiseorm,urllib,urllib3,wsgi"
+    export OTEL_PYTHON_DISABLED_INSTRUMENTATIONS="aio-pika,aiohttp-client,aiohttp-server,aiopg,asgi,asyncio,asyncpg,boto,boto3,cassandra,celery,confluent_kafka,dbapi,django,elasticsearch,falcon,fastapi,flask,grpc_client,grpc_server,grpc_aio_client,grpc_aio_server,httpx,jinja2,kafka,logging,mysql,mysqlclient,pika,psycopg,psycopg2,pymemcache,pymongo,pymysql,pyramid,redis,remoulade,requests,sklearn,sqlalchemy,sqlite3,starlette,system_metrics,threading,tornado,tortoiseorm,urllib,urllib3,wsgi"
 fi
 
 # - Use a wrapper because AWS Lambda's `python3 /var/runtime/bootstrap.py` will

lambda-layer/src/otel_wrapper.py

@@ -36,8 +36,14 @@
 
 import os
 from importlib import import_module
+from typing import Any
 
-from opentelemetry.instrumentation.aws_lambda import AwsLambdaInstrumentor
+from opentelemetry.context import Context
+from opentelemetry.instrumentation.aws_lambda import _X_AMZN_TRACE_ID, AwsLambdaInstrumentor
+from opentelemetry.propagate import get_global_textmap
+from opentelemetry.propagators.aws import AwsXRayPropagator
+from opentelemetry.propagators.aws.aws_xray_propagator import TRACE_HEADER_KEY
+from opentelemetry.trace import get_current_span
 
 
 def modify_module_name(module_name):
@@ -49,7 +55,26 @@ class HandlerError(Exception):
     pass
 
 
-AwsLambdaInstrumentor().instrument()
+def custom_event_context_extractor(lambda_event: Any) -> Context:
+    xray_env_var = os.environ.get(_X_AMZN_TRACE_ID)
+    lambda_trace_context = AwsXRayPropagator().extract({TRACE_HEADER_KEY: xray_env_var})
+    parent_span_context = get_current_span(lambda_trace_context).get_span_context()
+
+    if parent_span_context is None or not parent_span_context.is_valid:
+        headers = None
+        try:
+            headers = lambda_event["headers"]
+        except (TypeError, KeyError):
+            pass
+        if not isinstance(headers, dict):
+            headers = {}
+
+        return get_global_textmap().extract(headers)
+
+    return lambda_trace_context
+
+
+AwsLambdaInstrumentor().instrument(event_context_extractor=custom_event_context_extractor)
 
 path = os.environ.get("ORIG_HANDLER")
 

lambda-layer/src/tests/test_lambda_instrumentation.py

@@ -53,12 +53,17 @@ def __init__(self, aws_request_id, invoked_function_arn):
 MOCK_XRAY_TRACE_ID = 0x5FB7331105E8BB83207FA31D4D9CDB4C
 MOCK_XRAY_TRACE_ID_STR = f"{MOCK_XRAY_TRACE_ID:x}"
 MOCK_XRAY_PARENT_SPAN_ID = 0x3328B8445A6DBAD2
+MOCK_XRAY_LAMBDA_LINEAGE = "Lineage=01cfa446:0"
 MOCK_XRAY_TRACE_CONTEXT_COMMON = (
     f"Root={TRACE_ID_VERSION}-{MOCK_XRAY_TRACE_ID_STR[:TRACE_ID_FIRST_PART_LENGTH]}"
     f"-{MOCK_XRAY_TRACE_ID_STR[TRACE_ID_FIRST_PART_LENGTH:]};Parent={MOCK_XRAY_PARENT_SPAN_ID:x}"
 )
-MOCK_XRAY_TRACE_CONTEXT_SAMPLED = f"{MOCK_XRAY_TRACE_CONTEXT_COMMON};Sampled=1"
-MOCK_XRAY_TRACE_CONTEXT_NOT_SAMPLED = f"{MOCK_XRAY_TRACE_CONTEXT_COMMON};Sampled=0"
+MOCK_XRAY_TRACE_CONTEXT_SAMPLED = f"{MOCK_XRAY_TRACE_CONTEXT_COMMON};Sampled=1;{MOCK_XRAY_LAMBDA_LINEAGE}"
+MOCK_XRAY_TRACE_CONTEXT_NOT_SAMPLED = f"{MOCK_XRAY_TRACE_CONTEXT_COMMON};Sampled=0;{MOCK_XRAY_LAMBDA_LINEAGE}"
+MOCK_XRAY_TRACE_CONTEXT_PASSTHROUGH = (
+    f"Root={TRACE_ID_VERSION}-{MOCK_XRAY_TRACE_ID_STR[:TRACE_ID_FIRST_PART_LENGTH]}"
+    f"-{MOCK_XRAY_TRACE_ID_STR[TRACE_ID_FIRST_PART_LENGTH:]};{MOCK_XRAY_LAMBDA_LINEAGE}"
+)
 
 # See more:
 # https://www.w3.org/TR/trace-context/#examples-of-http-traceparent-headers
@@ -232,7 +237,7 @@ def test_parent_context_from_lambda_event(self):
             {
                 **os.environ,
                 # NOT Active Tracing
-                _X_AMZN_TRACE_ID: MOCK_XRAY_TRACE_CONTEXT_NOT_SAMPLED,
+                _X_AMZN_TRACE_ID: MOCK_XRAY_TRACE_CONTEXT_PASSTHROUGH,
                 # NOT using the X-Ray Propagator
                 OTEL_PROPAGATORS: "tracecontext",
             },