Added new Lambda Multiple Span Flag

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.

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

aws-opentelemetry-distro/src/amazon/opentelemetry/distro/aws_lambda_span_processor.py

@@ -0,0 +1,44 @@
+# 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
+from opentelemetry.sdk.trace import ReadableSpan, Span, SpanProcessor
+from opentelemetry.trace import SpanKind
+
+
+class AwsLambdaSpanProcessor(SpanProcessor):
+    def __init__(self, instrumentation_names=None):
+        self.instrumentation_names = set(instrumentation_names or ["opentelemetry.instrumentation.flask"])
+        self.parent_lambda_span = None
+
+    @override
+    def on_start(self, span: Span, parent_context: Optional[Context] = None) -> None:
+        scope = getattr(span, "instrumentation_scope", None)
+        if span.kind == SpanKind.SERVER and scope.name == "opentelemetry.instrumentation.aws_lambda":
+            self.parent_lambda_span = span
+
+        if (
+            scope.name in self.instrumentation_names
+            and self.parent_lambda_span.get_span_context().span_id == span.parent.span_id
+        ):
+            if span.kind == SpanKind.INTERNAL:
+                span._kind = SpanKind.SERVER
+            self.parent_lambda_span.set_attribute(AWS_TRACE_LAMBDA_FLAG_MULTIPLE_SERVER, True)
+        return
+
+    @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

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

aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_aws_lambda_span_processor.py

@@ -0,0 +1,82 @@
+# 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.trace import Span, SpanContext, SpanKind
+
+
+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_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.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.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_span_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()

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]