Add instrumentation for AWS Lambda Service - Implementation (Part 2/2) (#777)

* Add instrumentation for AWS Lambda Service - Implementation

* Lambda is CONSUMER SQS trace if 'Records' key in Lambda event

* More robust check of SQS by indexing and catching

* Explicitly catch errors we expect when determinig if SQS triggered Lambda

CHANGELOG.md

@@ -12,6 +12,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   ([#742](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/742))
 - `opentelemetry-instrumentation` Add `setuptools` to `install_requires`
   ([#781](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/781))
+- `opentelemetry-instrumentation-aws-lambda` Add instrumentation for AWS Lambda Service - Implementation (Part 2/2)
+  ([#777](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/777))
 
 ### Fixed
 

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

@@ -11,3 +11,266 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+
+"""
+The opentelemetry-instrumentation-aws-lambda package provides an `Instrumentor`
+to traces calls whithin a Python AWS Lambda function.
+
+Usage
+-----
+
+.. code:: python
+
+    # Copy this snippet into an AWS Lambda function
+
+    import boto3
+    from opentelemetry.instrumentation.botocore import AwsBotocoreInstrumentor
+    from opentelemetry.instrumentation.aws_lambda import AwsLambdaInstrumentor
+
+
+    # Enable instrumentation
+    AwsBotocoreInstrumentor().instrument()
+    AwsLambdaInstrumentor().instrument()
+
+    # Lambda function
+    def lambda_handler(event, context):
+        s3 = boto3.resource('s3')
+        for bucket in s3.buckets.all():
+            print(bucket.name)
+
+        return "200 OK"
+
+API
+---
+
+The `instrument` method accepts the following keyword args:
+
+tracer_provider (TracerProvider) - an optional tracer provider
+event_context_extractor (Callable) - a function that returns an OTel Trace
+    Context given the Lambda Event the AWS Lambda was invoked with
+    this function signature is: def event_context_extractor(lambda_event: Any) -> Context
+for example:
+
+.. code:: python
+
+    from opentelemetry.instrumentation.aws_lambda import AwsLambdaInstrumentor
+
+    def custom_event_context_extractor(lambda_event):
+        # If the `TraceContextTextMapPropagator` is the global propagator, we
+        # can use it to parse out the context from the HTTP Headers.
+        return get_global_textmap().extract(lambda_event["foo"]["headers"])
+
+    AwsLambdaInstrumentor().instrument(
+        event_context_extractor=custom_event_context_extractor
+    )
+"""
+
+import logging
+import os
+from importlib import import_module
+from typing import Any, Callable, Collection
+
+from wrapt import wrap_function_wrapper
+
+from opentelemetry.context.context import Context
+from opentelemetry.instrumentation.aws_lambda.package import _instruments
+from opentelemetry.instrumentation.aws_lambda.version import __version__
+from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
+from opentelemetry.instrumentation.utils import unwrap
+from opentelemetry.propagate import get_global_textmap
+from opentelemetry.propagators.aws.aws_xray_propagator import (
+    TRACE_HEADER_KEY,
+    AwsXRayPropagator,
+)
+from opentelemetry.semconv.resource import ResourceAttributes
+from opentelemetry.semconv.trace import SpanAttributes
+from opentelemetry.trace import (
+    SpanKind,
+    TracerProvider,
+    get_tracer,
+    get_tracer_provider,
+)
+from opentelemetry.trace.propagation import get_current_span
+
+logger = logging.getLogger(__name__)
+
+_HANDLER = "_HANDLER"
+_X_AMZN_TRACE_ID = "_X_AMZN_TRACE_ID"
+ORIG_HANDLER = "ORIG_HANDLER"
+
+
+def _default_event_context_extractor(lambda_event: Any) -> Context:
+    """Default way of extracting the context from the Lambda Event.
+
+    Assumes the Lambda Event is a map with the headers under the 'headers' key.
+    This is the mapping to use when the Lambda is invoked by an API Gateway
+    REST API where API Gateway is acting as a pure proxy for the request.
+
+    See more:
+    https://docs.aws.amazon.com/apigateway/latest/developerguide/set-up-lambda-proxy-integrations.html#api-gateway-simple-proxy-for-lambda-input-format
+
+    Args:
+        lambda_event: user-defined, so it could be anything, but this
+            method counts on it being a map with a 'headers' key
+    Returns:
+        A Context with configuration found in the event.
+    """
+    try:
+        headers = lambda_event["headers"]
+    except (TypeError, KeyError):
+        logger.debug(
+            "Extracting context from Lambda Event failed: either enable X-Ray active tracing or configure API Gateway to trigger this Lambda function as a pure proxy. Otherwise, generated spans will have an invalid (empty) parent context."
+        )
+        headers = {}
+    return get_global_textmap().extract(headers)
+
+
+def _determine_parent_context(
+    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
+
+        Args:
+            lambda_event: user-defined, so it could be anything, but this
+                method counts it being a map with a 'headers' key
+        Returns:
+            A Context with configuration found in the carrier.
+        """
+    parent_context = None
+
+    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:
+        parent_context = _default_event_context_extractor(lambda_event)
+
+    return parent_context
+
+
+def _instrument(
+    wrapped_module_name,
+    wrapped_function_name,
+    event_context_extractor: Callable[[Any], Context],
+    tracer_provider: TracerProvider = None,
+):
+    def _instrumented_lambda_handler_call(
+        call_wrapped, instance, args, kwargs
+    ):
+        orig_handler_name = ".".join(
+            [wrapped_module_name, wrapped_function_name]
+        )
+
+        lambda_event = args[0]
+
+        parent_context = _determine_parent_context(
+            lambda_event, event_context_extractor
+        )
+
+        # See more:
+        # https://docs.aws.amazon.com/lambda/latest/dg/with-sqs.html
+        try:
+            if lambda_event["Records"][0]["eventSource"] == "aws:sqs":
+                span_kind = SpanKind.CONSUMER
+        except (IndexError, KeyError, TypeError):
+            span_kind = SpanKind.SERVER
+
+        tracer = get_tracer(__name__, __version__, tracer_provider)
+
+        with tracer.start_as_current_span(
+            name=orig_handler_name, context=parent_context, kind=span_kind,
+        ) as span:
+            if span.is_recording():
+                lambda_context = args[1]
+                # NOTE: The specs mention an exception here, allowing the
+                # `ResourceAttributes.FAAS_ID` attribute to be set as a span
+                # attribute instead of a resource attribute.
+                #
+                # See more:
+                # https://github.com/open-telemetry/opentelemetry-specification/blob/main/specification/trace/semantic_conventions/faas.md#example
+                span.set_attribute(
+                    ResourceAttributes.FAAS_ID,
+                    lambda_context.invoked_function_arn,
+                )
+                span.set_attribute(
+                    SpanAttributes.FAAS_EXECUTION,
+                    lambda_context.aws_request_id,
+                )
+
+            result = call_wrapped(*args, **kwargs)
+
+        _tracer_provider = tracer_provider or get_tracer_provider()
+        try:
+            # NOTE: `force_flush` before function quit in case of Lambda freeze.
+            # Assumes we are using the OpenTelemetry SDK implementation of the
+            # `TracerProvider`.
+            _tracer_provider.force_flush()
+        except Exception:  # pylint: disable=broad-except
+            logger.error(
+                "TracerProvider was missing `force_flush` method. This is necessary in case of a Lambda freeze and would exist in the OTel SDK implementation."
+            )
+
+        return result
+
+    wrap_function_wrapper(
+        wrapped_module_name,
+        wrapped_function_name,
+        _instrumented_lambda_handler_call,
+    )
+
+
+class AwsLambdaInstrumentor(BaseInstrumentor):
+    def instrumentation_dependencies(self) -> Collection[str]:
+        return _instruments
+
+    def _instrument(self, **kwargs):
+        """Instruments Lambda Handlers on AWS Lambda.
+
+        See more:
+        https://github.com/open-telemetry/opentelemetry-specification/blob/main/specification/trace/semantic_conventions/instrumentation/aws-lambda.md#instrumenting-aws-lambda
+
+        Args:
+            **kwargs: Optional arguments
+                ``tracer_provider``: a TracerProvider, defaults to global
+                ``event_context_extractor``: a method which takes the Lambda
+                    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.
+        """
+        lambda_handler = os.environ.get(ORIG_HANDLER, os.environ.get(_HANDLER))
+        # pylint: disable=attribute-defined-outside-init
+        (
+            self._wrapped_module_name,
+            self._wrapped_function_name,
+        ) = lambda_handler.rsplit(".", 1)
+
+        _instrument(
+            self._wrapped_module_name,
+            self._wrapped_function_name,
+            event_context_extractor=kwargs.get(
+                "event_context_extractor", _default_event_context_extractor
+            ),
+            tracer_provider=kwargs.get("tracer_provider"),
+        )
+
+    def _uninstrument(self, **kwargs):
+        unwrap(
+            import_module(self._wrapped_module_name),
+            self._wrapped_function_name,
+        )

instrumentation/opentelemetry-instrumentation-aws-lambda/tests/mocks/lambda_function.py

@@ -0,0 +1,17 @@
+# Copyright The OpenTelemetry Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+def handler(event, context):
+    return "200 ok"

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

@@ -0,0 +1,247 @@
+# Copyright The OpenTelemetry Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import os
+from importlib import import_module
+from unittest import mock
+
+from opentelemetry.environment_variables import OTEL_PROPAGATORS
+from opentelemetry.instrumentation.aws_lambda import (
+    _HANDLER,
+    _X_AMZN_TRACE_ID,
+    AwsLambdaInstrumentor,
+)
+from opentelemetry.propagate import get_global_textmap
+from opentelemetry.propagators.aws.aws_xray_propagator import (
+    TRACE_ID_FIRST_PART_LENGTH,
+    TRACE_ID_VERSION,
+)
+from opentelemetry.semconv.resource import ResourceAttributes
+from opentelemetry.semconv.trace import SpanAttributes
+from opentelemetry.test.test_base import TestBase
+from opentelemetry.trace import SpanKind
+from opentelemetry.trace.propagation.tracecontext import (
+    TraceContextTextMapPropagator,
+)
+
+
+class MockLambdaContext:
+    def __init__(self, aws_request_id, invoked_function_arn):
+        self.invoked_function_arn = invoked_function_arn
+        self.aws_request_id = aws_request_id
+
+
+MOCK_LAMBDA_CONTEXT = MockLambdaContext(
+    aws_request_id="mock_aws_request_id",
+    invoked_function_arn="arn://mock-lambda-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_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}"
+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"
+)
+
+# See more:
+# https://www.w3.org/TR/trace-context/#examples-of-http-traceparent-headers
+
+MOCK_W3C_TRACE_ID = 0x5CE0E9A56015FEC5AADFA328AE398115
+MOCK_W3C_PARENT_SPAN_ID = 0xAB54A98CEB1F0AD2
+MOCK_W3C_TRACE_CONTEXT_SAMPLED = (
+    f"00-{MOCK_W3C_TRACE_ID:x}-{MOCK_W3C_PARENT_SPAN_ID:x}-01"
+)
+
+MOCK_W3C_TRACE_STATE_KEY = "vendor_specific_key"
+MOCK_W3C_TRACE_STATE_VALUE = "test_value"
+
+
+def mock_execute_lambda(event=None):
+    """Mocks the AWS Lambda execution.
+
+    NOTE: We don't use `moto`'s `mock_lambda` because we are not instrumenting
+    calls to AWS Lambda using the AWS SDK. Instead, we are instrumenting AWS
+    Lambda itself.
+
+    See more:
+    https://docs.aws.amazon.com/lambda/latest/dg/runtimes-modify.html#runtime-wrapper
+
+    Args:
+        event: The Lambda event which may or may not be used by instrumentation.
+    """
+
+    module_name, handler_name = os.environ[_HANDLER].rsplit(".", 1)
+    handler_module = import_module(module_name.replace("/", "."))
+    getattr(handler_module, handler_name)(event, MOCK_LAMBDA_CONTEXT)
+
+
+class TestAwsLambdaInstrumentor(TestBase):
+    """AWS Lambda Instrumentation Testsuite"""
+
+    def setUp(self):
+        super().setUp()
+        self.common_env_patch = mock.patch.dict(
+            "os.environ", {_HANDLER: "mocks.lambda_function.handler"},
+        )
+        self.common_env_patch.start()
+
+        # NOTE: Whether AwsLambdaInstrumentor().instrument() is run is decided
+        # by each test case. It depends on if the test is for auto or manual
+        # instrumentation.
+
+    def tearDown(self):
+        super().tearDown()
+        self.common_env_patch.stop()
+        AwsLambdaInstrumentor().uninstrument()
+
+    def test_active_tracing(self):
+        test_env_patch = mock.patch.dict(
+            "os.environ",
+            {
+                **os.environ,
+                # Using Active tracing
+                _X_AMZN_TRACE_ID: MOCK_XRAY_TRACE_CONTEXT_SAMPLED,
+            },
+        )
+        test_env_patch.start()
+
+        AwsLambdaInstrumentor().instrument()
+
+        mock_execute_lambda()
+
+        spans = self.memory_exporter.get_finished_spans()
+
+        assert spans
+
+        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.assertEqual(span.kind, SpanKind.SERVER)
+        self.assertSpanHasAttributes(
+            span,
+            {
+                ResourceAttributes.FAAS_ID: MOCK_LAMBDA_CONTEXT.invoked_function_arn,
+                SpanAttributes.FAAS_EXECUTION: MOCK_LAMBDA_CONTEXT.aws_request_id,
+            },
+        )
+
+        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)
+
+        test_env_patch.stop()
+
+    def test_parent_context_from_lambda_event(self):
+        test_env_patch = mock.patch.dict(
+            "os.environ",
+            {
+                **os.environ,
+                # NOT Active Tracing
+                _X_AMZN_TRACE_ID: MOCK_XRAY_TRACE_CONTEXT_NOT_SAMPLED,
+                # NOT using the X-Ray Propagator
+                OTEL_PROPAGATORS: "tracecontext",
+            },
+        )
+        test_env_patch.start()
+
+        AwsLambdaInstrumentor().instrument()
+
+        mock_execute_lambda(
+            {
+                "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",
+                }
+            }
+        )
+
+        spans = self.memory_exporter.get_finished_spans()
+
+        assert spans
+
+        self.assertEqual(len(spans), 1)
+        span = spans[0]
+        self.assertEqual(span.get_span_context().trace_id, MOCK_W3C_TRACE_ID)
+
+        parent_context = span.parent
+        self.assertEqual(
+            parent_context.trace_id, span.get_span_context().trace_id
+        )
+        self.assertEqual(parent_context.span_id, MOCK_W3C_PARENT_SPAN_ID)
+        self.assertEqual(len(parent_context.trace_state), 3)
+        self.assertEqual(
+            parent_context.trace_state.get(MOCK_W3C_TRACE_STATE_KEY),
+            MOCK_W3C_TRACE_STATE_VALUE,
+        )
+        self.assertTrue(parent_context.is_remote)
+
+        test_env_patch.stop()
+
+    def test_using_custom_extractor(self):
+        def custom_event_context_extractor(lambda_event):
+            return get_global_textmap().extract(lambda_event["foo"]["headers"])
+
+        test_env_patch = mock.patch.dict(
+            "os.environ",
+            {
+                **os.environ,
+                # NOT Active Tracing
+                _X_AMZN_TRACE_ID: MOCK_XRAY_TRACE_CONTEXT_NOT_SAMPLED,
+                # NOT using the X-Ray Propagator
+                OTEL_PROPAGATORS: "tracecontext",
+            },
+        )
+        test_env_patch.start()
+
+        AwsLambdaInstrumentor().instrument(
+            event_context_extractor=custom_event_context_extractor,
+        )
+
+        mock_execute_lambda(
+            {
+                "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",
+                    }
+                }
+            }
+        )
+
+        spans = self.memory_exporter.get_finished_spans()
+
+        assert spans
+
+        self.assertEqual(len(spans), 1)
+        span = spans[0]
+        self.assertEqual(span.get_span_context().trace_id, MOCK_W3C_TRACE_ID)
+
+        parent_context = span.parent
+        self.assertEqual(
+            parent_context.trace_id, span.get_span_context().trace_id
+        )
+        self.assertEqual(parent_context.span_id, MOCK_W3C_PARENT_SPAN_ID)
+        self.assertEqual(len(parent_context.trace_state), 3)
+        self.assertEqual(
+            parent_context.trace_state.get(MOCK_W3C_TRACE_STATE_KEY),
+            MOCK_W3C_TRACE_STATE_VALUE,
+        )
+        self.assertTrue(parent_context.is_remote)
+
+        test_env_patch.stop()

tox.ini

@@ -25,6 +25,9 @@ envlist =
     py3{6,7,8,9,10}-test-instrumentation-aiopg
     ; instrumentation-aiopg intentionally excluded from pypy3
 
+    ; opentelemetry-instrumentation-aws-lambda
+    py3{6,7,8,9}-test-instrumentation-aws-lambda
+
     ; opentelemetry-instrumentation-botocore
     py3{6,7,8,9,10}-test-instrumentation-botocore
     pypy3-test-instrumentation-botocore
@@ -227,6 +230,7 @@ changedir =
   test-instrumentation-aiopg: instrumentation/opentelemetry-instrumentation-aiopg/tests
   test-instrumentation-asgi: instrumentation/opentelemetry-instrumentation-asgi/tests
   test-instrumentation-asyncpg: instrumentation/opentelemetry-instrumentation-asyncpg/tests
+  test-instrumentation-aws-lambda: instrumentation/opentelemetry-instrumentation-aws-lambda/tests
   test-instrumentation-boto: instrumentation/opentelemetry-instrumentation-boto/tests
   test-instrumentation-botocore: instrumentation/opentelemetry-instrumentation-botocore/tests
   test-instrumentation-celery: instrumentation/opentelemetry-instrumentation-celery/tests
@@ -288,6 +292,8 @@ commands_pre =
 
   asyncpg: pip install {toxinidir}/instrumentation/opentelemetry-instrumentation-asyncpg[test]
 
+  aws-lambda: pip install {toxinidir}/instrumentation/opentelemetry-instrumentation-aws-lambda[test]
+
   boto: pip install {toxinidir}/instrumentation/opentelemetry-instrumentation-botocore[test]
   boto: pip install {toxinidir}/instrumentation/opentelemetry-instrumentation-boto[test]
 
@@ -434,6 +440,7 @@ commands_pre =
   python -m pip install -e {toxinidir}/instrumentation/opentelemetry-instrumentation-tornado[test]
   python -m pip install -e {toxinidir}/instrumentation/opentelemetry-instrumentation-mysql[test]
   python -m pip install -e {toxinidir}/instrumentation/opentelemetry-instrumentation-httpx[test]
+  python -m pip install -e {toxinidir}/instrumentation/opentelemetry-instrumentation-aws-lambda[test]
   python -m pip install -e {toxinidir}/exporter/opentelemetry-exporter-datadog[test]
   python -m pip install -e {toxinidir}/sdk-extension/opentelemetry-sdk-extension-aws[test]
   python -m pip install -e {toxinidir}/propagator/opentelemetry-propagator-aws-xray[test]