Update OpenCensus shim to update OpenCensus execution context (open-telemetry#3231)

aabmass · web-flow · commit 17aa1e862c45 · 2023-04-04T01:52:41.000+05:30
* Update OpenCensus shim to update OpenCensus execution context

Previously, the shim was only writing to OTel context. This updates it to write to both OC and OTel.

Also updated the code to use the `span_context` provided when instrumentation instantiates the `Tracer`.
diff --git a/shim/opentelemetry-opencensus-shim/src/opentelemetry/shim/opencensus/_patch.py b/shim/opentelemetry-opencensus-shim/src/opentelemetry/shim/opencensus/_patch.py
@@ -12,9 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from functools import lru_cache
 from logging import getLogger
 from typing import Optional
 
+from opencensus.trace.span_context import SpanContext
 from opencensus.trace.tracer import Tracer
 from opencensus.trace.tracers.noop_tracer import NoopTracer
 
@@ -33,10 +35,19 @@ def install_shim(
         __version__,
         tracer_provider=tracer_provider,
     )
-    shim_tracer = ShimTracer(NoopTracer(), otel_tracer=otel_tracer)
 
-    def fget_tracer(self) -> ShimTracer:
-        return shim_tracer
+    @lru_cache()
+    def cached_shim_tracer(span_context: SpanContext) -> ShimTracer:
+        return ShimTracer(
+            NoopTracer(),
+            oc_span_context=span_context,
+            otel_tracer=otel_tracer,
+        )
+
+    def fget_tracer(self: Tracer) -> ShimTracer:
+        # self.span_context is how instrumentations pass propagated context into OpenCensus e.g.
+        # https://github.com/census-instrumentation/opencensus-python/blob/fd064f438c5e490d25b004ee2545be55d2e28679/contrib/opencensus-ext-flask/opencensus/ext/flask/flask_middleware.py#L147-L153
+        return cached_shim_tracer(self.span_context)
 
     def fset_tracer(self, value) -> None:
         # ignore attempts to set the value
@@ -45,8 +56,8 @@ def fset_tracer(self, value) -> None:
     # Tracer's constructor sets self.tracer to either a NoopTracer or ContextTracer depending
     # on sampler:
     # https://github.com/census-instrumentation/opencensus-python/blob/2e08df591b507612b3968be8c2538dedbf8fab37/opencensus/trace/tracer.py#L63.
-    # We monkeypatch Tracer.tracer with a property to return the shim instance instead. This
-    # makes all instances of Tracer (even those already created) use the ShimTracer singleton.
+    # We monkeypatch Tracer.tracer with a property to return a shim instance instead. This
+    # makes all instances of Tracer (even those already created) use a ShimTracer.
     Tracer.tracer = property(fget_tracer, fset_tracer)
     _logger.info("Installed OpenCensus shim")
 
diff --git a/shim/opentelemetry-opencensus-shim/src/opentelemetry/shim/opencensus/_shim_span.py b/shim/opentelemetry-opencensus-shim/src/opentelemetry/shim/opencensus/_shim_span.py
@@ -17,7 +17,8 @@
 from typing import TYPE_CHECKING
 
 import wrapt
-from opencensus.trace.base_span import BaseSpan
+from opencensus.trace import execution_context
+from opencensus.trace.blank_span import BlankSpan
 from opencensus.trace.span import SpanKind
 from opencensus.trace.status import Status
 from opencensus.trace.time_event import MessageEvent
@@ -62,7 +63,7 @@ def _opencensus_time_to_nanos(timestamp: str) -> int:
 class ShimSpan(wrapt.ObjectProxy):
     def __init__(
         self,
-        wrapped: BaseSpan,
+        wrapped: BlankSpan,
         *,
         otel_span: trace.Span,
         shim_tracer: "ShimTracer",
@@ -158,6 +159,9 @@ def __exit__(self, exception_type, exception_value, traceback):
         )
         # OpenCensus Span.__exit__() calls Tracer.end_span()
         # https://github.com/census-instrumentation/opencensus-python/blob/2e08df591b507612b3968be8c2538dedbf8fab37/opencensus/trace/span.py#L390
-        # but that would cause the OTel span to be ended twice. Instead just detach it from
-        # context directly.
+        # but that would cause the OTel span to be ended twice. Instead, this code just copies
+        # the context teardown from that method.
         context.detach(self._self_token)
+        execution_context.set_current_span(
+            self._self_shim_tracer.current_span()
+        )
diff --git a/shim/opentelemetry-opencensus-shim/src/opentelemetry/shim/opencensus/_shim_tracer.py b/shim/opentelemetry-opencensus-shim/src/opentelemetry/shim/opencensus/_shim_tracer.py
@@ -15,15 +15,19 @@
 import logging
 
 import wrapt
+from opencensus.trace import execution_context
 from opencensus.trace.blank_span import BlankSpan
+from opencensus.trace.span_context import SpanContext
 from opencensus.trace.tracers.base import Tracer as BaseTracer
+from opencensus.trace.tracestate import Tracestate
 
 from opentelemetry import context, trace
 from opentelemetry.shim.opencensus._shim_span import ShimSpan
 
 _logger = logging.getLogger(__name__)
 
 _SHIM_SPAN_KEY = context.create_key("opencensus-shim-span-key")
+_SAMPLED = trace.TraceFlags(trace.TraceFlags.SAMPLED)
 
 
 def set_shim_span_in_context(
@@ -36,12 +40,57 @@ def get_shim_span_in_context() -> ShimSpan:
     return context.get_value(_SHIM_SPAN_KEY)
 
 
+def set_oc_span_in_context(
+    oc_span_context: SpanContext, ctx: context.Context
+) -> context.Context:
+    """Returns a new OTel context based on ctx with oc_span_context set as the current span"""
+
+    # If no SpanContext is passed to the opencensus.trace.tracer.Tracer, it creates a new one
+    # with a random trace ID and a None span ID to be the parent:
+    # https://github.com/census-instrumentation/opencensus-python/blob/2e08df591b507612b3968be8c2538dedbf8fab37/opencensus/trace/tracer.py#L47.
+    #
+    # OpenTelemetry considers this an invalid SpanContext and will ignore it, so we can just
+    # return early
+    if oc_span_context.span_id is None:
+        return ctx
+
+    trace_id = int(oc_span_context.trace_id, 16)
+    span_id = int(oc_span_context.span_id, 16)
+    is_remote = oc_span_context.from_header
+    trace_flags = (
+        _SAMPLED if oc_span_context.trace_options.get_enabled() else None
+    )
+    trace_state = (
+        trace.TraceState(tuple(oc_span_context.tracestate.items()))
+        # OC SpanContext does not validate this type
+        if isinstance(oc_span_context.tracestate, Tracestate)
+        else None
+    )
+
+    return trace.set_span_in_context(
+        trace.NonRecordingSpan(
+            trace.SpanContext(
+                trace_id=trace_id,
+                span_id=span_id,
+                is_remote=is_remote,
+                trace_flags=trace_flags,
+                trace_state=trace_state,
+            )
+        )
+    )
+
+
 # pylint: disable=abstract-method
 class ShimTracer(wrapt.ObjectProxy):
     def __init__(
-        self, wrapped: BaseTracer, *, otel_tracer: trace.Tracer
+        self,
+        wrapped: BaseTracer,
+        *,
+        oc_span_context: SpanContext,
+        otel_tracer: trace.Tracer
     ) -> None:
         super().__init__(wrapped)
+        self._self_oc_span_context = oc_span_context
         self._self_otel_tracer = otel_tracer
 
     # For now, finish() is not implemented by the shim. It would require keeping a list of all
@@ -53,7 +102,15 @@ def span(self, name="span"):
         return self.start_span(name=name)
 
     def start_span(self, name="span"):
-        span = self._self_otel_tracer.start_span(name)
+        parent_ctx = context.get_current()
+        # If there is no current span in context, use the one provided to the OC Tracer at
+        # creation time
+        if trace.get_current_span(parent_ctx) is trace.INVALID_SPAN:
+            parent_ctx = set_oc_span_in_context(
+                self._self_oc_span_context, parent_ctx
+            )
+
+        span = self._self_otel_tracer.start_span(name, context=parent_ctx)
         shim_span = ShimSpan(
             BlankSpan(name=name, context_tracer=self),
             otel_span=span,
@@ -67,20 +124,27 @@ def start_span(self, name="span"):
         # equivalent to the below. This can cause context to leak but is equivalent.
         # pylint: disable=protected-access
         shim_span._self_token = context.attach(ctx)
+        # Also set it in OC's context, equivalent to
+        # https://github.com/census-instrumentation/opencensus-python/blob/2e08df591b507612b3968be8c2538dedbf8fab37/opencensus/trace/tracers/context_tracer.py#L94
+        execution_context.set_current_span(shim_span)
         return shim_span
 
     def end_span(self):
-        """Finishes the current span in the context and pops restores the context from before
-        the span was started.
+        """Finishes the current span in the context and restores the context from before the
+        span was started.
         """
         span = self.current_span()
         if not span:
             _logger.warning("No active span, cannot do end_span.")
             return
 
         span.finish()
+
         # pylint: disable=protected-access
         context.detach(span._self_token)
+        # Also reset the OC execution_context, equivalent to
+        # https://github.com/census-instrumentation/opencensus-python/blob/2e08df591b507612b3968be8c2538dedbf8fab37/opencensus/trace/tracers/context_tracer.py#L114-L117
+        execution_context.set_current_span(self.current_span())
 
     # pylint: disable=no-self-use
     def current_span(self):
diff --git a/shim/opentelemetry-opencensus-shim/tests/test_shim.py b/shim/opentelemetry-opencensus-shim/tests/test_shim.py
@@ -16,15 +16,21 @@
 import unittest
 from unittest.mock import patch
 
+from opencensus.trace import trace_options, tracestate
 from opencensus.trace.blank_span import BlankSpan as OcBlankSpan
 from opencensus.trace.link import Link as OcLink
 from opencensus.trace.span import SpanKind
+from opencensus.trace.span_context import SpanContext
 from opencensus.trace.tracer import Tracer as OcTracer
 from opencensus.trace.tracers.noop_tracer import NoopTracer as OcNoopTracer
 
+from opentelemetry import context, trace
 from opentelemetry.shim.opencensus import install_shim, uninstall_shim
 from opentelemetry.shim.opencensus._shim_span import ShimSpan
-from opentelemetry.shim.opencensus._shim_tracer import ShimTracer
+from opentelemetry.shim.opencensus._shim_tracer import (
+    ShimTracer,
+    set_oc_span_in_context,
+)
 
 
 class TestShim(unittest.TestCase):
@@ -126,3 +132,78 @@ def test_shim_span_contextmanager_calls_does_not_call_end(self):
                 pass
 
         spy_otel_span.end.assert_not_called()
+
+    def test_set_oc_span_in_context_no_span_id(self):
+        # This won't create a span ID and is the default behavior if you don't pass a context
+        # when creating the Tracer
+        ctx = set_oc_span_in_context(SpanContext(), context.get_current())
+        self.assertIs(trace.get_current_span(ctx), trace.INVALID_SPAN)
+
+    def test_set_oc_span_in_context_ids(self):
+        ctx = set_oc_span_in_context(
+            SpanContext(
+                trace_id="ace0216bab2b7ba249761dbb19c871b7",
+                span_id="1fead89ecf242225",
+            ),
+            context.get_current(),
+        )
+        span_ctx = trace.get_current_span(ctx).get_span_context()
+
+        self.assertEqual(
+            trace.format_trace_id(span_ctx.trace_id),
+            "ace0216bab2b7ba249761dbb19c871b7",
+        )
+        self.assertEqual(
+            trace.format_span_id(span_ctx.span_id), "1fead89ecf242225"
+        )
+
+    def test_set_oc_span_in_context_remote(self):
+        for is_from_remote in True, False:
+            ctx = set_oc_span_in_context(
+                SpanContext(
+                    trace_id="ace0216bab2b7ba249761dbb19c871b7",
+                    span_id="1fead89ecf242225",
+                    from_header=is_from_remote,
+                ),
+                context.get_current(),
+            )
+            span_ctx = trace.get_current_span(ctx).get_span_context()
+            self.assertEqual(span_ctx.is_remote, is_from_remote)
+
+    def test_set_oc_span_in_context_traceoptions(self):
+        for oc_trace_options, expect in [
+            # Not sampled
+            (
+                trace_options.TraceOptions("0"),
+                trace.TraceFlags(trace.TraceFlags.DEFAULT),
+            ),
+            # Sampled
+            (
+                trace_options.TraceOptions("1"),
+                trace.TraceFlags(trace.TraceFlags.SAMPLED),
+            ),
+        ]:
+            ctx = set_oc_span_in_context(
+                SpanContext(
+                    trace_id="ace0216bab2b7ba249761dbb19c871b7",
+                    span_id="1fead89ecf242225",
+                    trace_options=oc_trace_options,
+                ),
+                context.get_current(),
+            )
+            span_ctx = trace.get_current_span(ctx).get_span_context()
+            self.assertEqual(span_ctx.trace_flags, expect)
+
+    def test_set_oc_span_in_context_tracestate(self):
+        ctx = set_oc_span_in_context(
+            SpanContext(
+                trace_id="ace0216bab2b7ba249761dbb19c871b7",
+                span_id="1fead89ecf242225",
+                tracestate=tracestate.Tracestate({"hello": "tracestate"}),
+            ),
+            context.get_current(),
+        )
+        span_ctx = trace.get_current_span(ctx).get_span_context()
+        self.assertEqual(
+            span_ctx.trace_state, trace.TraceState([("hello", "tracestate")])
+        )
diff --git a/shim/opentelemetry-opencensus-shim/tests/test_shim_with_sdk.py b/shim/opentelemetry-opencensus-shim/tests/test_shim_with_sdk.py
