Updated PR

Author: AsakerMohd

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

@@ -5,29 +5,31 @@
 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.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"])
-        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:
+        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_span.kind == SpanKind.SERVER and parent_scope.name == "opentelemetry.instrumentation.aws_lambda":
                 span._kind = SpanKind.SERVER
-            self.parent_lambda_span.set_attribute(AWS_TRACE_LAMBDA_FLAG_MULTIPLE_SERVER, True)
+                parent_span.set_attribute(AWS_TRACE_LAMBDA_FLAG_MULTIPLE_SERVER, True)
+
         return
 
     @override

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

@@ -5,7 +5,9 @@
 
 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, get_value, set_value
 from opentelemetry.trace import Span, SpanContext, SpanKind
+from opentelemetry.trace.propagation import _SPAN_KEY
 
 
 class TestAwsLambdaSpanProcessor(TestCase):
@@ -20,6 +22,8 @@ def setUp(self):
         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)
 
@@ -35,7 +39,7 @@ def test_lambda_span_multiple_server_flag_internal_api(self, mock_span_class):
         flask_span.kind = SpanKind.INTERNAL
         flask_span.parent = self.lambda_span_context
 
-        self.processor.on_start(flask_span)
+        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])
@@ -53,7 +57,7 @@ def test_lambda_span_multiple_server_flag_server_api(self, mock_span_class):
         flask_span.kind = SpanKind.SERVER
         flask_span.parent = self.lambda_span_context
 
-        self.processor.on_start(flask_span)
+        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])
@@ -71,7 +75,7 @@ def test_lambda_span_single_server_span(self, mock_span_class):
         flask_span.kind = SpanKind.CLIENT
         flask_span.parent = self.lambda_span_context
 
-        self.processor.on_start(flask_span, 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()