clear context on Laminar.force_flush() (#215)

* clear context on Laminar.force_flush()

* tests to verify context is properly cleared

Author: dinmukhamedm

src/lmnr/opentelemetry_lib/tracing/__init__.py

@@ -12,12 +12,11 @@
 )
 from lmnr.opentelemetry_lib.tracing.context import (
     attach_context,
+    clear_context,
     pop_span_context as ctx_pop_span_context,
     get_current_context,
     get_token_stack,
     push_span_context as ctx_push_span_context,
-    _isolated_token_stack,
-    _isolated_token_stack_storage,
     set_token_stack,
 )
 
@@ -230,14 +229,7 @@ def clear(cls):
         if isinstance(cls.instance._span_processor, LaminarSpanProcessor):
             cls.instance._span_processor.clear()
         # Clear the isolated context state for clean test state
-        try:
-            _isolated_token_stack.set([])
-        except LookupError:
-            pass
-        if hasattr(_isolated_token_stack_storage, "token_stack"):
-            _isolated_token_stack_storage.token_stack = []
-        # Reset the isolated context to a fresh state
-        attach_context(Context())
+        clear_context()
 
     def shutdown(self):
         if self._tracer_provider is None:
@@ -254,6 +246,9 @@ def force_reinit_processor(self):
         if isinstance(self._span_processor, LaminarSpanProcessor):
             self._span_processor.force_flush()
             self._span_processor.force_reinit()
+            # Clear the isolated context to prevent subsequent invocations
+            # (e.g., in Lambda) from continuing traces from previous invocations
+            clear_context()
         else:
             self._logger.warning(
                 "Not using LaminarSpanProcessor, cannot force reinit processor"

src/lmnr/opentelemetry_lib/tracing/context.py

@@ -141,3 +141,29 @@ def push_span_context(context: Context) -> None:
     token_stack = get_token_stack().copy()
     token_stack.append(token)
     set_token_stack(token_stack)
+
+
+def clear_context() -> None:
+    """Clear the isolated context and token stack.
+
+    This is primarily used during force_flush operations in Lambda-like
+    environments to ensure subsequent invocations don't continue traces
+    from previous invocations.
+
+    Warning: This should only be called when you're certain no spans are
+    actively being processed, as it will reset all context state.
+    """
+    # Clear the token stack first
+    try:
+        _isolated_token_stack.set([])
+    except LookupError:
+        pass
+
+    # Clear thread-local storage if it exists
+    if hasattr(_isolated_token_stack_storage, "token_stack"):
+        _isolated_token_stack_storage.token_stack = []
+
+    # Reset the context to a fresh empty context
+    # This doesn't require manually detaching tokens since we're
+    # intentionally resetting everything to a clean state
+    _ISOLATED_RUNTIME_CONTEXT._current_context.set(Context())

src/lmnr/sdk/laminar.py

@@ -892,7 +892,9 @@ def flush(cls) -> bool:
 
     @classmethod
     def force_flush(cls):
-        """Force flush the internal tracer.
+        """Force flush the internal tracer. WARNING: Any active spans are
+        removed from context; that is, spans started afterwards will start
+        a new trace.
 
         Actually shuts down the span processor and re-initializes it as long
         as it is a LaminarSpanProcessor. This is not recommended in production

tests/test_context.py

@@ -0,0 +1,84 @@
+from opentelemetry.trace.span import INVALID_SPAN_ID
+from lmnr import Laminar, observe
+from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
+
+
+def test_clear_context_observe(span_exporter: InMemorySpanExporter):
+    @observe()
+    def inner():
+        Laminar.set_trace_user_id("test_user_id_2")
+        return "inner"
+
+    @observe()
+    def outer():
+        Laminar.set_trace_user_id("test_user_id_1")
+        Laminar.force_flush()
+        return inner()
+
+    outer()
+    spans = span_exporter.get_finished_spans()
+    assert len(spans) == 2
+    outer_span = [s for s in spans if s.name == "outer"][0]
+    inner_span = [s for s in spans if s.name == "inner"][0]
+    assert (
+        outer_span.attributes["lmnr.association.properties.user_id"] == "test_user_id_1"
+    )
+    assert (
+        inner_span.attributes["lmnr.association.properties.user_id"] == "test_user_id_2"
+    )
+
+    assert inner_span.parent is None or inner_span.parent.span_id == INVALID_SPAN_ID
+    assert (
+        inner_span.get_span_context().trace_id != outer_span.get_span_context().trace_id
+    )
+
+
+def test_clear_context_start_as_current_span(span_exporter: InMemorySpanExporter):
+    with Laminar.start_as_current_span("outer"):
+        Laminar.set_trace_user_id("test_user_id_1")
+        Laminar.force_flush()
+        with Laminar.start_as_current_span("inner"):
+            Laminar.set_trace_user_id("test_user_id_2")
+            pass
+
+    spans = span_exporter.get_finished_spans()
+    assert len(spans) == 2
+    outer_span = [s for s in spans if s.name == "outer"][0]
+    inner_span = [s for s in spans if s.name == "inner"][0]
+    assert (
+        outer_span.attributes["lmnr.association.properties.user_id"] == "test_user_id_1"
+    )
+    assert (
+        inner_span.attributes["lmnr.association.properties.user_id"] == "test_user_id_2"
+    )
+
+    assert inner_span.parent is None or inner_span.parent.span_id == INVALID_SPAN_ID
+    assert (
+        inner_span.get_span_context().trace_id != outer_span.get_span_context().trace_id
+    )
+
+
+def test_clear_context_start_active_span(span_exporter: InMemorySpanExporter):
+    span = Laminar.start_active_span("outer")
+    Laminar.set_trace_user_id("test_user_id_1")
+    Laminar.force_flush()
+    span2 = Laminar.start_active_span("inner")
+    Laminar.set_trace_user_id("test_user_id_2")
+    span2.end()
+    span.end()
+
+    spans = span_exporter.get_finished_spans()
+    assert len(spans) == 2
+    outer_span = [s for s in spans if s.name == "outer"][0]
+    inner_span = [s for s in spans if s.name == "inner"][0]
+    assert (
+        outer_span.attributes["lmnr.association.properties.user_id"] == "test_user_id_1"
+    )
+    assert (
+        inner_span.attributes["lmnr.association.properties.user_id"] == "test_user_id_2"
+    )
+
+    assert inner_span.parent is None or inner_span.parent.span_id == INVALID_SPAN_ID
+    assert (
+        inner_span.get_span_context().trace_id != outer_span.get_span_context().trace_id
+    )