Moving contextvars and threadlocal context implementations to the API (open-telemetry#419)

codeboten · c24t · web-flow · commit d824f19192eb · 2020-02-18T10:31:47.000-08:00
While keeping the flexibility of the interface to support additional implementations of the RuntimeContext, moving the ContextVarsRuntimeContext and ThreadLocalRuntimeContext into the API ensures a useful context implementation for users without requiring additional configuration. This changes checks the version number to determine which implementation to use, unless it is overridden by the OPENTELEMETRY_CONTEXT environment variable. Removing the DefaultRuntimeContext as part of this change.

Signed-off-by: Alex Boten &lt;aboten@lightstep.com&gt;
Co-authored-by: Chris Kleinknecht &lt;libc@google.com&gt;
diff --git a/opentelemetry-api/setup.py b/opentelemetry-api/setup.py
@@ -44,7 +44,10 @@
     include_package_data=True,
     long_description=open("README.rst").read(),
     long_description_content_type="text/x-rst",
-    install_requires=["typing; python_version<'3.5'"],
+    install_requires=[
+        "typing; python_version<'3.5'",
+        "aiocontextvars; python_version<'3.7' and python_version>='3.5'",
+    ],
     extras_require={},
     license="Apache-2.0",
     package_dir={"": "src"},
@@ -58,9 +61,12 @@
     zip_safe=False,
     entry_points={
         "opentelemetry_context": [
-            "default_context = "
-            "opentelemetry.context.default_context:"
-            "DefaultRuntimeContext",
+            "contextvars_context = "
+            "opentelemetry.context.contextvars_context:"
+            "ContextVarsRuntimeContext",
+            "threadlocal_context = "
+            "opentelemetry.context.threadlocal_context:"
+            "ThreadLocalRuntimeContext",
         ]
     },
 )
diff --git a/opentelemetry-api/src/opentelemetry/context/__init__.py b/opentelemetry-api/src/opentelemetry/context/__init__.py
@@ -15,6 +15,7 @@
 import logging
 import typing
 from os import environ
+from sys import version_info
 
 from pkg_resources import iter_entry_points
 
@@ -84,9 +85,14 @@ def get_current() -> Context:
     if _RUNTIME_CONTEXT is None:
         # FIXME use a better implementation of a configuration manager to avoid having
         # to get configuration values straight from environment variables
