Added new Lambda Multiple Span Flag (aws-observability#334)

*Description of changes:*
Added a new lambda span processor that 
1. Overwrites the INTERNAL API Span (example:
`opentelemetry.instrumentation.flask`) to SERVER
2. adds a new flag to Lambda Spans if there are multiple server spans in
the trace.

Also updated the Operation Name logic so that if the span is an
`opentelemetry.instrumentation.flask` span, the operation name is the
API Route instead of the FunctionHandler operation name. This makes it
easier to identify that the Span is an API span in the Application
Signals console.

*Testing:*
Added unit test to test processor. Also tested by creating a lambda
layer. The lambda span does have the new flag and the operation name for
the API Span is the API Route.

By submitting this pull request, I confirm that you can use, modify,
copy, and redistribute this contribution, under the terms of your
choice.

Author: AsakerMohd

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,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

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,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()

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]