diff --git a/lambda-layer/src/Dockerfile b/lambda-layer/src/Dockerfile index 49f1a957b..c403ec403 100644 --- a/lambda-layer/src/Dockerfile +++ b/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 && \ diff --git a/lambda-layer/src/otel-instrument b/lambda-layer/src/otel-instrument index 3e51ff7b1..37d5008c0 100755 --- a/lambda-layer/src/otel-instrument +++ b/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 diff --git a/lambda-layer/src/otel_wrapper.py b/lambda-layer/src/otel_wrapper.py index 295410406..8ea565e49 100644 --- a/lambda-layer/src/otel_wrapper.py +++ b/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") diff --git a/lambda-layer/src/tests/test_lambda_instrumentation.py b/lambda-layer/src/tests/test_lambda_instrumentation.py index 0b44d1deb..1b8676b04 100644 --- a/lambda-layer/src/tests/test_lambda_instrumentation.py +++ b/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", },