+        if version_info < (3, 5):
+            # contextvars are not supported in 3.4, use thread-local storage
+            default_context = "threadlocal_context"
+        else:
+            default_context = "contextvars_context"
 
         configured_context = environ.get(
-            "OPENTELEMETRY_CONTEXT", "default_context"
+            "OPENTELEMETRY_CONTEXT", default_context
         )  # type: str
         try:
             _RUNTIME_CONTEXT = next(
diff --git a/opentelemetry-api/src/opentelemetry/context/aiocontextvarsfix.py b/opentelemetry-api/src/opentelemetry/context/aiocontextvarsfix.py
@@ -55,7 +55,7 @@ def _get_event_loop():
     asyncio._set_running_loop = asyncio.events._set_running_loop
 
 # noinspection PyUnresolvedReferences
-import aiocontextvars  # pylint: disable=unused-import,wrong-import-position # noqa # isort:skip
+import aiocontextvars  # pylint: disable=import-error,unused-import,wrong-import-position # noqa # isort:skip
 
 
 def _run_coroutine_threadsafe(coro, loop):
diff --git a/opentelemetry-api/src/opentelemetry/context/contextvars_context.py b/opentelemetry-api/src/opentelemetry/context/contextvars_context.py
@@ -14,14 +14,13 @@
 from contextvars import ContextVar
 from sys import version_info
 
-from opentelemetry.context import Context
-from opentelemetry.context.context import RuntimeContext
+from opentelemetry.context.context import Context, RuntimeContext
 
 if (3, 5, 3) <= version_info < (3, 7):
-    import aiocontextvars  # type: ignore # pylint:disable=unused-import
+    import aiocontextvars  # type: ignore # pylint:disable=unused-import,import-error
 
 elif (3, 4) < version_info <= (3, 5, 2):
-    import opentelemetry.sdk.context.aiocontextvarsfix  # pylint:disable=unused-import
+    import opentelemetry.context.aiocontextvarsfix  # pylint:disable=unused-import
 
 
 class ContextVarsRuntimeContext(RuntimeContext):
diff --git a/opentelemetry-api/src/opentelemetry/context/default_context.py b/opentelemetry-api/src/opentelemetry/context/default_context.py
diff --git a/opentelemetry-api/src/opentelemetry/context/threadlocal_context.py b/opentelemetry-api/src/opentelemetry/context/threadlocal_context.py
@@ -14,7 +14,7 @@
 
 import threading
 
-from opentelemetry.context import Context, RuntimeContext
+from opentelemetry.context.context import Context, RuntimeContext
 
 
 class ThreadLocalRuntimeContext(RuntimeContext):
@@ -38,7 +38,10 @@ def get_current(self) -> Context:
             setattr(
                 self._current_context, self._CONTEXT_KEY, Context(),
             )
-        return getattr(self._current_context, self._CONTEXT_KEY)
+        context = getattr(
+            self._current_context, self._CONTEXT_KEY
+        )  # type: Context
+        return context
 
 
 __all__ = ["ThreadLocalRuntimeContext"]
diff --git a/opentelemetry-api/tests/context/test_contextvars_context.py b/opentelemetry-api/tests/context/test_contextvars_context.py
@@ -0,0 +1,68 @@
+# Copyright 2020, OpenTelemetry Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+from unittest.mock import patch
+
+from opentelemetry import context
+
+try:
+    import contextvars  # pylint: disable=unused-import
+    from opentelemetry.context.contextvars_context import (
+        ContextVarsRuntimeContext,
+    )
+except ImportError:
+    raise unittest.SkipTest("contextvars not available")
+
+
+def do_work() -> None:
+    context.set_current(context.set_value("say", "bar"))
+
+
+class TestContextVarsContext(unittest.TestCase):
+    def setUp(self):
+        self.previous_context = context.get_current()
+
+    def tearDown(self):
+        context.set_current(self.previous_context)
+
+    @patch(
+        "opentelemetry.context._RUNTIME_CONTEXT", ContextVarsRuntimeContext()  # type: ignore
+    )
+    def test_context(self):
+        self.assertIsNone(context.get_value("say"))
+        empty = context.get_current()
+        second = context.set_value("say", "foo")
+
+        self.assertEqual(context.get_value("say", context=second), "foo")
+
+        do_work()
+        self.assertEqual(context.get_value("say"), "bar")
+        third = context.get_current()
+
+        self.assertIsNone(context.get_value("say", context=empty))
+        self.assertEqual(context.get_value("say", context=second), "foo")
+        self.assertEqual(context.get_value("say", context=third), "bar")
+
+    @patch(
+        "opentelemetry.context._RUNTIME_CONTEXT", ContextVarsRuntimeContext()  # type: ignore
+    )
+    def test_set_value(self):
+        first = context.set_value("a", "yyy")
+        second = context.set_value("a", "zzz")
+        third = context.set_value("a", "---", first)
+        self.assertEqual("yyy", context.get_value("a", context=first))
+        self.assertEqual("zzz", context.get_value("a", context=second))
+        self.assertEqual("---", context.get_value("a", context=third))
+        self.assertEqual(None, context.get_value("a"))
diff --git a/opentelemetry-api/tests/context/test_threadlocal_context.py b/opentelemetry-api/tests/context/test_threadlocal_context.py
@@ -16,9 +16,7 @@
 from unittest.mock import patch
 
 from opentelemetry import context
-from opentelemetry.sdk.context.threadlocal_context import (
-    ThreadLocalRuntimeContext,
-)
+from opentelemetry.context.threadlocal_context import ThreadLocalRuntimeContext
 
 
 def do_work() -> None:
@@ -33,7 +31,7 @@ def tearDown(self):
         context.set_current(self.previous_context)
 
     @patch(
-        "opentelemetry.context._RUNTIME_CONTEXT", ThreadLocalRuntimeContext()
+        "opentelemetry.context._RUNTIME_CONTEXT", ThreadLocalRuntimeContext()  # type: ignore
     )
     def test_context(self):
         self.assertIsNone(context.get_value("say"))
@@ -51,7 +49,7 @@ def test_context(self):
         self.assertEqual(context.get_value("say", context=third), "bar")
 
     @patch(
-        "opentelemetry.context._RUNTIME_CONTEXT", ThreadLocalRuntimeContext()
+        "opentelemetry.context._RUNTIME_CONTEXT", ThreadLocalRuntimeContext()  # type: ignore
     )
     def test_set_value(self):
         first = context.set_value("a", "yyy")
diff --git a/opentelemetry-sdk/setup.py b/opentelemetry-sdk/setup.py
@@ -44,7 +44,7 @@
     include_package_data=True,
     long_description=open("README.rst").read(),
     long_description_content_type="text/x-rst",
-    install_requires=["opentelemetry-api==0.4.dev0", "aiocontextvars"],
+    install_requires=["opentelemetry-api==0.4.dev0"],
     extras_require={},
     license="Apache-2.0",
     package_dir={"": "src"},
@@ -56,14 +56,4 @@
         "/tree/master/opentelemetry-sdk"
     ),
     zip_safe=False,
