Skip to content

Commit 5b2e693

Browse files
codebotenc24t
codeboten
and
c24t
authored
Adding attach/detach methods as per spec (open-telemetry#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 open-telemetry#420 Co-authored-by: Chris Kleinknecht <libc@google.com>
1 parent a4e7a9a commit 5b2e693

File tree

12 files changed

+226
-174
lines changed

12 files changed

+226
-174
lines changed

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

Lines changed: 77 additions & 68 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@
1414

1515
import logging
1616
import typing
17+
from functools import wraps
1718
from os import environ
1819
from sys import version_info
1920

@@ -25,6 +26,47 @@
2526
_RUNTIME_CONTEXT = None # type: typing.Optional[RuntimeContext]
2627

2728

29+
_F = typing.TypeVar("_F", bound=typing.Callable[..., typing.Any])
30+
31+
32+
def _load_runtime_context(func: _F) -> _F:
33+
"""A decorator used to initialize the global RuntimeContext
34+
35+
Returns:
36+
A wrapper of the decorated method.
37+
"""
38+
39+
@wraps(func) # type: ignore
40+
def wrapper(
41+
*args: typing.Tuple[typing.Any, typing.Any],
42+
**kwargs: typing.Dict[typing.Any, typing.Any]
43+
) -> typing.Optional[typing.Any]:
44+
global _RUNTIME_CONTEXT # pylint: disable=global-statement
45+
if _RUNTIME_CONTEXT is None:
46+
# FIXME use a better implementation of a configuration manager to avoid having
47+
# to get configuration values straight from environment variables
48+
if version_info < (3, 5):
49+
# contextvars are not supported in 3.4, use thread-local storage
50+
default_context = "threadlocal_context"
51+
else:
52+
default_context = "contextvars_context"
53+
54+
configured_context = environ.get(
55+
"OPENTELEMETRY_CONTEXT", default_context
56+
) # type: str
57+
try:
58+
_RUNTIME_CONTEXT = next(
59+
iter_entry_points(
60+
"opentelemetry_context", configured_context
61+
)
62+
).load()()
63+
except Exception: # pylint: disable=broad-except
64+
logger.error("Failed to load context: %s", configured_context)
65+
return func(*args, **kwargs) # type: ignore
66+
67+
return wrapper # type:ignore
68+
69+
2870
def get_value(key: str, context: typing.Optional[Context] = None) -> "object":
2971
"""To access the local state of a concern, the RuntimeContext API
3072
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":
3375
Args:
3476
key: The key of the value to retrieve.
3577
context: The context from which to retrieve the value, if None, the current context is used.
78+
79+
Returns:
80+
The value associated with the key.
3681
"""
3782
return context.get(key) if context is not None else get_current().get(key)
3883

@@ -46,91 +91,55 @@ def set_value(
4691
which contains the new value.
4792
4893
Args:
49-
key: The key of the entry to set
50-
value: The value of the entry to set
51-
context: The context to copy, if None, the current context is used
52-
"""
53-
if context is None:
54-
context = get_current()
55-
new_values = context.copy()
56-
new_values[key] = value
57-
return Context(new_values)
94+
key: The key of the entry to set.
95+
value: The value of the entry to set.
96+
context: The context to copy, if None, the current context is used.
5897
59-
60-
def remove_value(
61-
key: str, context: typing.Optional[Context] = None
62-
) -> Context:
63-
"""To remove a value, this method returns a new context with the key
64-
cleared. Note that the removed value still remains present in the old
65-
context.
66-
67-
Args:
68-
key: The key of the entry to remove
69-
context: The context to copy, if None, the current context is used
98+
Returns:
99+
A new `Context` containing the value set.
70100
"""
71101
if context is None:
72102
context = get_current()
73103
new_values = context.copy()
74-
new_values.pop(key, None)
104+
new_values[key] = value
75105
return Context(new_values)
76106

77107

108+
@_load_runtime_context # type: ignore
78109
def get_current() -> Context:
79110
"""To access the context associated with program execution,
80-
the RuntimeContext API provides a function which takes no arguments
81-
and returns a RuntimeContext.
82-
"""
83-
84-
global _RUNTIME_CONTEXT # pylint: disable=global-statement
85-
if _RUNTIME_CONTEXT is None:
86-
# FIXME use a better implementation of a configuration manager to avoid having
87-
# to get configuration values straight from environment variables
88-
if version_info < (3, 5):
89-
# contextvars are not supported in 3.4, use thread-local storage
90-
default_context = "threadlocal_context"
91-
else:
92-
default_context = "contextvars_context"
93-
94-
configured_context = environ.get(
95-
"OPENTELEMETRY_CONTEXT", default_context
96-
) # type: str
97-
try:
98-
_RUNTIME_CONTEXT = next(
99-
iter_entry_points("opentelemetry_context", configured_context)
100-
).load()()
101-
except Exception: # pylint: disable=broad-except
102-
logger.error("Failed to load context: %s", configured_context)
111+
the Context API provides a function which takes no arguments
112+
and returns a Context.
103113
114+
Returns:
115+
The current `Context` object.
116+
"""
104117
return _RUNTIME_CONTEXT.get_current() # type:ignore
105118

106119

107-
def set_current(context: Context) -> Context:
108-
"""To associate a context with program execution, the Context
109-
API provides a function which takes a Context.
120+
@_load_runtime_context # type: ignore
121+
def attach(context: Context) -> object:
122+
"""Associates a Context with the caller's current execution unit. Returns
123+
a token that can be used to restore the previous Context.
110124
111125
Args:
112-
context: The context to use as current.
113-
"""
114-
old_context = get_current()
115-
_RUNTIME_CONTEXT.set_current(context) # type:ignore
116-
return old_context
117-
126+
context: The Context to set as current.
118127
119-
def with_current_context(
120-
func: typing.Callable[..., "object"]
121-
) -> typing.Callable[..., "object"]:
122-
"""Capture the current context and apply it to the provided func."""
128+
Returns:
129+
A token that can be used with `detach` to reset the context.
130+
"""
131+
return _RUNTIME_CONTEXT.attach(context) # type:ignore
123132

124-
caller_context = get_current()
125133

126-
def call_with_current_context(
127-
*args: "object", **kwargs: "object"
128-
) -> "object":
129-
try:
130-
backup = get_current()
131-
set_current(caller_context)
132-
return func(*args, **kwargs)
133-
finally:
134-
set_current(backup)
134+
@_load_runtime_context # type: ignore
135+
def detach(token: object) -> None:
136+
"""Resets the Context associated with the caller's current execution unit
137+
to the value it had before attaching a specified Context.
135138
136-
return call_with_current_context
139+
Args:
140+
token: The Token that was returned by a previous call to attach a Context.
141+
"""
142+
try:
143+
_RUNTIME_CONTEXT.detach(token) # type: ignore
144+
except Exception: # pylint: disable=broad-except
145+
logger.error("Failed to detach context")

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

Lines changed: 11 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -29,8 +29,9 @@ class RuntimeContext(ABC):
2929
"""
3030

3131
@abstractmethod
32-
def set_current(self, context: Context) -> None:
33-
""" Sets the current `Context` object.
32+
def attach(self, context: Context) -> object:
33+
""" Sets the current `Context` object. Returns a
34+
token that can be used to reset to the previous `Context`.
3435
3536
Args:
3637
context: The Context to set.
@@ -40,5 +41,13 @@ def set_current(self, context: Context) -> None:
4041
def get_current(self) -> Context:
4142
""" Returns the current `Context` object. """
4243

44+
@abstractmethod
45+
def detach(self, token: object) -> None:
46+
""" Resets Context to a previous value
47+
48+
Args:
49+
token: A reference to a previous Context.
50+
"""
51+
4352

4453
__all__ = ["Context", "RuntimeContext"]

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

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -35,13 +35,17 @@ def __init__(self) -> None:
3535
self._CONTEXT_KEY, default=Context()
3636
)
3737

38-
def set_current(self, context: Context) -> None:
39-
"""See `opentelemetry.context.RuntimeContext.set_current`."""
40-
self._current_context.set(context)
38+
def attach(self, context: Context) -> object:
39+
"""See `opentelemetry.context.RuntimeContext.attach`."""
40+
return self._current_context.set(context)
4141

4242
def get_current(self) -> Context:
4343
"""See `opentelemetry.context.RuntimeContext.get_current`."""
4444
return self._current_context.get()
4545

46+
def detach(self, token: object) -> None:
47+
"""See `opentelemetry.context.RuntimeContext.detach`."""
48+
self._current_context.reset(token) # type: ignore
49+
4650

4751
__all__ = ["ContextVarsRuntimeContext"]

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

Lines changed: 15 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -23,14 +23,20 @@ class ThreadLocalRuntimeContext(RuntimeContext):
2323
implementation is available for usage with Python 3.4.
2424
"""
2525

26+
class Token:
27+
def __init__(self, context: Context) -> None:
28+
self._context = context
29+
2630
_CONTEXT_KEY = "current_context"
2731

2832
def __init__(self) -> None:
2933
self._current_context = threading.local()
3034

31-
def set_current(self, context: Context) -> None:
32-
"""See `opentelemetry.context.RuntimeContext.set_current`."""
35+
def attach(self, context: Context) -> object:
36+
"""See `opentelemetry.context.RuntimeContext.attach`."""
37+
current = self.get_current()
3338
setattr(self._current_context, self._CONTEXT_KEY, context)
39+
return self.Token(current)
3440

3541
def get_current(self) -> Context:
3642
"""See `opentelemetry.context.RuntimeContext.get_current`."""
@@ -43,5 +49,12 @@ def get_current(self) -> Context:
4349
) # type: Context
4450
return context
4551

52+
def detach(self, token: object) -> None:
53+
"""See `opentelemetry.context.RuntimeContext.detach`."""
54+
if not isinstance(token, self.Token):
55+
raise ValueError("invalid token")
56+
# pylint: disable=protected-access
57+
setattr(self._current_context, self._CONTEXT_KEY, token._context)
58+
4659

4760
__all__ = ["ThreadLocalRuntimeContext"]

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@
1717
import typing
1818
from contextlib import contextmanager
1919

20-
from opentelemetry.context import get_value, set_current, set_value
20+
from opentelemetry.context import attach, get_value, set_value
2121
from opentelemetry.context.context import Context
2222

2323
PRINTABLE = frozenset(
@@ -142,4 +142,4 @@ def distributed_context_from_context(
142142
def with_distributed_context(
143143
dctx: DistributedContext, context: typing.Optional[Context] = None
144144
) -> None:
145-
set_current(set_value(_DISTRIBUTED_CONTEXT_KEY, dctx, context=context))
145+
attach(set_value(_DISTRIBUTED_CONTEXT_KEY, dctx, context=context))

‎opentelemetry-api/tests/context/base_context.py

Lines changed: 77 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,77 @@
1+
# Copyright 2020, OpenTelemetry Authors
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
import unittest
16+
from logging import ERROR
17+
18+
from opentelemetry import context
19+
20+
21+
def do_work() -> None:
22+
context.attach(context.set_value("say", "bar"))
23+
24+
25+
class ContextTestCases:
26+
class BaseTest(unittest.TestCase):
27+
def setUp(self) -> None:
28+
self.previous_context = context.get_current()
29+
30+
def tearDown(self) -> None:
31+
context.attach(self.previous_context)
32+
33+
def test_context(self):
34+
self.assertIsNone(context.get_value("say"))
35+
empty = context.get_current()
36+
second = context.set_value("say", "foo")
37+
38+
self.assertEqual(context.get_value("say", context=second), "foo")
39+
40+
do_work()
41+
self.assertEqual(context.get_value("say"), "bar")
42+
third = context.get_current()
43+
44+
self.assertIsNone(context.get_value("say", context=empty))
45+
self.assertEqual(context.get_value("say", context=second), "foo")
46+
self.assertEqual(context.get_value("say", context=third), "bar")
47+
48+
def test_set_value(self):
49+
first = context.set_value("a", "yyy")
50+
second = context.set_value("a", "zzz")
51+
third = context.set_value("a", "---", first)
52+
self.assertEqual("yyy", context.get_value("a", context=first))
53+
self.assertEqual("zzz", context.get_value("a", context=second))
54+
self.assertEqual("---", context.get_value("a", context=third))
55+
self.assertEqual(None, context.get_value("a"))
56+
57+
def test_attach(self):
58+
context.attach(context.set_value("a", "yyy"))
59+
60+
token = context.attach(context.set_value("a", "zzz"))
61+
self.assertEqual("zzz", context.get_value("a"))
62+
63+
context.detach(token)
64+
self.assertEqual("yyy", context.get_value("a"))
65+
66+
with self.assertLogs(level=ERROR):
67+
context.detach("some garbage")
68+
69+
def test_detach_out_of_order(self):
70+
t1 = context.attach(context.set_value("c", 1))
71+
self.assertEqual(context.get_current(), {"c": 1})
72+
t2 = context.attach(context.set_value("c", 2))
73+
self.assertEqual(context.get_current(), {"c": 2})
74+
context.detach(t1)
75+
self.assertEqual(context.get_current(), {})
76+
context.detach(t2)
77+
self.assertEqual(context.get_current(), {"c": 1})

‎opentelemetry-api/tests/context/test_context.py

Lines changed: 5 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -19,12 +19,12 @@
1919

2020

2121
def do_work() -> None:
22-
context.set_current(context.set_value("say", "bar"))
22+
context.attach(context.set_value("say", "bar"))
2323

2424

2525
class TestContext(unittest.TestCase):
2626
def setUp(self):
27-
context.set_current(Context())
27+
context.attach(Context())
2828

2929
def test_context(self):
3030
self.assertIsNone(context.get_value("say"))
@@ -55,11 +55,10 @@ def test_context_is_immutable(self):
5555
context.get_current()["test"] = "cant-change-immutable"
5656

5757
def test_set_current(self):
58-
context.set_current(context.set_value("a", "yyy"))
58+
context.attach(context.set_value("a", "yyy"))
5959

60-
old_context = context.set_current(context.set_value("a", "zzz"))
61-
self.assertEqual("yyy", context.get_value("a", context=old_context))
60+
token = context.attach(context.set_value("a", "zzz"))
6261
self.assertEqual("zzz", context.get_value("a"))
6362

64-
context.set_current(old_context)
63+
context.detach(token)
6564
self.assertEqual("yyy", context.get_value("a"))

0 commit comments

Comments
 (0)