From bab662a4101c2d93e5f29217d04ac5f2823930e6 Mon Sep 17 00:00:00 2001
From: Varsha GS <varsha.gs@ibm.com>
Date: Mon, 19 Aug 2024 16:37:09 +0530
Subject: [PATCH] fix: context propagation in nested spans

Signed-off-by: Varsha GS <varsha.gs@ibm.com>
---
 src/instana/tracer.py |  4 ++--
 tests/test_tracer.py  | 30 ++++++++++++++++++++++++++----
 2 files changed, 28 insertions(+), 6 deletions(-)

diff --git a/src/instana/tracer.py b/src/instana/tracer.py
index f7ffdcfc..fc5301f9 100644
--- a/src/instana/tracer.py
+++ b/src/instana/tracer.py
@@ -30,7 +30,7 @@
 from instana.recorder import StanRecorder
 from instana.sampling import InstanaSampler, Sampler
 from instana.span.kind import EXIT_SPANS
-from instana.span.span import InstanaSpan
+from instana.span.span import InstanaSpan, get_current_span
 from instana.span_context import SpanContext
 from instana.util.ids import generate_id
 
@@ -118,7 +118,7 @@ def start_span(
         record_exception: bool = True,
         set_status_on_exception: bool = True,
     ) -> InstanaSpan:
-        parent_context = span_context
+        parent_context = span_context if span_context else get_current_span().get_span_context()
 
         if parent_context is not None and not isinstance(parent_context, SpanContext):
             raise TypeError("parent_context must be an Instana SpanContext or None.")
diff --git a/tests/test_tracer.py b/tests/test_tracer.py
index 62f3bc15..16c2042f 100644
--- a/tests/test_tracer.py
+++ b/tests/test_tracer.py
@@ -1,15 +1,15 @@
 # (c) Copyright IBM Corp. 2024
 
-from opentelemetry.trace import set_span_in_context
-from opentelemetry.trace.span import _SPAN_ID_MAX_VALUE, INVALID_SPAN_ID
 import pytest
+
+from opentelemetry.trace.span import _SPAN_ID_MAX_VALUE
+
 from instana.agent.test import TestAgent
 from instana.recorder import StanRecorder
 from instana.sampling import InstanaSampler
-from instana.span.span import InstanaSpan
+from instana.span.span import InstanaSpan, get_current_span, INVALID_SPAN_ID, INVALID_SPAN
 from instana.span_context import SpanContext
 from instana.tracer import InstanaTracer, InstanaTracerProvider
-from opentelemetry.trace.span import _SPAN_ID_MAX_VALUE, INVALID_SPAN_ID
 
 
 def test_tracer_defaults(tracer_provider: InstanaTracerProvider) -> None:
@@ -99,6 +99,28 @@ def test_tracer_start_as_current_span(tracer_provider: InstanaTracerProvider) ->
         assert span.name == span_name
 
 
+def test_tracer_nested_span(tracer_provider: InstanaTracerProvider) -> None:
+    tracer = InstanaTracer(
+        tracer_provider.sampler,
+        tracer_provider._span_processor,
+        tracer_provider._exporter,
+        tracer_provider._propagators,
+    )
+    parent_span_name = "parent-span"
+    child_span_name = "child-span"
+    with tracer.start_as_current_span(name=parent_span_name) as pspan:
+        assert get_current_span() is pspan
+        with tracer.start_as_current_span(name=child_span_name) as cspan:
+            assert get_current_span() is cspan
+            assert cspan.parent_id == pspan.context.span_id
+        # child span goes out of scope
+        assert cspan.end_time is not None
+        assert get_current_span() is pspan
+    # parent span goes out of scope
+    assert pspan.end_time is not None
+    assert get_current_span() is INVALID_SPAN
+
+
 def test_tracer_create_span_context(
     span_context: SpanContext, tracer_provider: InstanaTracerProvider
 ) -> None: