refactor: Asyncio instrumentation

Signed-off-by: Paulo Vital <[email protected]>

Author: pvital

src/instana/__init__.py

@@ -163,7 +163,7 @@ def boot_agent():
 
         # Import & initialize instrumentation
         from instana.instrumentation import (
-            # asyncio,  # noqa: F401
+            asyncio,  # noqa: F401
             boto3_inst,  # noqa: F401
             # cassandra_inst,  # noqa: F401
             # couchbase_inst,  # noqa: F401

src/instana/configurator.py

@@ -5,16 +5,13 @@
 This file contains a config object that will hold configuration options for the package.
 Defaults are set and can be overridden after package load.
 """
-from .util import DictionaryOfStan
+
+from instana.util import DictionaryOfStan
 
 # La Protagonista
 config = DictionaryOfStan()
 
 
 # This option determines if tasks created via asyncio (with ensure_future or create_task) will
 # automatically carry existing context into the created task.
-config['asyncio_task_context_propagation']['enabled'] = False
-
-
-
-
+config["asyncio_task_context_propagation"]["enabled"] = False

src/instana/instrumentation/asyncio.py

@@ -2,46 +2,89 @@
 # (c) Copyright Instana Inc. 2019
 
 
+import time
+from contextlib import contextmanager
+from typing import Any, Callable, Dict, Iterator, Tuple
+
 import wrapt
-from opentracing.scope_managers.constants import ACTIVE_ATTR
-from opentracing.scope_managers.contextvars import no_parent_scope
+from opentelemetry.trace import use_span
+from opentelemetry.trace.status import StatusCode
 
-from ..configurator import config
-from ..log import logger
-from ..singletons import async_tracer
+from instana.configurator import config
+from instana.log import logger
+from instana.span.span import InstanaSpan
+from instana.util.traceutils import get_tracer_tuple, tracing_is_off
 
 try:
     import asyncio
 
     @wrapt.patch_function_wrapper("asyncio", "ensure_future")
-    def ensure_future_with_instana(wrapped, instance, argv, kwargs):
-        if config["asyncio_task_context_propagation"]["enabled"] is False:
-            with no_parent_scope():
-                return wrapped(*argv, **kwargs)
-
-        scope = async_tracer.scope_manager.active
-        task = wrapped(*argv, **kwargs)
-
-        if scope is not None:
-            setattr(task, ACTIVE_ATTR, scope)
+    def ensure_future_with_instana(
+        wrapped: Callable[..., asyncio.ensure_future],
+        instance: object,
+        argv: Tuple[object, Tuple[object, ...]],
+        kwargs: Dict[str, Any],
+    ) -> object:
+        if (
+            not config["asyncio_task_context_propagation"]["enabled"]
+            or tracing_is_off()
+        ):
+            return wrapped(*argv, **kwargs)
 
-        return task
+        with _start_as_current_async_span() as span:
+            try:
+                span.set_status(StatusCode.OK)
+                return wrapped(*argv, **kwargs)
+            except Exception as exc:
+                logger.debug(f"asyncio ensure_future_with_instana error: {exc}")
 
     if hasattr(asyncio, "create_task"):
 
         @wrapt.patch_function_wrapper("asyncio", "create_task")
-        def create_task_with_instana(wrapped, instance, argv, kwargs):
-            if config["asyncio_task_context_propagation"]["enabled"] is False:
-                with no_parent_scope():
+        def create_task_with_instana(
+            wrapped: Callable[..., asyncio.create_task],
+            instance: object,
+            argv: Tuple[object, Tuple[object, ...]],
+            kwargs: Dict[str, Any],
+        ) -> object:
+            if (
+                not config["asyncio_task_context_propagation"]["enabled"]
+                or tracing_is_off()
+            ):
+                return wrapped(*argv, **kwargs)
+
+            with _start_as_current_async_span() as span:
+                try:
+                    span.set_status(StatusCode.OK)
                     return wrapped(*argv, **kwargs)
+                except Exception as exc:
+                    logger.debug(f"asyncio create_task_with_instana error: {exc}")
 
-            scope = async_tracer.scope_manager.active
-            task = wrapped(*argv, **kwargs)
+    @contextmanager
+    def _start_as_current_async_span() -> Iterator[InstanaSpan]:
+        """
+        Creates and yield a special InstanaSpan to only propagate the Asyncio
+        context.
+        """
+        tracer, parent_span, _ = get_tracer_tuple()
+        parent_context = parent_span.get_span_context() if parent_span else None
 
-            if scope is not None:
-                setattr(task, ACTIVE_ATTR, scope)
+        _time = time.time_ns()
 
-            return task
+        span = InstanaSpan(
+            name="asyncio",
+            context=parent_context,
+            span_processor=tracer.span_processor,
+            start_time=_time,
+            end_time=_time,
+        )
+        with use_span(
+            span,
+            end_on_exit=False,
+            record_exception=False,
+            set_status_on_exception=False,
+        ) as span:
+            yield span
 
     logger.debug("Instrumenting asyncio")
 except ImportError:

src/instana/span/kind.py

@@ -6,7 +6,7 @@
 
 EXIT_KIND = ("exit", "client", "producer", SpanKind.CLIENT, SpanKind.PRODUCER)
 
-LOCAL_SPANS = ("render", SpanKind.INTERNAL)
+LOCAL_SPANS = ("asyncio", "render", SpanKind.INTERNAL)
 
 HTTP_SPANS = (
     "aiohttp-client",

src/instana/util/traceutils.py

@@ -22,12 +22,13 @@ def extract_custom_headers(tracing_span, headers) -> None:
 
 def get_active_tracer() -> Optional[InstanaTracer]:
     try:
-        # ToDo: Might have to add additional stuff when testing with async and tornado tracer
         current_span = get_current_span()
-        if current_span and current_span.is_recording():
-            return tracer
-        else:
+        if current_span:
+            # asyncio Spans are used as NonRecording Spans solely for context propagation
+            if current_span.is_recording() or current_span.name == "asyncio":
+                return tracer
             return None
+        return None
     except Exception:
         # Do not try to log this with instana, as there is no active tracer and there will be an infinite loop at least
         # for PY2