Trace continuation strategy prototype 2

opentelemetry-api/src/opentelemetry/trace/__init__.py

@@ -73,6 +73,7 @@
 
 from opentelemetry import context as context_api
 from opentelemetry.attributes import BoundedAttributes
+from opentelemetry.context import get_current
 from opentelemetry.context.context import Context
 from opentelemetry.environment_variables import OTEL_PYTHON_TRACER_PROVIDER
 from opentelemetry.trace.propagation import (
@@ -402,6 +403,31 @@ def function():
             The newly-created span.
         """
 
+    @_agnosticcontextmanager
+    @abstractmethod
+    def apply_egress_continuation(
+        self,
+        kind: SpanKind,
+        context: Context | None = None,
+        attributes: types.Attributes = None,
+    ) -> Iterator[Context]:
+        """Context manager for getting a Context informed with continuation hints
+        for egress propagation based on span kind and attributes.
+
+        Example::
+
+            with tracer.apply_egress_continuation(kind=Span.CLIENT, attributes=attributes) as injection_context:
+                propagate.inject(carrier, context=injection_context, setter=setter)
+
+        Args:
+            context: An optional Context. Defaults to the global context.
+            kind: The span's kind.
+            attributes: The span's attributes.
+
+        Yields:
+            A context that contains egress continuation hints for propagators
+        """
+
 
 class ProxyTracer(Tracer):
     # pylint: disable=W0222,signature-differs
@@ -442,6 +468,18 @@ def start_as_current_span(self, *args, **kwargs) -> Iterator[Span]:
         with self._tracer.start_as_current_span(*args, **kwargs) as span:  # type: ignore
             yield span
 
+    @_agnosticcontextmanager
+    def apply_egress_continuation(
+        self,
+        kind: SpanKind,
+        context: Context | None = None,
+        attributes: types.Attributes = None,
+    ) -> Iterator[Context]:
+        with self._tracer.apply_egress_continuation(
+            kind, context, attributes
+        ) as injection_context:
+            yield injection_context
+
 
 class NoOpTracer(Tracer):
     """The default Tracer, used when no Tracer implementation is available.
@@ -507,6 +545,16 @@ def start_as_current_span(
         ) as span:
             yield span
 
+    @_agnosticcontextmanager
+    def apply_egress_continuation(
+        self,
+        kind: SpanKind,
+        context: Context | None = None,
+        attributes: types.Attributes = None,
+    ) -> Iterator[Context]:
+        injected_context = context if context else get_current()
+        yield injected_context
+
 
 @deprecated("You should use NoOpTracer. Deprecated since version 1.9.0.")
 class _DefaultTracer(NoOpTracer):

opentelemetry-api/src/opentelemetry/trace/propagation/tracecontext.py

@@ -3,12 +3,44 @@
 #
 import re
 
+from opentelemetry import context as context_api
 from opentelemetry import trace
 from opentelemetry.context.context import Context
 from opentelemetry.propagators import textmap
 from opentelemetry.trace import format_span_id, format_trace_id
 from opentelemetry.trace.span import TraceState
 
+_SUPPRESS_TRACE_CONTEXT_INJECTION_KEY = context_api.create_key(
+    "suppress-trace-context-injection"
+)
+
+
+def suppress_trace_context_injection(
+    context: Context | None = None,
+) -> Context:
+    """Returns a context that suppresses W3C Trace Context injection."""
+    return context_api.set_value(
+        _SUPPRESS_TRACE_CONTEXT_INJECTION_KEY, True, context
+    )
+
+
+def enable_trace_context_injection(
+    context: Context | None = None,
+) -> Context:
+    """Returns a context that allows W3C Trace Context injection."""
+    return context_api.set_value(
+        _SUPPRESS_TRACE_CONTEXT_INJECTION_KEY, False, context
+    )
+
+
+def is_trace_context_injection_suppressed(
+    context: Context | None = None,
+) -> bool:
+    """Returns whether W3C Trace Context injection is suppressed."""
+    return bool(
+        context_api.get_value(_SUPPRESS_TRACE_CONTEXT_INJECTION_KEY, context)
+    )
+
 
 class TraceContextTextMapPropagator(textmap.TextMapPropagator):
     """Extracts and injects using w3c TraceContext's headers."""
@@ -84,6 +116,9 @@ def inject(
 
         See `opentelemetry.propagators.textmap.TextMapPropagator.inject`
         """
+        if is_trace_context_injection_suppressed(context):
+            return
+
         span = trace.get_current_span(context)
         span_context = span.get_span_context()
         if span_context == trace.INVALID_SPAN_CONTEXT:

opentelemetry-api/tests/trace/propagation/test_tracecontexthttptextformat.py

@@ -170,6 +170,34 @@ def test_propagate_invalid_context(self):
         FORMAT.inject(output, context=ctx)
         self.assertFalse("traceparent" in output)
 
+    def test_suppress_trace_context_injection(self):
+        """Do not propagate trace context when injection is suppressed."""
+        output: dict[str, str] = {}
+        span = trace.NonRecordingSpan(
+            trace.SpanContext(self.TRACE_ID, self.SPAN_ID, is_remote=False)
+        )
+        ctx = trace.set_span_in_context(span)
+        ctx = tracecontext.suppress_trace_context_injection(ctx)
+
+        FORMAT.inject(output, context=ctx)
+
+        self.assertFalse("traceparent" in output)
+        self.assertFalse("tracestate" in output)
+
+    def test_enable_trace_context_injection(self):
+        """Propagate trace context after inherited suppression is cleared."""
+        output: dict[str, str] = {}
+        span = trace.NonRecordingSpan(
+            trace.SpanContext(self.TRACE_ID, self.SPAN_ID, is_remote=False)
+        )
+        ctx = trace.set_span_in_context(span)
+        ctx = tracecontext.suppress_trace_context_injection(ctx)
+        ctx = tracecontext.enable_trace_context_injection(ctx)
+
+        FORMAT.inject(output, context=ctx)
+
+        self.assertTrue("traceparent" in output)
+
     def test_tracestate_empty_header(self):
         """Test tracestate with an additional empty header (should be ignored)"""
         span = trace.get_current_span(

opentelemetry-api/tests/trace/test_tracer.py

@@ -5,14 +5,20 @@
 import asyncio
 from unittest import TestCase
 
+from opentelemetry.context import get_current
 from opentelemetry.trace import (
     INVALID_SPAN,
+    NonRecordingSpan,
     NoOpTracer,
     Span,
+    SpanContext,
+    SpanKind,
     Tracer,
     _agnosticcontextmanager,
     get_current_span,
+    set_span_in_context,
 )
+from opentelemetry.trace.propagation import tracecontext
 
 
 class TestTracer(TestCase):
@@ -27,6 +33,45 @@ def test_start_as_current_span_context_manager(self):
         with self.tracer.start_as_current_span("") as span:
             self.assertIsInstance(span, Span)
 
+    def test_apply_egress_continuation_noop_tracer_preserves_context(self):
+        propagator = tracecontext.TraceContextTextMapPropagator()
+        span_context = SpanContext(
+            trace_id=0x000000000000000000000000DEADBEEF,
+            span_id=0x00000000DEADBEF0,
+            is_remote=False,
+        )
+        context = set_span_in_context(NonRecordingSpan(span_context))
+
+        with self.tracer.apply_egress_continuation(
+            kind=SpanKind.CLIENT,
+            context=context,
+            attributes={"server.address": "api.example.com"},
+        ) as injection_context:
+            carrier: dict[str, str] = {}
+            propagator.inject(carrier, context=injection_context)
+
+        self.assertIn("traceparent", carrier)
+
+    def test_apply_egress_continuation_noop_tracer_preserves_suppression(self):
+        propagator = tracecontext.TraceContextTextMapPropagator()
+        span_context = SpanContext(
+            trace_id=0x000000000000000000000000DEADBEEF,
+            span_id=0x00000000DEADBEF0,
+            is_remote=False,
+        )
+        context = set_span_in_context(NonRecordingSpan(span_context))
+        context = tracecontext.suppress_trace_context_injection(context)
+
+        with self.tracer.apply_egress_continuation(
+            kind=SpanKind.CLIENT,
+            context=context,
+            attributes={"server.address": "api.example.com"},
+        ) as injection_context:
+            carrier: dict[str, str] = {}
+            propagator.inject(carrier, context=injection_context)
+
+        self.assertNotIn("traceparent", carrier)
+
     def test_start_as_current_span_decorator(self):
         # using a list to track the mock call order
         calls = []
@@ -41,6 +86,10 @@ def start_as_current_span(self, *args, **kwargs):  # type: ignore
                 yield INVALID_SPAN
                 calls.append(9)
 
+            @_agnosticcontextmanager  # pylint: disable=protected-access
+            def apply_egress_continuation(self, *args, **kwargs):  # type: ignore
+                yield get_current()
+
         mock_tracer = MockTracer()
 
         # test 1 : sync function