diff --git a/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/_aws_attribute_keys.py b/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/_aws_attribute_keys.py index 879f86653..23ba661af 100644 --- a/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/_aws_attribute_keys.py +++ b/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/_aws_attribute_keys.py @@ -13,6 +13,7 @@ AWS_SDK_DESCENDANT: str = "aws.sdk.descendant" AWS_CONSUMER_PARENT_SPAN_KIND: str = "aws.consumer.parent.span.kind" AWS_TRACE_FLAG_SAMPLED: str = "aws.trace.flag.sampled" +AWS_TRACE_LAMBDA_FLAG_MULTIPLE_SERVER: str = "aws.trace.lambda.multiple-server" AWS_CLOUDFORMATION_PRIMARY_IDENTIFIER: str = "aws.remote.resource.cfn.primary.identifier" # AWS_#_NAME attributes are not supported in python as they are not part of the Semantic Conventions. diff --git a/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/_aws_span_processing_util.py b/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/_aws_span_processing_util.py index 24aaa68dc..21e19afa9 100644 --- a/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/_aws_span_processing_util.py +++ b/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/_aws_span_processing_util.py @@ -57,7 +57,8 @@ def get_ingress_operation(__, span: ReadableSpan) -> str: with the first API path parameter" if the default span name is None, UnknownOperation or http.method value. """ operation: str = span.name - if _AWS_LAMBDA_FUNCTION_NAME in os.environ: + scope = getattr(span, "instrumentation_scope", None) + if _AWS_LAMBDA_FUNCTION_NAME in os.environ and scope.name != "opentelemetry.instrumentation.flask": operation = os.environ.get(_AWS_LAMBDA_FUNCTION_NAME) + "/FunctionHandler" elif should_use_internal_operation(span): operation = INTERNAL_OPERATION diff --git a/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/aws_lambda_span_processor.py b/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/aws_lambda_span_processor.py new file mode 100644 index 000000000..d29b7bbbb --- /dev/null +++ b/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/aws_lambda_span_processor.py @@ -0,0 +1,47 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 +from typing import Optional + +from typing_extensions import override + +from amazon.opentelemetry.distro._aws_attribute_keys import AWS_TRACE_LAMBDA_FLAG_MULTIPLE_SERVER +from opentelemetry.context import Context, get_value +from opentelemetry.sdk.trace import ReadableSpan, Span, SpanProcessor +from opentelemetry.trace import SpanKind +from opentelemetry.trace.propagation import _SPAN_KEY + + +class AwsLambdaSpanProcessor(SpanProcessor): + def __init__(self, instrumentation_names=None): + self.instrumentation_names = set(instrumentation_names or ["opentelemetry.instrumentation.flask"]) + + @override + def on_start(self, span: Span, parent_context: Optional[Context] = None) -> None: + + scope = getattr(span, "instrumentation_scope", None) + if scope.name in self.instrumentation_names: + parent_span = get_value(_SPAN_KEY, context=parent_context) + + if parent_span is None: + return + + parent_scope = getattr(parent_span, "instrumentation_scope", None) + if parent_scope.name == "opentelemetry.instrumentation.aws_lambda": + span._kind = SpanKind.SERVER + parent_span.set_attribute(AWS_TRACE_LAMBDA_FLAG_MULTIPLE_SERVER, True) + + return + + # pylint: disable=no-self-use + @override + def on_end(self, span: ReadableSpan) -> None: + return + + @override + def shutdown(self) -> None: + self.force_flush() + + # pylint: disable=no-self-use + @override + def force_flush(self, timeout_millis: int = None) -> bool: + return True diff --git a/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/aws_opentelemetry_configurator.py b/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/aws_opentelemetry_configurator.py index d5cea0989..f3a897128 100644 --- a/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/aws_opentelemetry_configurator.py +++ b/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/aws_opentelemetry_configurator.py @@ -16,6 +16,7 @@ AttributePropagatingSpanProcessorBuilder, ) from amazon.opentelemetry.distro.aws_batch_unsampled_span_processor import BatchUnsampledSpanProcessor +from amazon.opentelemetry.distro.aws_lambda_span_processor import AwsLambdaSpanProcessor from amazon.opentelemetry.distro.aws_metric_attributes_span_exporter_builder import ( AwsMetricAttributesSpanExporterBuilder, ) @@ -343,6 +344,7 @@ def _customize_span_processors(provider: TracerProvider, resource: Resource) -> # Export 100% spans and not export Application-Signals metrics if on Lambda. if _is_lambda_environment(): _export_unsampled_span_for_lambda(provider, resource) + provider.add_span_processor(AwsLambdaSpanProcessor()) return # Construct meterProvider diff --git a/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_aws_lambda_span_processor.py b/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_aws_lambda_span_processor.py new file mode 100644 index 000000000..202e2364b --- /dev/null +++ b/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_aws_lambda_span_processor.py @@ -0,0 +1,86 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 +from unittest import TestCase +from unittest.mock import MagicMock, patch + +from amazon.opentelemetry.distro._aws_attribute_keys import AWS_TRACE_LAMBDA_FLAG_MULTIPLE_SERVER +from amazon.opentelemetry.distro.aws_lambda_span_processor import AwsLambdaSpanProcessor +from opentelemetry.context import Context, set_value +from opentelemetry.trace import Span, SpanContext, SpanKind +from opentelemetry.trace.propagation import _SPAN_KEY + + +class TestAwsLambdaSpanProcessor(TestCase): + + def setUp(self): + self.processor = AwsLambdaSpanProcessor() + self.lambda_span: Span = MagicMock() + self.lambda_span.instrumentation_scope.name = "opentelemetry.instrumentation.aws_lambda" + self.lambda_span.kind = SpanKind.SERVER + + self.lambda_span_context: SpanContext = MagicMock() + self.lambda_span_context.trace_id = "ABC" + self.lambda_span_context.span_id = "lambda_id" + + self.lambda_context: Context = set_value(_SPAN_KEY, self.lambda_span) + + self.lambda_span.get_span_context.return_value = self.lambda_span_context + self.processor.on_start(self.lambda_span) + + def tearDown(self): + self.processor.on_end(self.lambda_span) + self.processor.shutdown() + + @patch("opentelemetry.sdk.trace.Span") + def test_lambda_span_multiple_server_flag_internal_api(self, mock_span_class): + + flask_span = mock_span_class.return_value + flask_span.instrumentation_scope.name = "opentelemetry.instrumentation.flask" + flask_span.kind = SpanKind.INTERNAL + flask_span.parent = self.lambda_span_context + + self.processor.on_start(flask_span, self.lambda_context) + + self.assertEqual(flask_span._kind, SpanKind.SERVER) + self.assertIn(AWS_TRACE_LAMBDA_FLAG_MULTIPLE_SERVER, self.lambda_span.set_attribute.call_args_list[0][0][0]) + + self.processor.on_end(flask_span) + self.processor.on_end(self.lambda_span) + + self.processor.shutdown() + + @patch("opentelemetry.sdk.trace.Span") + def test_lambda_span_multiple_server_flag_server_api(self, mock_span_class): + + flask_span = mock_span_class.return_value + flask_span.instrumentation_scope.name = "opentelemetry.instrumentation.flask" + flask_span.kind = SpanKind.SERVER + flask_span.parent = self.lambda_span_context + + self.processor.on_start(flask_span, self.lambda_context) + + self.assertEqual(flask_span.kind, SpanKind.SERVER) + self.assertIn(AWS_TRACE_LAMBDA_FLAG_MULTIPLE_SERVER, self.lambda_span.set_attribute.call_args_list[0][0][0]) + + self.processor.on_end(flask_span) + self.processor.on_end(self.lambda_span) + + self.processor.shutdown() + + @patch("opentelemetry.sdk.trace.Span") + def test_lambda_span_single_server_span(self, mock_span_class): + + flask_span = mock_span_class.return_value + flask_span.instrumentation_scope.name = "opentelemetry.instrumentation.http" + flask_span.kind = SpanKind.CLIENT + flask_span.parent = self.lambda_span_context + + self.processor.on_start(flask_span, self.lambda_context) + + self.assertEqual(flask_span.kind, SpanKind.CLIENT) + flask_span.set_attribute.assert_not_called() + + self.processor.on_end(flask_span) + self.processor.on_end(self.lambda_span) + + self.processor.shutdown() diff --git a/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_aws_opentelementry_configurator.py b/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_aws_opentelementry_configurator.py index cb29d0533..d2963c4ec 100644 --- a/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_aws_opentelementry_configurator.py +++ b/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_aws_opentelementry_configurator.py @@ -318,7 +318,7 @@ def test_customize_span_processors_lambda(self): os.environ.setdefault("OTEL_AWS_APPLICATION_SIGNALS_ENABLED", "True") os.environ.setdefault("AWS_LAMBDA_FUNCTION_NAME", "myLambdaFunc") _customize_span_processors(mock_tracer_provider, Resource.get_empty()) - self.assertEqual(mock_tracer_provider.add_span_processor.call_count, 2) + self.assertEqual(mock_tracer_provider.add_span_processor.call_count, 3) first_processor: SpanProcessor = mock_tracer_provider.add_span_processor.call_args_list[0].args[0] self.assertIsInstance(first_processor, AttributePropagatingSpanProcessor) second_processor: SpanProcessor = mock_tracer_provider.add_span_processor.call_args_list[1].args[0]