From 58dc2000af0af899032e558117b9babc80bc3f83 Mon Sep 17 00:00:00 2001
From: Aaron Ramshaw <aaronramshaw@gmail.com>
Date: Wed, 10 Apr 2019 09:07:05 +1000
Subject: [PATCH] Execution context - use thread locals if gevent/eventlet has
 monkey patched it #453

* Extract out logic to decide which backing to use for execution_context
into a function init_execution_context
* Check if gevent or eventlet has monkey patched _threading.local, if it
has then use elasticapm.context.threadlocal as the backing.
---
 elasticapm/context/__init__.py | 37 ++++++++++++++++++++++++++++++++++
 elasticapm/traces.py           |  6 ++----
 tests/context/__init__.py      |  0
 tests/context/test_context.py  | 24 ++++++++++++++++++++++
 4 files changed, 63 insertions(+), 4 deletions(-)
 create mode 100644 tests/context/__init__.py
 create mode 100644 tests/context/test_context.py

diff --git a/elasticapm/context/__init__.py b/elasticapm/context/__init__.py
index 7e2b340e6..5b32c1ba5 100644
--- a/elasticapm/context/__init__.py
+++ b/elasticapm/context/__init__.py
@@ -27,3 +27,40 @@
 #  CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
 #  OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 #  OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+
+def init_execution_context():
+    # If _threading_local has been monkeypatched (by gevent or eventlet), then
+    # we should assume it's use as this will be the most "green-thread safe"
+    if threading_local_monkey_patched():
+        from elasticapm.context.threadlocal import execution_context
+
+        return execution_context
+
+    try:
+        from elasticapm.context.contextvars import execution_context
+    except ImportError:
+        from elasticapm.context.threadlocal import execution_context
+    return execution_context
+
+
+def threading_local_monkey_patched():
+    # Returns True if thread locals have been patched by either gevent of
+    # eventlet
+    try:
+        from gevent.monkey import is_object_patched
+    except ImportError:
+        pass
+    else:
+        if is_object_patched("_threading", "local"):
+            return True
+
+    try:
+        from eventlet.patcher import is_monkey_patched
+    except ImportError:
+        pass
+    else:
+        if is_monkey_patched("thread"):
+            return True
+
+    return False
diff --git a/elasticapm/traces.py b/elasticapm/traces.py
index 5e8c52c34..60e7d8796 100644
--- a/elasticapm/traces.py
+++ b/elasticapm/traces.py
@@ -37,6 +37,7 @@
 
 from elasticapm.conf import constants
 from elasticapm.conf.constants import SPAN, TRANSACTION
+from elasticapm.context import init_execution_context
 from elasticapm.utils import compat, encoding, get_name_from_func
 from elasticapm.utils.disttracing import TraceParent, TracingOptions
 
@@ -51,10 +52,7 @@
 TAG_RE = re.compile('[.*"]')
 
 
-try:
-    from elasticapm.context.contextvars import execution_context
-except ImportError:
-    from elasticapm.context.threadlocal import execution_context
+execution_context = init_execution_context()
 
 
 class Transaction(object):
diff --git a/tests/context/__init__.py b/tests/context/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/tests/context/test_context.py b/tests/context/test_context.py
new file mode 100644
index 000000000..8aa5b9e65
--- /dev/null
+++ b/tests/context/test_context.py
@@ -0,0 +1,24 @@
+import sys
+
+import elasticapm.context
+from elasticapm.context.threadlocal import ThreadLocalContext
+
+
+def test_execution_context_backing():
+    execution_context = elasticapm.context.init_execution_context()
+
+    if sys.version_info[0] == 3 and sys.version_info[1] >= 7:
+        from elasticapm.context.contextvars import ContextVarsContext
+
+        assert isinstance(execution_context, ContextVarsContext)
+    else:
+        assert isinstance(execution_context, ThreadLocalContext)
+
+
+def test_execution_context_monkeypatched(monkeypatch):
+    with monkeypatch.context() as m:
+        m.setattr(elasticapm.context, "threading_local_monkey_patched", lambda: True)
+        execution_context = elasticapm.context.init_execution_context()
+
+    # Should always use ThreadLocalContext when thread local is monkey patched
+    assert isinstance(execution_context, ThreadLocalContext)