implement support for asyncio [double-plus-WIP]

Signed-off-by: Benjamin Wohlwend <[email protected]>

Author: beniwohli

elasticapm/__init__.py

@@ -8,6 +8,7 @@
 :copyright: (c) 2010 by the Sentry Team, see AUTHORS for more details.
 :license: BSD, see LICENSE for more details.
 """
+import sys
 
 __all__ = ('VERSION', 'Client')
 
@@ -20,6 +21,11 @@
 from elasticapm.base import Client
 from elasticapm.conf import setup_logging  # noqa: F401
 from elasticapm.instrumentation.control import instrument, uninstrument  # noqa: F401
-from elasticapm.traces import capture_span, set_context, set_custom_context  # noqa: F401
+from elasticapm.traces import set_context, set_custom_context  # noqa: F401
 from elasticapm.traces import set_transaction_name, set_user_context, tag  # noqa: F401
 from elasticapm.traces import set_transaction_result  # noqa: F401
+
+if sys.version_info >= (3, 5):
+    from elasticapm.contrib.asyncio.traces import async_capture_span as capture_span # noqa: F401
+else:
+    from elasticapm.traces import capture_span # noqa: F401

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

@@ -0,0 +1,16 @@
+from elasticapm.traces import capture_span, error_logger, get_transaction
+
+
+class async_capture_span(capture_span):
+    async def __aenter__(self):
+        transaction = get_transaction()
+        if transaction and transaction.is_sampled:
+            transaction.begin_span(self.name, self.type, context=self.extra, leaf=self.leaf)
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        transaction = get_transaction()
+        if transaction and transaction.is_sampled:
+            try:
+                transaction.end_span(self.skip_frames)
+            except LookupError:
+                error_logger.info('ended non-existing span %s of type %s', self.name, self.type)

elasticapm/contrib/asyncio/traces.py

@@ -14,31 +14,16 @@
 
 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):
@@ -54,11 +39,9 @@ def __init__(self, frames_collector_func, transaction_type="custom", is_sampled=
         self._frames_collector_func = frames_collector_func
 
         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 = {}
 
@@ -69,46 +52,35 @@ 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()
+        if parent_span and parent_span.leaf:
+            span = DroppedSpan(parent_span, leaf=True)
+        elif self.max_spans and self._span_counter > self.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, 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, name, span_type, start, context, leaf)
+            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.spans.append(span)
-
+        set_span(span.parent)
         return span
 
     def to_dict(self):
@@ -162,14 +134,22 @@ 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, collect_frequency, sample_rate=1.0, max_spans=0, max_queue_size=None,
                  span_frames_min_duration=None, ignore_patterns=None):
@@ -219,7 +199,8 @@ def begin_transaction(self, transaction_type):
         is_sampled = self._sample_rate == 1.0 or self._sample_rate > random.random()
         transaction = Transaction(self._frames_collector_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
+
+        set_transaction(transaction)
         return transaction
 
     def _should_ignore(self, transaction_name):
@@ -270,7 +251,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)