Thanks to visit codestin.com
Credit goes to github.com

Skip to content

Commit ba7d8e9

Browse files
author
Alex Boten
committed
Adding attach/detach methods as per spec
This change updates the Context API with the following: - removes the `remove_value` method - removes the `set_current` method - adds `attach` and `detach` methods Signed-off-by: Alex Boten <[email protected]>
1 parent 120ae29 commit ba7d8e9

File tree

11 files changed

+115
-63
lines changed

11 files changed

+115
-63
lines changed

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

Lines changed: 22 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -57,24 +57,6 @@ def set_value(
5757
return Context(new_values)
5858

5959

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
70-
"""
71-
if context is None:
72-
context = get_current()
73-
new_values = context.copy()
74-
new_values.pop(key, None)
75-
return Context(new_values)
76-
77-
7860
def get_current() -> Context:
7961
"""To access the context associated with program execution,
8062
the RuntimeContext API provides a function which takes no arguments
@@ -104,16 +86,29 @@ def get_current() -> Context:
10486
return _RUNTIME_CONTEXT.get_current() # type:ignore
10587

10688

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.
89+
def attach(context: Context) -> object:
90+
"""Associates a Context with the caller's current execution unit. Returns
91+
a token that can be used to restore the previous Context.
92+
93+
Args:
94+
context: The Context to set as current.
95+
"""
96+
get_current()
97+
98+
return _RUNTIME_CONTEXT.set_current(context) # type:ignore
99+
100+
101+
def detach(token: object) -> None:
102+
"""Resets the Context associated with the caller's current execution unit
103+
to the value it had before attaching a specified Context.
110104
111105
Args:
112-
context: The context to use as current.
106+
token: The Token that was returned by a previous call to attach a Context.
113107
"""
114-
old_context = get_current()
115-
_RUNTIME_CONTEXT.set_current(context) # type:ignore
116-
return old_context
108+
try:
109+
_RUNTIME_CONTEXT.reset(token) # type: ignore
110+
except ValueError:
111+
logger.error("Failed to detach context")
117112

118113

119114
def with_current_context(
@@ -127,10 +122,9 @@ def call_with_current_context(
127122
*args: "object", **kwargs: "object"
128123
) -> "object":
129124
try:
130-
backup = get_current()
131-
set_current(caller_context)
125+
token = attach(caller_context)
132126
return func(*args, **kwargs)
133127
finally:
134-
set_current(backup)
128+
detach(token)
135129

136130
return call_with_current_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 set_current(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 reset(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: 9 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,7 @@
1111
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
1212
# See the License for the specific language governing permissions and
1313
# limitations under the License.
14-
from contextvars import ContextVar
14+
from contextvars import ContextVar, Token
1515
from sys import version_info
1616

1717
from opentelemetry.context.context import Context, RuntimeContext
@@ -35,13 +35,19 @@ def __init__(self) -> None:
3535
self._CONTEXT_KEY, default=Context()
3636
)
3737

38-
def set_current(self, context: Context) -> None:
38+
def set_current(self, context: Context) -> object:
3939
"""See `opentelemetry.context.RuntimeContext.set_current`."""
40-
self._current_context.set(context)
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 reset(self, token: object) -> None:
47+
"""See `opentelemetry.context.RuntimeContext.reset`."""
48+
if not isinstance(token, Token):
49+
raise ValueError("invalid token")
50+
self._current_context.reset(token)
51+
4652

4753
__all__ = ["ContextVarsRuntimeContext"]

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

Lines changed: 11 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -28,9 +28,12 @@ class ThreadLocalRuntimeContext(RuntimeContext):
2828
def __init__(self) -> None:
2929
self._current_context = threading.local()
3030

31-
def set_current(self, context: Context) -> None:
31+
def set_current(self, context: Context) -> object:
3232
"""See `opentelemetry.context.RuntimeContext.set_current`."""
33+
current = self.get_current()
34+
setattr(self._current_context, str(id(current)), current)
3335
setattr(self._current_context, self._CONTEXT_KEY, context)
36+
return str(id(current))
3437

3538
def get_current(self) -> Context:
3639
"""See `opentelemetry.context.RuntimeContext.get_current`."""
@@ -43,5 +46,12 @@ def get_current(self) -> Context:
4346
) # type: Context
4447
return context
4548

49+
def reset(self, token: object) -> None:
50+
"""See `opentelemetry.context.RuntimeContext.reset`."""
51+
if not hasattr(self._current_context, str(token)):
52+
raise ValueError("invalid token")
53+
context = getattr(self._current_context, str(token)) # type: Context
54+
setattr(self._current_context, self._CONTEXT_KEY, context)
55+
4656

4757
__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/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"))

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

Lines changed: 23 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@
1313
# limitations under the License.
1414

1515
import unittest
16+
from logging import ERROR
1617
from unittest.mock import patch
1718

1819
from opentelemetry import context
@@ -27,18 +28,19 @@
2728

2829

2930
def do_work() -> None:
30-
context.set_current(context.set_value("say", "bar"))
31+
context.attach(context.set_value("say", "bar"))
3132

3233

3334
class TestContextVarsContext(unittest.TestCase):
3435
def setUp(self):
3536
self.previous_context = context.get_current()
3637

3738
def tearDown(self):
38-
context.set_current(self.previous_context)
39+
context.attach(self.previous_context)
3940

4041
@patch(
41-
"opentelemetry.context._RUNTIME_CONTEXT", ContextVarsRuntimeContext() # type: ignore
42+
"opentelemetry.context._RUNTIME_CONTEXT", # type: ignore
43+
ContextVarsRuntimeContext(),
4244
)
4345
def test_context(self):
4446
self.assertIsNone(context.get_value("say"))
@@ -56,7 +58,8 @@ def test_context(self):
5658
self.assertEqual(context.get_value("say", context=third), "bar")
5759

5860
@patch(
59-
"opentelemetry.context._RUNTIME_CONTEXT", ContextVarsRuntimeContext() # type: ignore
61+
"opentelemetry.context._RUNTIME_CONTEXT", # type: ignore
62+
ContextVarsRuntimeContext(),
6063
)
6164
def test_set_value(self):
6265
first = context.set_value("a", "yyy")
@@ -66,3 +69,19 @@ def test_set_value(self):
6669
self.assertEqual("zzz", context.get_value("a", context=second))
6770
self.assertEqual("---", context.get_value("a", context=third))
6871
self.assertEqual(None, context.get_value("a"))
72+
73+
@patch(
74+
"opentelemetry.context._RUNTIME_CONTEXT", # type: ignore
75+
ContextVarsRuntimeContext(),
76+
)
77+
def test_set_current(self):
78+
context.attach(context.set_value("a", "yyy"))
79+
80+
token = context.attach(context.set_value("a", "zzz"))
81+
self.assertEqual("zzz", context.get_value("a"))
82+
83+
context.detach(token)
84+
self.assertEqual("yyy", context.get_value("a"))
85+
86+
with self.assertLogs(level=ERROR):
87+
context.detach("some garbage")

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

Lines changed: 23 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -13,25 +13,27 @@
1313
# limitations under the License.
1414

1515
import unittest
16+
from logging import ERROR
1617
from unittest.mock import patch
1718

1819
from opentelemetry import context
1920
from opentelemetry.context.threadlocal_context import ThreadLocalRuntimeContext
2021

2122

2223
def do_work() -> None:
23-
context.set_current(context.set_value("say", "bar"))
24+
context.attach(context.set_value("say", "bar"))
2425

2526

2627
class TestThreadLocalContext(unittest.TestCase):
2728
def setUp(self):
2829
self.previous_context = context.get_current()
2930

3031
def tearDown(self):
31-
context.set_current(self.previous_context)
32+
context.attach(self.previous_context)
3233

3334
@patch(
34-
"opentelemetry.context._RUNTIME_CONTEXT", ThreadLocalRuntimeContext() # type: ignore
35+
"opentelemetry.context._RUNTIME_CONTEXT", # type: ignore
36+
ThreadLocalRuntimeContext(),
3537
)
3638
def test_context(self):
3739
self.assertIsNone(context.get_value("say"))
@@ -49,7 +51,8 @@ def test_context(self):
4951
self.assertEqual(context.get_value("say", context=third), "bar")
5052

5153
@patch(
52-
"opentelemetry.context._RUNTIME_CONTEXT", ThreadLocalRuntimeContext() # type: ignore
54+
"opentelemetry.context._RUNTIME_CONTEXT", # type: ignore
55+
ThreadLocalRuntimeContext(),
5356
)
5457
def test_set_value(self):
5558
first = context.set_value("a", "yyy")
@@ -59,3 +62,19 @@ def test_set_value(self):
5962
self.assertEqual("zzz", context.get_value("a", context=second))
6063
self.assertEqual("---", context.get_value("a", context=third))
6164
self.assertEqual(None, context.get_value("a"))
65+
66+
@patch(
67+
"opentelemetry.context._RUNTIME_CONTEXT", # type: ignore
68+
ThreadLocalRuntimeContext(),
69+
)
70+
def test_set_current(self):
71+
context.attach(context.set_value("a", "yyy"))
72+
73+
token = context.attach(context.set_value("a", "zzz"))
74+
self.assertEqual("zzz", context.get_value("a"))
75+
76+
context.detach(token)
77+
self.assertEqual("yyy", context.get_value("a"))
78+
79+
with self.assertLogs(level=ERROR):
80+
context.detach("some garbage")

‎opentelemetry-sdk/src/opentelemetry/sdk/trace/__init__.py

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -541,14 +541,13 @@ def use_span(
541541
) -> Iterator[trace_api.Span]:
542542
"""See `opentelemetry.trace.Tracer.use_span`."""
543543
try:
544-
context_snapshot = context_api.get_current()
545-
context_api.set_current(
544+
token = context_api.attach(
546545
context_api.set_value(self.source.key, span)
547546
)
548547
try:
549548
yield span
550549
finally:
551-
context_api.set_current(context_snapshot)
550+
context_api.detach(token)
552551

553552
except Exception as error: # pylint: disable=broad-except
554553
if (

‎opentelemetry-sdk/src/opentelemetry/sdk/trace/export/__init__.py

Lines changed: 5 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,7 @@
1919
import typing
2020
from enum import Enum
2121

22-
from opentelemetry.context import get_current, set_current, set_value
22+
from opentelemetry.context import attach, detach, get_current, set_value
2323
from opentelemetry.trace import DefaultSpan
2424
from opentelemetry.util import time_ns
2525

@@ -75,14 +75,13 @@ def on_start(self, span: Span) -> None:
7575
pass
7676

7777
def on_end(self, span: Span) -> None:
78-
backup_context = get_current()
79-
set_current(set_value("suppress_instrumentation", True))
78+
token = attach(set_value("suppress_instrumentation", True))
8079
try:
8180
self.span_exporter.export((span,))
8281
# pylint: disable=broad-except
8382
except Exception:
8483
logger.exception("Exception while exporting Span.")
85-
set_current(backup_context)
84+
detach(token)
8685

8786
def shutdown(self) -> None:
8887
self.span_exporter.shutdown()
@@ -202,16 +201,15 @@ def export(self) -> None:
202201
else:
203202
self.spans_list[idx] = span
204203
idx += 1
205-
backup_context = get_current()
206-
set_current(set_value("suppress_instrumentation", True))
204+
token = attach(set_value("suppress_instrumentation", True))
207205
try:
208206
# Ignore type b/c the Optional[None]+slicing is too "clever"
209207
# for mypy
210208
self.span_exporter.export(self.spans_list[:idx]) # type: ignore
211209
# pylint: disable=broad-except
212210
except Exception:
213211
logger.exception("Exception while exporting Span batch.")
214-
set_current(backup_context)
212+
detach(token)
215213

216214
if notify_flush:
217215
with self.flush_condition:

0 commit comments

Comments
 (0)