Allow to override global tracer_provider after tracers creation

anton-ryzhov · anton-ryzhov · commit 06c561e0344c · 2020-12-04T12:37:22.000+01:00
diff --git a/opentelemetry-api/src/opentelemetry/trace/__init__.py b/opentelemetry-api/src/opentelemetry/trace/__init__.py
@@ -76,6 +76,7 @@
 
 import abc
 import enum
+import functools
 import typing
 from contextlib import contextmanager
 from logging import getLogger
@@ -409,9 +410,80 @@ def use_span(
         yield
 
 
+class ProxyTracer(Tracer):
+    """Proxies all calls to current TracerProvider
+    """
+
+    def __init__(
+        self, get_current_tracer: typing.Callable[[], "Tracer"],
+    ) -> None:
+        self._get_current_tracer = get_current_tracer
+
+    def start_span(
+        self,
+        name: str,
+        context: typing.Optional[Context] = None,
+        kind: SpanKind = SpanKind.INTERNAL,
+        attributes: types.Attributes = None,
+        links: typing.Sequence[Link] = (),
+        start_time: typing.Optional[int] = None,
+        record_exception: bool = True,
+        set_status_on_exception: bool = True,
+    ) -> "Span":
+        return self._get_current_tracer().start_span(
+            name=name,
+            context=context,
+            kind=kind,
+            attributes=attributes,
+            links=links,
+            start_time=start_time,
+            record_exception=record_exception,
+            set_status_on_exception=set_status_on_exception,
+        )
+
+    @contextmanager  # type: ignore
+    def start_as_current_span(
+        self,
+        name: str,
+        context: typing.Optional[Context] = None,
+        kind: SpanKind = SpanKind.INTERNAL,
+        attributes: types.Attributes = None,
+        links: typing.Sequence[Link] = (),
+        start_time: typing.Optional[int] = None,
+        record_exception: bool = True,
+        set_status_on_exception: bool = True,
+    ) -> typing.Iterator["Span"]:
+        with self._get_current_tracer().start_as_current_span(
+            name=name,
+            context=context,
+            kind=kind,
+            attributes=attributes,
+            links=links,
+            start_time=start_time,
+            record_exception=record_exception,
+            set_status_on_exception=set_status_on_exception,
+        ) as span:
+            yield span
+
+    @contextmanager  # type: ignore
+    def use_span(
+        self, span: "Span", end_on_exit: bool = False,
+    ) -> typing.Iterator[None]:
+        with self._get_current_tracer().use_span(
+            span=span, end_on_exit=end_on_exit,
+        ) as context:
+            yield context
+
+
 _TRACER_PROVIDER = None
 
 
+@functools.lru_cache  # type: ignore
+def _get_current_tracer(*args: typing.Any, **kwargs: typing.Any,) -> "Tracer":
+    tracer_provider = get_tracer_provider()
+    return tracer_provider.get_tracer(*args, **kwargs)  # type: ignore
+
+
 def get_tracer(
     instrumenting_module_name: str,
     instrumenting_library_version: str = "",
@@ -425,7 +497,14 @@ def get_tracer(
     If tracer_provider is ommited the current configured one is used.
     """
     if tracer_provider is None:
-        tracer_provider = get_tracer_provider()
+        return ProxyTracer(
+            functools.partial(
+                _get_current_tracer,
+                instrumenting_module_name=instrumenting_module_name,
+                instrumenting_library_version=instrumenting_library_version,
+            )
+        )
+
     return tracer_provider.get_tracer(
         instrumenting_module_name, instrumenting_library_version
     )
@@ -439,11 +518,8 @@ def set_tracer_provider(tracer_provider: TracerProvider) -> None:
     """
     global _TRACER_PROVIDER  # pylint: disable=global-statement
 
-    if _TRACER_PROVIDER is not None:
-        logger.warning("Overriding of current TracerProvider is not allowed")
-        return
-
     _TRACER_PROVIDER = tracer_provider
+    _get_current_tracer.cache_clear()  # pylint: disable=no-member
 
 
 def get_tracer_provider() -> TracerProvider:
diff --git a/opentelemetry-api/tests/trace/test_globals.py b/opentelemetry-api/tests/trace/test_globals.py
@@ -1,5 +1,5 @@
 import unittest
-from unittest.mock import patch
+from unittest.mock import MagicMock, patch
 
 from opentelemetry import context, trace
 
@@ -13,13 +13,77 @@ def tearDown(self) -> None:
         self._patcher.stop()
 
     def test_get_tracer(self):
-        """trace.get_tracer should proxy to the global tracer provider."""
-        trace.get_tracer("foo", "var")
-        self._mock_tracer_provider.get_tracer.assert_called_with("foo", "var")
+        """trace.get_tracer should create a proxy to the global tracer provider."""
+        tracer = trace.get_tracer("foo", "var")
+        self._mock_tracer_provider.get_tracer.assert_not_called()
+        self.assertIsInstance(tracer, trace.ProxyTracer)
+
+        tracer.start_span("one")
+        tracer.start_span("two")
+        self._mock_tracer_provider.get_tracer.assert_called_once_with(
+            instrumenting_module_name="foo",
+            instrumenting_library_version="var",
+        )
+
         mock_provider = unittest.mock.Mock()
         trace.get_tracer("foo", "var", mock_provider)
         mock_provider.get_tracer.assert_called_with("foo", "var")
 
+    def test_set_tracer_provider(self):
+        """trace.get_tracer should update global tracer provider."""
+        self.assertIs(
+            trace._TRACER_PROVIDER,  # pylint: disable=protected-access
+            self._mock_tracer_provider,
+        )
+
+        tracer_provider1 = trace.DefaultTracerProvider()
+        trace.set_tracer_provider(tracer_provider1)
+        self.assertIs(
+            trace._TRACER_PROVIDER,  # pylint: disable=protected-access
+            tracer_provider1,
+        )
+
+        tracer_provider2 = trace.DefaultTracerProvider()
+        trace.set_tracer_provider(tracer_provider2)
+        self.assertIs(
+            trace._TRACER_PROVIDER,  # pylint: disable=protected-access
+            tracer_provider2,
+        )
+
+    @patch("opentelemetry.trace._load_trace_provider")
+    def test_get_tracer_provider(self, load_trace_provider_mock):
+        """trace.get_tracer should get or create a global tracer provider."""
+        load_trace_provider_mock.assert_not_called()
+
+        tracer_provider = trace.get_tracer_provider()
+        self.assertIs(
+            trace._TRACER_PROVIDER,  # pylint: disable=protected-access
+            tracer_provider,
+        )
+        load_trace_provider_mock.assert_not_called()
+
+        trace._TRACER_PROVIDER = None  # pylint: disable=protected-access
+        tracer_provider1 = trace.get_tracer_provider()
+        self.assertIsNotNone(
+            trace._TRACER_PROVIDER  # pylint: disable=protected-access
+        )
+        self.assertIs(
+            trace._TRACER_PROVIDER,  # pylint: disable=protected-access
+            tracer_provider1,
+        )
+        load_trace_provider_mock.assert_called_once_with("tracer_provider")
+
+        tracer_provider2 = trace.get_tracer_provider()
+        self.assertIsNotNone(
+            trace._TRACER_PROVIDER  # pylint: disable=protected-access
+        )
+        self.assertIs(
+            trace._TRACER_PROVIDER,  # pylint: disable=protected-access
+            tracer_provider2,
+        )
+        self.assertIs(tracer_provider1, tracer_provider2)
+        load_trace_provider_mock.assert_called_once_with("tracer_provider")
+
 
 class TestTracer(unittest.TestCase):
     def setUp(self):
@@ -38,3 +102,109 @@ def test_get_current_span(self):
         finally:
             context.detach(token)
         self.assertEqual(trace.get_current_span(), trace.INVALID_SPAN)
+
+
+class TestProxyTracer(unittest.TestCase):
+    def setUp(self):
+        self._patcher = patch("opentelemetry.trace._TRACER_PROVIDER")
+        self._patcher.start()
+
+        self.proxy_tracer = trace.get_tracer("foo", "var")
+        self.inner_tracer = MagicMock(wraps=trace.DefaultTracer())
+
+        tracer_provider = MagicMock(wraps=trace.DefaultTracerProvider())
+        tracer_provider.get_tracer.return_value = self.inner_tracer
+        trace.set_tracer_provider(tracer_provider)
+
+    def tearDown(self):
+        self._patcher.stop()
+
+    def test_start_span(self):
+        """ProxyTracer should call `start_span` on a real `Tracer`
+        """
+        self.inner_tracer.start_span.assert_not_called()
+
+        span = self.proxy_tracer.start_span("span1")
+        self.assertIs(span, trace.INVALID_SPAN)
+        self.inner_tracer.start_span.assert_called_once_with(
+            name="span1",
+            context=None,
+            kind=trace.SpanKind.INTERNAL,
+            attributes=None,
+            links=(),
+            start_time=None,
+            record_exception=True,
+            set_status_on_exception=True,
+        )
+
+    def test_start_as_current_span(self):
+        """ProxyTracer should call `start_as_current_span` on a real `Tracer`
+        """
+        self.inner_tracer.start_as_current_span.assert_not_called()
+
+        with self.proxy_tracer.start_as_current_span("span1") as span:
+            self.assertIs(span, trace.INVALID_SPAN)
+            self.inner_tracer.start_as_current_span.assert_called_once_with(
+                name="span1",
+                context=None,
+                kind=trace.SpanKind.INTERNAL,
+                attributes=None,
+                links=(),
+                start_time=None,
+                record_exception=True,
+                set_status_on_exception=True,
+            )
+
+        self.inner_tracer.start_as_current_span.reset_mock()
+
+        @self.proxy_tracer.start_as_current_span("span1")
+        def func(arg, kwarg):
+            self.assertEqual(arg, "argval")
+            self.assertEqual(kwarg, "kwargval")
+            return "retval"
+
+        self.inner_tracer.start_as_current_span.assert_not_called()
+
+        result = func("argval", kwarg="kwargval")
+        self.assertEqual(result, "retval")
+
+        self.inner_tracer.start_as_current_span.assert_called_once_with(
+            name="span1",
+            context=None,
+            kind=trace.SpanKind.INTERNAL,
+            attributes=None,
+            links=(),
+            start_time=None,
+            record_exception=True,
+            set_status_on_exception=True,
+        )
+
+    def test_use_span(self):
+        """ProxyTracer should call `use_span` on a real `Tracer`
+        """
+        self.inner_tracer.use_span.assert_not_called()
+
+        span = self.proxy_tracer.start_span("span1")
+
+        with self.proxy_tracer.use_span(span) as current_context:
+            self.assertIsNone(current_context)
+            self.inner_tracer.use_span.assert_called_once_with(
+                span=span, end_on_exit=False,
+            )
+
+        self.inner_tracer.use_span.reset_mock()
+
+        @self.proxy_tracer.use_span(span)
+        def func(arg, kwarg):
+            self.assertEqual(arg, "argval")
+            self.assertEqual(kwarg, "kwargval")
+            return "retval"
+
+        self.inner_tracer.use_span.assert_not_called()
+
+        result = func("argval", kwarg="kwargval")
+        self.assertEqual(result, "retval")
+
+        self.inner_tracer.use_span.assert_called_once_with(
+            span=span, end_on_exit=False,
+        )
diff --git a/opentelemetry-sdk/tests/trace/test_globals.py b/opentelemetry-sdk/tests/trace/test_globals.py
@@ -1,25 +1,17 @@
 # type:ignore
 import unittest
-from logging import WARNING
 
 from opentelemetry import trace
 from opentelemetry.sdk.trace import TracerProvider  # type:ignore
 
 
 class TestGlobals(unittest.TestCase):
-    def test_tracer_provider_override_warning(self):
-        """trace.set_tracer_provider should throw a warning when overridden"""
-        trace.set_tracer_provider(TracerProvider())
-        tracer_provider = trace.get_tracer_provider()
-        with self.assertLogs(level=WARNING) as test:
-            trace.set_tracer_provider(TracerProvider())
-            self.assertEqual(
-                test.output,
-                [
-                    (
-                        "WARNING:opentelemetry.trace:Overriding of current "
-                        "TracerProvider is not allowed"
-                    )
-                ],
-            )
+    def test_tracer_provider_override(self):
+        """trace.set_tracer_provider should override global tracer"""
+        tracer_provider = TracerProvider()
+        trace.set_tracer_provider(tracer_provider)
         self.assertIs(tracer_provider, trace.get_tracer_provider())
+
+        new_tracer_provider = TracerProvider()
+        trace.set_tracer_provider(new_tracer_provider)
+        self.assertIs(new_tracer_provider, trace.get_tracer_provider())
