refactored span stack and added contextvars as alternative for threadlocals (#291)

this is a partial "import" of the asyncio work done in #252, sans the
asyncio specific extensions

closes #291

Author: beniwohli

CHANGELOG.md

@@ -13,10 +13,15 @@ Further breaking changes:
  * Some settings now require a unit for duration or size. See documentation on
    configuration for more information.
 
+
+Other changes:
+ * on Python 3.7, use [contextvars](https://docs.python.org/3/library/contextvars.html) instead of threadlocals for storing
+   current transaction and span. This is a necessary precursor for full asyncio support. (#291)
+
 ## Unreleased
 
  * fixed an issue with detecting names of wrapped functions that are partials (#294)
- 
+  
 ## v3.0.1
 
 [Check the diff](https://github.com/elastic/apm-agent-python/compare/v3.0.0...v3.0.1)

elasticapm/base.py

@@ -12,6 +12,7 @@
 from __future__ import absolute_import
 
 import datetime
+import inspect
 import logging
 import os
 import platform
@@ -112,9 +113,12 @@ def __init__(self, config=None, **inline):
         else:
             skip_modules = ("elasticapm.",)
 
-        def frames_collector_func():
-            return self._get_stack_info_for_trace(
-                stacks.iter_stack_frames(skip_top_modules=skip_modules),
+        self.transaction_store = TransactionsStore(
+            frames_collector_func=lambda: list(
+                stacks.iter_stack_frames(start_frame=inspect.currentframe(), skip_top_modules=skip_modules)
+            ),
+            frames_processing_func=lambda frames: self._get_stack_info_for_trace(
+                frames,
                 library_frame_context_lines=self.config.source_lines_span_library_frames,
                 in_app_frame_context_lines=self.config.source_lines_span_app_frames,
                 with_locals=self.config.collect_local_variables in ("all", "transactions"),
@@ -126,12 +130,8 @@ def frames_collector_func():
                     ),
                     local_var,
                 ),
-            )
-
-        self.transaction_store = TransactionsStore(
-            frames_collector_func=frames_collector_func,
+            ),
             queue_func=self.queue,
-            collect_frequency=self.config.api_request_time,
             sample_rate=self.config.transaction_sample_rate,
             max_spans=self.config.transaction_max_spans,
             span_frames_min_duration=self.config.span_frames_min_duration,

elasticapm/context/__init__.py

@@ -0,0 +1,31 @@
+from __future__ import absolute_import
+
+import contextvars
+
+elasticapm_transaction_var = contextvars.ContextVar("elasticapm_transaction_var")
+elasticapm_span_var = contextvars.ContextVar("elasticapm_span_var")
+
+
+def get_transaction(clear=False):
+    try:
+        transaction = elasticapm_transaction_var.get()
+        if clear:
+            set_transaction(None)
+        return transaction
+    except LookupError:
+        return None
+
+
+def set_transaction(transaction):
+    elasticapm_transaction_var.set(transaction)
+
+
+def get_span():
+    try:
+        return elasticapm_span_var.get()
+    except LookupError:
+        return None
+
+
+def set_span(span):
+    elasticapm_span_var.set(span)

elasticapm/context/contextvars.py

@@ -0,0 +1,30 @@
+import threading
+
+thread_local = threading.local()
+thread_local.transaction = None
+elasticapm_span_var = None
+
+
+def get_transaction(clear=False):
+    """
+    Get the transaction registered for the current thread.
+
+    :return:
+    :rtype: Transaction
+    """
+    transaction = getattr(thread_local, "transaction", None)
+    if clear:
+        thread_local.transaction = None
+    return transaction
+
+
+def set_transaction(transaction):
+    thread_local.transaction = transaction
+
+
+def get_span():
+    return getattr(thread_local, "span", None)
+
+
+def set_span(span):
+    thread_local.span = span

elasticapm/context/threadlocal.py

@@ -15,43 +15,20 @@
 
 error_logger = logging.getLogger("elasticapm.errors")
 
-thread_local = threading.local()
-thread_local.transaction = None
-
-
 _time_func = timeit.default_timer
 
 
 TAG_RE = re.compile('^[^.*"]+$')
 
 
-DROPPED_SPAN = object()
-IGNORED_SPAN = object()
-
-
-def get_transaction(clear=False):
-    """
-    Get the transaction registered for the current thread.
-
-    :return:
-    :rtype: Transaction
-    """
-    transaction = getattr(thread_local, "transaction", None)
-    if clear:
-        thread_local.transaction = None
-    return transaction
+try:
+    from elasticapm.context.contextvars import get_transaction, set_transaction, get_span, set_span
+except ImportError:
+    from elasticapm.context.threadlocal import get_transaction, set_transaction, get_span, set_span
 
 
 class Transaction(object):
-    def __init__(
-        self,
-        frames_collector_func,
-        queue_func,
-        transaction_type="custom",
-        is_sampled=True,
-        max_spans=None,
-        span_frames_min_duration=None,
-    ):
+    def __init__(self, store, transaction_type="custom", is_sampled=True):
         self.id = str(uuid.uuid4())
         self.trace_id = None  # for later use in distributed tracing
         self.timestamp = datetime.datetime.utcnow()
@@ -60,15 +37,10 @@ def __init__(
         self.duration = None
         self.result = None
         self.transaction_type = transaction_type
-        self._frames_collector_func = frames_collector_func
-        self._queue_func = queue_func
+        self._store = store
 
         self.spans = []
-        self.span_stack = []
-        self.max_spans = max_spans
-        self.span_frames_min_duration = span_frames_min_duration
         self.dropped_spans = 0
-        self.ignore_subtree = False
         self.context = {}
         self.tags = {}
 
@@ -79,45 +51,40 @@ def end_transaction(self, skip_frames=8):
         self.duration = _time_func() - self.start_time
 
     def begin_span(self, name, span_type, context=None, leaf=False):
-        # If we were already called with `leaf=True`, we'll just push
-        # a placeholder on the stack.
-        if self.ignore_subtree:
-            self.span_stack.append(IGNORED_SPAN)
-            return None
-
-        if leaf:
-            self.ignore_subtree = True
-
-        self._span_counter += 1
-
-        if self.max_spans and self._span_counter > self.max_spans:
+        parent_span = get_span()
+        store = self._store
+        if parent_span and parent_span.leaf:
+            span = DroppedSpan(parent_span, leaf=True)
+        elif store.max_spans and self._span_counter > store.max_spans - 1:
             self.dropped_spans += 1
-            self.span_stack.append(DROPPED_SPAN)
-            return None
-
-        start = _time_func() - self.start_time
-        span = Span(self._span_counter - 1, self.id, self.trace_id, name, span_type, start, context)
-        self.span_stack.append(span)
+            span = DroppedSpan(parent_span)
+            self._span_counter += 1
+        else:
+            start = _time_func() - self.start_time
+            span = Span(self._span_counter, self.id, self.trace_id, name, span_type, start, context, leaf)
+            span.frames = store.frames_collector_func()
+            span.parent = parent_span
+            self._span_counter += 1
+        set_span(span)
         return span
 
     def end_span(self, skip_frames):
-        span = self.span_stack.pop()
-        if span is IGNORED_SPAN:
-            return None
-
-        self.ignore_subtree = False
-
-        if span is DROPPED_SPAN:
+        span = get_span()
+        if span is None:
+            raise LookupError()
+        if isinstance(span, DroppedSpan):
+            set_span(span.parent)
             return
 
         span.duration = _time_func() - span.start_time - self.start_time
 
-        if self.span_stack:
-            span.parent = self.span_stack[-1].idx
-
-        if not self.span_frames_min_duration or span.duration >= self.span_frames_min_duration:
-            span.frames = self._frames_collector_func()[skip_frames:]
-        self._queue_func(SPAN, span.to_dict())
+        if not self._store.span_frames_min_duration or span.duration >= self._store.span_frames_min_duration:
+            span.frames = self._store.frames_processing_func(span.frames)[skip_frames:]
+        else:
+            span.frames = None
+        self.spans.append(span)
+        set_span(span.parent)
+        self._store.queue_func(SPAN, span.to_dict())
         return span
 
     def to_dict(self):
@@ -185,30 +152,38 @@ def to_dict(self):
             "type": encoding.keyword_field(self.type),
             "start": self.start_time * 1000,  # milliseconds
             "duration": self.duration * 1000,  # milliseconds
-            "parent": self.parent,
+            "parent": self.parent.idx if self.parent else None,
             "context": self.context,
         }
         if self.frames:
             result["stacktrace"] = self.frames
         return result
 
 
+class DroppedSpan(object):
+    __slots__ = ("leaf", "parent")
+
+    def __init__(self, parent, leaf=False):
+        self.parent = parent
+        self.leaf = leaf
+
+
 class TransactionsStore(object):
     def __init__(
         self,
         frames_collector_func,
+        frames_processing_func,
         queue_func,
-        collect_frequency,
         sample_rate=1.0,
         max_spans=0,
         span_frames_min_duration=None,
         ignore_patterns=None,
     ):
         self.cond = threading.Condition()
-        self.collect_frequency = collect_frequency
         self.max_spans = max_spans
-        self._queue_func = queue_func
-        self._frames_collector_func = frames_collector_func
+        self.queue_func = queue_func
+        self.frames_processing_func = frames_processing_func
+        self.frames_collector_func = frames_collector_func
         self._transactions = []
         self._last_collect = _time_func()
         self._ignore_patterns = [re.compile(p) for p in ignore_patterns or []]
@@ -244,15 +219,8 @@ def begin_transaction(self, transaction_type):
         :returns the Transaction object
         """
         is_sampled = self._sample_rate == 1.0 or self._sample_rate > random.random()
-        transaction = Transaction(
-            self._frames_collector_func,
-            self._queue_func,
-            transaction_type,
-            max_spans=self.max_spans,
-            span_frames_min_duration=self.span_frames_min_duration,
-            is_sampled=is_sampled,
-        )
-        thread_local.transaction = transaction
+        transaction = Transaction(self, transaction_type, is_sampled=is_sampled)
+        set_transaction(transaction)
         return transaction
 
     def _should_ignore(self, transaction_name):
@@ -271,7 +239,7 @@ def end_transaction(self, result=None, transaction_name=None):
                 return
             if transaction.result is None:
                 transaction.result = result
-            self._queue_func(TRANSACTION, transaction.to_dict())
+            self.queue_func(TRANSACTION, transaction.to_dict())
         return transaction
 
 
@@ -303,7 +271,7 @@ def __exit__(self, exc_type, exc_val, exc_tb):
         if transaction and transaction.is_sampled:
             try:
                 transaction.end_span(self.skip_frames)
-            except IndexError:
+            except LookupError:
                 error_logger.info("ended non-existing span %s of type %s", self.name, self.type)
 
 

elasticapm/traces.py

@@ -158,7 +158,7 @@ def iter_traceback_frames(tb):
         tb = tb.tb_next
 
 
-def iter_stack_frames(frames=None, skip=0, skip_top_modules=()):
+def iter_stack_frames(frames=None, start_frame=None, skip=0, skip_top_modules=()):
     """
     Given an optional list of frames (defaults to current stack),
     iterates over all frames that do not contain the ``__traceback_hide__``
@@ -173,12 +173,13 @@ def iter_stack_frames(frames=None, skip=0, skip_top_modules=()):
     itself.
 
     :param frames: a list of frames, or None
+    :param start_frame: a Frame object or None
     :param skip: number of frames to skip from the beginning
     :param skip_top_modules: tuple of strings
 
     """
     if not frames:
-        frame = inspect.currentframe().f_back
+        frame = start_frame if start_frame is not None else inspect.currentframe().f_back
         frames = _walk_stack(frame)
     stop_ignoring = False
     for i, frame in enumerate(frames):

elasticapm/utils/stacks.py

@@ -586,7 +586,7 @@ def test_transaction_max_spans(elasticapm_client):
     spans = elasticapm_client.events[SPAN]
     assert all(span["transaction_id"] == transaction["id"] for span in spans)
 
-    assert transaction_obj.max_spans == 5
+    assert transaction_obj._store.max_spans == 5
     assert transaction_obj.dropped_spans == 10
     assert len(spans) == 5
     for span in spans: