Use contextvars if available (getsentry#6)

untitaker · web-flow · commit dcbb6275e32e · 2018-07-20T09:39:33.000+02:00
diff --git a/sentry_sdk/hub.py b/sentry_sdk/hub.py
@@ -1,26 +1,25 @@
 import sys
 import copy
 import linecache
-from threading import local
 from contextlib import contextmanager
 
 from ._compat import with_metaclass
 from .scope import Scope
 from .utils import exceptions_from_error_tuple, create_event, \
-    skip_internal_frames
+    skip_internal_frames, ContextVar
 
 
-_local = local()
+_local = ContextVar('sentry_current_hub')
 
 
 class HubMeta(type):
 
     @property
     def current(self):
-        try:
-            rv = _local.hub
-        except AttributeError:
-            _local.hub = rv = Hub(GLOBAL_HUB)
+        rv = _local.get(None)
+        if rv is None:
+            rv = Hub(GLOBAL_HUB)
+            _local.set(rv)
         return rv
 
     @property
@@ -32,10 +31,10 @@ class _HubManager(object):
 
     def __init__(self, hub):
         self._old = Hub.current
-        _local.hub = hub
+        _local.set(hub)
 
     def __exit__(self, exc_type, exc_value, tb):
-        _local.hub = self._old
+        _local.set(self._old)
 
 
 class _ScopeManager(object):
diff --git a/sentry_sdk/integrations/django/__init__.py b/sentry_sdk/integrations/django/__init__.py
@@ -1,6 +1,6 @@
 from __future__ import absolute_import
 
-from threading import Lock, local
+from threading import Lock
 
 from django.conf import settings
 from django.core import signals
@@ -10,6 +10,7 @@
     from django.core.urlresolvers import resolve
 
 from sentry_sdk import get_current_hub, configure_scope, capture_exception
+from sentry_sdk.utils import ContextVar
 
 
 try:
@@ -23,15 +24,15 @@
 def _get_transaction_from_request(request):
     return resolve(request.path).func.__name__
 
-_request_scope = local()
+_request_scope = ContextVar('sentry_django_request_scope')
 
 
 # request_started (or any other signal) cannot be used because the request is
 # not yet available
 class SentryMiddleware(MiddlewareMixin):
     def process_request(self, request):
-        assert getattr(_request_scope, 'manager', None) is None, 'race condition'
-        _request_scope.manager = get_current_hub().push_scope().__enter__()
+        assert _request_scope.get(None) is None, 'race condition'
+        _request_scope.set(get_current_hub().push_scope().__enter__())
 
         try:
             with configure_scope() as scope:
@@ -41,9 +42,10 @@ def process_request(self, request):
 
 
 def _request_finished(*args, **kwargs):
-    assert getattr(_request_scope, 'manager', None) is not None, 'race condition'
-    _request_scope.manager.__exit__(None, None, None)
-    _request_scope.manager = None
+    val = _request_scope.get(None)
+    assert val is not None, 'race condition'
+    val.__exit__(None, None, None)
+    _request_scope.set(None)
 
 
 def _got_request_exception(request=None, **kwargs):
diff --git a/sentry_sdk/utils.py b/sentry_sdk/utils.py
@@ -339,3 +339,23 @@ def to_string(value):
         return text_type(value)
     except UnicodeDecodeError:
         return repr(value)[1:-1]
+
+
+try:
+    from contextvars import ContextVar
+except ImportError:
+    from threading import local
+    _nothing = object()
+
+    class ContextVar(object):
+        # Super-limited impl of ContextVar
+
+        def __init__(self, name):
+            self._name = name
+            self._local = local()
+
+        def get(self, default):
+            return getattr(self._local, 'value', default)
+
+        def set(self, value):
+            setattr(self._local, 'value', value)
