Adding attach/detach methods as per spec (#429)

This change updates the Context API with the following:

- removes the remove_value method
- removes the set_current method
- adds attach and detach methods

Fixes #420

Co-authored-by: Chris Kleinknecht <[email protected]>

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,47 @@
 _RUNTIME_CONTEXT = None  # type: typing.Optional[RuntimeContext]
 
 
+_F = typing.TypeVar("_F", bound=typing.Callable[..., typing.Any])
+
+
+def _load_runtime_context(func: _F) -> _F:
+    """A decorator used to initialize the global RuntimeContext
+
+    Returns:
+        A wrapper of the decorated method.
+    """
+
+    @wraps(func)  # type: ignore
+    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,
@@ -33,6 +75,9 @@ def get_value(key: str, context: typing.Optional[Context] = None) -> "object":
     Args:
         key: The key of the value to retrieve.
         context: The context from which to retrieve the value, if None, the current context is used.
+
+    Returns:
+        The value associated with the key.
     """
     return context.get(key) if context is not None else get_current().get(key)
 
@@ -46,91 +91,55 @@ def set_value(
     which contains the new value.
 
     Args:
-        key: The key of the entry to set
-        value: The value of the entry to set
-        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[key] = value
-    return Context(new_values)
+        key: The key of the entry to set.
+        value: The value of the entry to set.
+        context: The context to copy, if None, the current context is used.
 
-
-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
+    Returns:
+        A new `Context` containing the value set.
     """
     if context is None:
         context = get_current()
     new_values = context.copy()
-    new_values.pop(key, None)
+    new_values[key] = value
     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.
-    """
-
-    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)
+    the Context API provides a function which takes no arguments
+    and returns a Context.
 
+    Returns:
+        The current `Context` object.
+    """
     return _RUNTIME_CONTEXT.get_current()  # 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 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 use as current.
-    """
-    old_context = get_current()
-    _RUNTIME_CONTEXT.set_current(context)  # type:ignore
-    return old_context
-
+        context: The Context to set as current.
 
-def with_current_context(
-    func: typing.Callable[..., "object"]
-) -> typing.Callable[..., "object"]:
-    """Capture the current context and apply it to the provided func."""
+    Returns:
+        A token that can be used with `detach` to reset the context.
+    """
+    return _RUNTIME_CONTEXT.attach(context)  # type:ignore
 
-    caller_context = get_current()
 
-    def call_with_current_context(
-        *args: "object", **kwargs: "object"
-    ) -> "object":
-        try:
-            backup = get_current()
-            set_current(caller_context)
-            return func(*args, **kwargs)
-        finally:
-            set_current(backup)
+@_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.
 
-    return call_with_current_context
+    Args:
+        token: The Token that was returned by a previous call to attach a Context.
+    """
+    try:
+        _RUNTIME_CONTEXT.detach(token)  # type: ignore
+    except Exception:  # pylint: disable=broad-except
+        logger.error("Failed to detach 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/base_context.py

@@ -0,0 +1,77 @@
+# 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 logging import ERROR
+
+from opentelemetry import context
+
+
+def do_work() -> None:
+    context.attach(context.set_value("say", "bar"))
+
+
+class ContextTestCases:
+    class BaseTest(unittest.TestCase):
+        def setUp(self) -> None:
+            self.previous_context = context.get_current()
+
+        def tearDown(self) -> None:
+            context.attach(self.previous_context)
+
+        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")
+
+        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"))
+
+        def test_attach(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")
+
+        def test_detach_out_of_order(self):
+            t1 = context.attach(context.set_value("c", 1))
+            self.assertEqual(context.get_current(), {"c": 1})
+            t2 = context.attach(context.set_value("c", 2))
+            self.assertEqual(context.get_current(), {"c": 2})
+            context.detach(t1)
+            self.assertEqual(context.get_current(), {})
+            context.detach(t2)
+            self.assertEqual(context.get_current(), {"c": 1})

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"))