Skip to content

Commit eff8f78

Browse files
authored
(5) Add reset() to thread local ContextVar and no-op copy_context() (#2570)
Improves the capabilities of our threadlocal context variables, we use when no "real" context variables are available (for example in old Python versions) - Adds a no-op [copy_context](https://docs.python.org/3/library/contextvars.html#contextvars.copy_context) function to use in environments without context vars - Adds [reset](https://docs.python.org/3/library/contextvars.html#contextvars.ContextVar.reset) functionality to our threadlocal based context vars. This is preparation work for refactoring how we deal with Hubs and Scopes in the future.
1 parent b4183bd commit eff8f78

File tree

2 files changed

+39
-13
lines changed

2 files changed

+39
-13
lines changed

sentry_sdk/utils.py

Lines changed: 38 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
import logging
55
import math
66
import os
7+
import random
78
import re
89
import subprocess
910
import sys
@@ -1248,24 +1249,49 @@ def _make_threadlocal_contextvars(local):
12481249
class ContextVar(object):
12491250
# Super-limited impl of ContextVar
12501251

1251-
def __init__(self, name):
1252-
# type: (str) -> None
1252+
def __init__(self, name, default=None):
1253+
# type: (str, Any) -> None
12531254
self._name = name
1255+
self._default = default
12541256
self._local = local()
1257+
self._original_local = local()
12551258

1256-
def get(self, default):
1259+
def get(self, default=None):
12571260
# type: (Any) -> Any
1258-
return getattr(self._local, "value", default)
1261+
return getattr(self._local, "value", default or self._default)
12591262

12601263
def set(self, value):
1261-
# type: (Any) -> None
1264+
# type: (Any) -> Any
1265+
token = str(random.getrandbits(64))
1266+
original_value = self.get()
1267+
setattr(self._original_local, token, original_value)
12621268
self._local.value = value
1269+
return token
1270+
1271+
def reset(self, token):
1272+
# type: (Any) -> None
1273+
self._local.value = getattr(self._original_local, token)
1274+
del self._original_local[token]
12631275

12641276
return ContextVar
12651277

12661278

1279+
def _make_noop_copy_context():
1280+
# type: () -> Callable[[], Any]
1281+
class NoOpContext:
1282+
def run(self, func, *args, **kwargs):
1283+
# type: (Callable[..., Any], *Any, **Any) -> Any
1284+
return func(*args, **kwargs)
1285+
1286+
def copy_context():
1287+
# type: () -> NoOpContext
1288+
return NoOpContext()
1289+
1290+
return copy_context
1291+
1292+
12671293
def _get_contextvars():
1268-
# type: () -> Tuple[bool, type]
1294+
# type: () -> Tuple[bool, type, Callable[[], Any]]
12691295
"""
12701296
Figure out the "right" contextvars installation to use. Returns a
12711297
`contextvars.ContextVar`-like class with a limited API.
@@ -1281,28 +1307,28 @@ def _get_contextvars():
12811307
# `aiocontextvars` is absolutely required for functional
12821308
# contextvars on Python 3.6.
12831309
try:
1284-
from aiocontextvars import ContextVar
1310+
from aiocontextvars import ContextVar, copy_context
12851311

1286-
return True, ContextVar
1312+
return True, ContextVar, copy_context
12871313
except ImportError:
12881314
pass
12891315
else:
12901316
# On Python 3.7 contextvars are functional.
12911317
try:
1292-
from contextvars import ContextVar
1318+
from contextvars import ContextVar, copy_context
12931319

1294-
return True, ContextVar
1320+
return True, ContextVar, copy_context
12951321
except ImportError:
12961322
pass
12971323

12981324
# Fall back to basic thread-local usage.
12991325

13001326
from threading import local
13011327

1302-
return False, _make_threadlocal_contextvars(local)
1328+
return False, _make_threadlocal_contextvars(local), _make_noop_copy_context()
13031329

13041330

1305-
HAS_REAL_CONTEXTVARS, ContextVar = _get_contextvars()
1331+
HAS_REAL_CONTEXTVARS, ContextVar, copy_context = _get_contextvars()
13061332

13071333
CONTEXTVARS_ERROR_MESSAGE = """
13081334

tests/utils/test_contextvars.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,7 @@ def test_leaks(maybe_monkeypatched_threading):
1212

1313
from sentry_sdk import utils
1414

15-
_, ContextVar = utils._get_contextvars() # noqa: N806
15+
_, ContextVar, _ = utils._get_contextvars() # noqa: N806
1616

1717
ts = []
1818

0 commit comments

Comments
 (0)