Adding attach/detach methods as per spec

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

@@ -14,6 +14,7 @@
 
 import logging
 import typing
+from functools import wraps
 from os import environ
 from sys import version_info
 
@@ -25,6 +26,44 @@
 _RUNTIME_CONTEXT = None  # type: typing.Optional[RuntimeContext]
 
 
+_F = typing.TypeVar("_F", bound=typing.Callable[..., typing.Any])
+
+
+def _load_runtime_context(func: _F) -> _F:
+    """Initializes the global RuntimeContext
+    """
+
+    @wraps(func)
+    def wrapper(
+        *args: typing.Tuple[typing.Any, typing.Any],
+        **kwargs: typing.Dict[typing.Any, typing.Any]
+    ) -> typing.Optional[typing.Any]:
+        global _RUNTIME_CONTEXT  # pylint: disable=global-statement
+        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
+            )  # type: str
+            try:
+                _RUNTIME_CONTEXT = next(
+                    iter_entry_points(
+                        "opentelemetry_context", configured_context
+                    )
+                ).load()()
+            except Exception:  # pylint: disable=broad-except
+                logger.error("Failed to load context: %s", configured_context)
+        return func(*args, **kwargs)  # type: ignore
+
+    return wrapper  # type:ignore
+
+
 def get_value(key: str, context: typing.Optional[Context] = None) -> "object":
     """To access the local state of a concern, the RuntimeContext API
     provides a function which takes a context and a key as input,
@@ -57,63 +96,38 @@ def set_value(
     return Context(new_values)
 
 
-def remove_value(
-    key: str, context: typing.Optional[Context] = None
-) -> Context:
-    """To remove a value, this method returns a new context with the key
-    cleared. Note that the removed value still remains present in the old
-    context.
-
-    Args:
-        key: The key of the entry to remove
-        context: The context to copy, if None, the current context is used
-    """
-    if context is None:
-        context = get_current()
-    new_values = context.copy()
-    new_values.pop(key, None)
-    return Context(new_values)
-
-
+@_load_runtime_context  # type: ignore
 def get_current() -> Context:
     """To access the context associated with program execution,
     the RuntimeContext API provides a function which takes no arguments
     and returns a RuntimeContext.
     """
+    return _RUNTIME_CONTEXT.get_current()  # type:ignore
 
-    global _RUNTIME_CONTEXT  # pylint: disable=global-statement
-    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
-        )  # type: str
-        try:
-            _RUNTIME_CONTEXT = next(
-                iter_entry_points("opentelemetry_context", configured_context)
-            ).load()()
-        except Exception:  # pylint: disable=broad-except
-            logger.error("Failed to load context: %s", configured_context)
 
-    return _RUNTIME_CONTEXT.get_current()  # type:ignore
+@_load_runtime_context  # type: ignore
+def attach(context: Context) -> object:
+    """Associates a Context with the caller's current execution unit. Returns
+    a token that can be used to restore the previous Context.
+
+    Args:
+        context: The Context to set as current.
+    """
+    return _RUNTIME_CONTEXT.attach(context)  # type:ignore
 
 
-def set_current(context: Context) -> Context:
-    """To associate a context with program execution, the Context
-    API provides a function which takes a Context.
+@_load_runtime_context  # type: ignore
+def detach(token: object) -> None:
+    """Resets the Context associated with the caller's current execution unit
+    to the value it had before attaching a specified Context.
 
     Args:
-        context: The context to use as current.
+        token: The Token that was returned by a previous call to attach a Context.
     """
