Lambda Microservice changes

Author: AsakerMohd

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 opentelemetry.context import Context
+from opentelemetry.sdk.trace import ReadableSpan, Span, SpanProcessor
+from opentelemetry.trace import SpanKind
+from amazon.opentelemetry.distro._aws_attribute_keys import AWS_TRACE_LAMBDA_FLAG_MULTIPLE_SERVER
+
+class AwsLambdaSpanProcessor(SpanProcessor):
+    def __init__(self, instrumentation_names=None):
+        """
+        :param instrumentation_names: A set or list of instrumentation scope names
+            for which we want to mark as SERVER spans if they are INTERNAL.
+        """
+        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 span.kind == SpanKind.INTERNAL and scope.name in self.instrumentation_names:
+            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 force_flush(self, timeout_millis: int = 30000) -> bool:
+        """Flush any buffered data."""
+        return True
+
+    @override
+    def shutdown(self) -> None:
+        """Clean up."""
+        self.force_flush()

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

@@ -24,6 +24,7 @@
 from amazon.opentelemetry.distro.otlp_udp_exporter import OTLPUdpSpanExporter
 from amazon.opentelemetry.distro.sampler.aws_xray_remote_sampler import AwsXRayRemoteSampler
 from amazon.opentelemetry.distro.scope_based_exporter import ScopeBasedPeriodicExportingMetricReader
+from amazon.opentelemetry.distro.aws_lambda_span_processor import AwsLambdaSpanProcessor
 from amazon.opentelemetry.distro.scope_based_filtering_view import ScopeBasedRetainingView
 from opentelemetry.exporter.otlp.proto.http.metric_exporter import OTLPMetricExporter as OTLPHttpOTLPMetricExporter
 from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
@@ -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