-    entry_points={
-        "opentelemetry_context": [
-            "contextvars_context = "
-            "opentelemetry.sdk.context.contextvars_context:"
-            "ContextVarsRuntimeContext",
-            "threadlocal_context = "
-            "opentelemetry.sdk.context.threadlocal_context:"
-            "ThreadLocalRuntimeContext",
-        ]
-    },
 )
diff --git a/opentelemetry-sdk/tests/context/test_asyncio.py b/opentelemetry-sdk/tests/context/test_asyncio.py
@@ -25,7 +25,7 @@
 
 try:
     import contextvars  # pylint: disable=unused-import
-    from opentelemetry.sdk.context.contextvars_context import (
+    from opentelemetry.context.contextvars_context import (
         ContextVarsRuntimeContext,
     )
 except ImportError:
@@ -53,10 +53,6 @@ def stop_loop_when(loop, cond_func, timeout=5.0):
     loop.call_later(0.1, stop_loop_when, loop, cond_func, timeout)
 
 
-def do_work() -> None:
-    context.set_current(context.set_value("say", "bar"))
-
-
 class TestAsyncio(unittest.TestCase):
     @asyncio.coroutine
     def task(self, name):
@@ -114,41 +110,3 @@ def test_with_asyncio(self):
             if span is expected_parent:
                 continue
             self.assertEqual(span.parent, expected_parent)
-
-
-class TestContextVarsContext(unittest.TestCase):
-    def setUp(self):
-        self.previous_context = context.get_current()
-
-    def tearDown(self):
-        context.set_current(self.previous_context)
-
-    @patch(
-        "opentelemetry.context._RUNTIME_CONTEXT", ContextVarsRuntimeContext()
-    )
-    def test_context(self):
-        self.assertIsNone(context.get_value("say"))
-        empty = context.get_current()
-        second = context.set_value("say", "foo")
-
-        self.assertEqual(context.get_value("say", context=second), "foo")
-
-        do_work()
-        self.assertEqual(context.get_value("say"), "bar")
-        third = context.get_current()
-
-        self.assertIsNone(context.get_value("say", context=empty))
-        self.assertEqual(context.get_value("say", context=second), "foo")
-        self.assertEqual(context.get_value("say", context=third), "bar")
-
-    @patch(
-        "opentelemetry.context._RUNTIME_CONTEXT", ContextVarsRuntimeContext()
-    )
-    def test_set_value(self):
-        first = context.set_value("a", "yyy")
-        second = context.set_value("a", "zzz")
-        third = context.set_value("a", "---", first)
-        self.assertEqual("yyy", context.get_value("a", context=first))
-        self.assertEqual("zzz", context.get_value("a", context=second))
-        self.assertEqual("---", context.get_value("a", context=third))
-        self.assertEqual(None, context.get_value("a"))
diff --git a/opentelemetry-sdk/tests/context/test_threads.py b/opentelemetry-sdk/tests/context/test_threads.py