-    old_context = get_current()
-    _RUNTIME_CONTEXT.set_current(context)  # type:ignore
-    return old_context
+    try:
+        _RUNTIME_CONTEXT.detach(token)  # type: ignore
+    except Exception:  # pylint: disable=broad-except
+        logger.error("Failed to detach context")
 
 
 def with_current_context(
@@ -127,10 +141,9 @@ def call_with_current_context(
         *args: "object", **kwargs: "object"
     ) -> "object":
         try:
-            backup = get_current()
-            set_current(caller_context)
+            token = attach(caller_context)
             return func(*args, **kwargs)
         finally:
-            set_current(backup)
+            detach(token)
 
     return call_with_current_context

opentelemetry-api/src/opentelemetry/context/context.py

@@ -29,8 +29,9 @@ class RuntimeContext(ABC):
     """
 
     @abstractmethod
-    def set_current(self, context: Context) -> None:
-        """ Sets the current `Context` object.
+    def attach(self, context: Context) -> object:
+        """ Sets the current `Context` object. Returns a
+        token that can be used to reset to the previous `Context`.
 
         Args:
             context: The Context to set.
@@ -40,5 +41,13 @@ def set_current(self, context: Context) -> None:
     def get_current(self) -> Context:
         """ Returns the current `Context` object. """
 
+    @abstractmethod
+    def detach(self, token: object) -> None:
+        """ Resets Context to a previous value
+
+        Args:
+            token: A reference to a previous Context.
+        """
+
 
 __all__ = ["Context", "RuntimeContext"]

opentelemetry-api/src/opentelemetry/context/contextvars_context.py

@@ -35,13 +35,17 @@ def __init__(self) -> None:
             self._CONTEXT_KEY, default=Context()
         )
 
-    def set_current(self, context: Context) -> None:
-        """See `opentelemetry.context.RuntimeContext.set_current`."""
-        self._current_context.set(context)
+    def attach(self, context: Context) -> object:
+        """See `opentelemetry.context.RuntimeContext.attach`."""
+        return self._current_context.set(context)
 
     def get_current(self) -> Context:
         """See `opentelemetry.context.RuntimeContext.get_current`."""
         return self._current_context.get()
 
+    def detach(self, token: object) -> None:
+        """See `opentelemetry.context.RuntimeContext.detach`."""
+        self._current_context.reset(token)  # type: ignore
+
 
 __all__ = ["ContextVarsRuntimeContext"]

opentelemetry-api/src/opentelemetry/context/threadlocal_context.py

@@ -23,14 +23,20 @@ class ThreadLocalRuntimeContext(RuntimeContext):
     implementation is available for usage with Python 3.4.
     """
 
+    class Token:
+        def __init__(self, context: Context) -> None:
+            self._context = context
+
     _CONTEXT_KEY = "current_context"
 
     def __init__(self) -> None:
         self._current_context = threading.local()
 
-    def set_current(self, context: Context) -> None:
-        """See `opentelemetry.context.RuntimeContext.set_current`."""
+    def attach(self, context: Context) -> object:
+        """See `opentelemetry.context.RuntimeContext.attach`."""
+        current = self.get_current()
         setattr(self._current_context, self._CONTEXT_KEY, context)
+        return self.Token(current)
 
     def get_current(self) -> Context:
         """See `opentelemetry.context.RuntimeContext.get_current`."""
@@ -43,5 +49,12 @@ def get_current(self) -> Context:
         )  # type: Context
         return context
 
+    def detach(self, token: object) -> None:
+        """See `opentelemetry.context.RuntimeContext.detach`."""
+        if not isinstance(token, self.Token):
+            raise ValueError("invalid token")
+        # pylint: disable=protected-access
+        setattr(self._current_context, self._CONTEXT_KEY, token._context)
+
 
 __all__ = ["ThreadLocalRuntimeContext"]

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

@@ -17,7 +17,7 @@
 import typing
 from contextlib import contextmanager
 
-from opentelemetry.context import get_value, set_current, set_value
+from opentelemetry.context import attach, get_value, set_value
 from opentelemetry.context.context import Context
 
 PRINTABLE = frozenset(
@@ -142,4 +142,4 @@ def distributed_context_from_context(
 def with_distributed_context(
     dctx: DistributedContext, context: typing.Optional[Context] = None
 ) -> None:
-    set_current(set_value(_DISTRIBUTED_CONTEXT_KEY, dctx, context=context))
+    attach(set_value(_DISTRIBUTED_CONTEXT_KEY, dctx, context=context))

opentelemetry-api/tests/context/test_context.py

@@ -19,12 +19,12 @@
 
 
 def do_work() -> None:
-    context.set_current(context.set_value("say", "bar"))
+    context.attach(context.set_value("say", "bar"))
 
 
 class TestContext(unittest.TestCase):
     def setUp(self):
-        context.set_current(Context())
+        context.attach(Context())
 
     def test_context(self):
         self.assertIsNone(context.get_value("say"))
@@ -55,11 +55,10 @@ def test_context_is_immutable(self):
             context.get_current()["test"] = "cant-change-immutable"
 
     def test_set_current(self):
-        context.set_current(context.set_value("a", "yyy"))
+        context.attach(context.set_value("a", "yyy"))
 
-        old_context = context.set_current(context.set_value("a", "zzz"))
-        self.assertEqual("yyy", context.get_value("a", context=old_context))
+        token = context.attach(context.set_value("a", "zzz"))
         self.assertEqual("zzz", context.get_value("a"))
 
-        context.set_current(old_context)
+        context.detach(token)
         self.assertEqual("yyy", context.get_value("a"))

opentelemetry-api/tests/context/test_contextvars_context.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import unittest
+from logging import ERROR
 from unittest.mock import patch
 
 from opentelemetry import context
@@ -27,18 +28,19 @@
 
 
 def do_work() -> None:
-    context.set_current(context.set_value("say", "bar"))
+    context.attach(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)
+        context.attach(self.previous_context)
 
     @patch(
-        "opentelemetry.context._RUNTIME_CONTEXT", ContextVarsRuntimeContext()  # type: ignore
+        "opentelemetry.context._RUNTIME_CONTEXT",  # type: ignore
+        ContextVarsRuntimeContext(),
     )
     def test_context(self):
         self.assertIsNone(context.get_value("say"))
@@ -56,7 +58,8 @@ def test_context(self):
         self.assertEqual(context.get_value("say", context=third), "bar")
 
     @patch(
-        "opentelemetry.context._RUNTIME_CONTEXT", ContextVarsRuntimeContext()  # type: ignore
+        "opentelemetry.context._RUNTIME_CONTEXT",  # type: ignore
+        ContextVarsRuntimeContext(),
     )
     def test_set_value(self):
         first = context.set_value("a", "yyy")
@@ -66,3 +69,19 @@ def test_set_value(self):
         self.assertEqual("zzz", context.get_value("a", context=second))
         self.assertEqual("---", context.get_value("a", context=third))
         self.assertEqual(None, context.get_value("a"))
+
+    @patch(
+        "opentelemetry.context._RUNTIME_CONTEXT",  # type: ignore
+        ContextVarsRuntimeContext(),
+    )
+    def test_set_current(self):
+        context.attach(context.set_value("a", "yyy"))
+
+        token = context.attach(context.set_value("a", "zzz"))
+        self.assertEqual("zzz", context.get_value("a"))
+
+        context.detach(token)
+        self.assertEqual("yyy", context.get_value("a"))
+
+        with self.assertLogs(level=ERROR):
+            context.detach("some garbage")

opentelemetry-api/tests/context/test_threadlocal_context.py

@@ -13,25 +13,27 @@
 # limitations under the License.
 
 import unittest
+from logging import ERROR
 from unittest.mock import patch
 
 from opentelemetry import context
 from opentelemetry.context.threadlocal_context import ThreadLocalRuntimeContext
 
 
 def do_work() -> None:
-    context.set_current(context.set_value("say", "bar"))
+    context.attach(context.set_value("say", "bar"))
 
 
 class TestThreadLocalContext(unittest.TestCase):
     def setUp(self):
         self.previous_context = context.get_current()
 
     def tearDown(self):
-        context.set_current(self.previous_context)
+        context.attach(self.previous_context)
 
     @patch(
-        "opentelemetry.context._RUNTIME_CONTEXT", ThreadLocalRuntimeContext()  # type: ignore
+        "opentelemetry.context._RUNTIME_CONTEXT",  # type: ignore
+        ThreadLocalRuntimeContext(),
     )
     def test_context(self):
         self.assertIsNone(context.get_value("say"))
@@ -49,7 +51,8 @@ def test_context(self):
         self.assertEqual(context.get_value("say", context=third), "bar")
 
     @patch(
-        "opentelemetry.context._RUNTIME_CONTEXT", ThreadLocalRuntimeContext()  # type: ignore
+        "opentelemetry.context._RUNTIME_CONTEXT",  # type: ignore
+        ThreadLocalRuntimeContext(),
     )
     def test_set_value(self):
         first = context.set_value("a", "yyy")
@@ -59,3 +62,19 @@ def test_set_value(self):
         self.assertEqual("zzz", context.get_value("a", context=second))
         self.assertEqual("---", context.get_value("a", context=third))
         self.assertEqual(None, context.get_value("a"))
+
+    @patch(
+        "opentelemetry.context._RUNTIME_CONTEXT",  # type: ignore
+        ThreadLocalRuntimeContext(),
+    )
+    def test_set_current(self):
+        context.attach(context.set_value("a", "yyy"))
+
+        token = context.attach(context.set_value("a", "zzz"))
+        self.assertEqual("zzz", context.get_value("a"))
+
+        context.detach(token)
+        self.assertEqual("yyy", context.get_value("a"))
+
+        with self.assertLogs(level=ERROR):
+            context.detach("some garbage")