refactor: AIOHTTP server and client instrumentation

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

Author: pvital

src/instana/__init__.py

@@ -184,11 +184,10 @@ def boot_agent():
             # sanic_inst,  # noqa: F401
             urllib3,  # noqa: F401
         )
-
-        # from instana.instrumentation.aiohttp import (
-        #     client,  # noqa: F401
-        #     server,  # noqa: F401
-        # )
+        from instana.instrumentation.aiohttp import (
+            client,  # noqa: F401
+            server,  # noqa: F401
+        )
         # from instana.instrumentation.aws import lambda_inst  # noqa: F401
         # from instana.instrumentation.celery import hooks  # noqa: F401
         from instana.instrumentation.django import middleware  # noqa: F401

src/instana/instrumentation/aiohttp/client.py

@@ -2,85 +2,111 @@
 # (c) Copyright Instana Inc. 2019
 
 
-import opentracing
+from types import SimpleNamespace
+from typing import TYPE_CHECKING, Any, Awaitable, Callable, Dict, Tuple
 import wrapt
 
-from ...log import logger
-from ...singletons import agent, async_tracer
-from ...util.secrets import strip_secrets_from_query
-from ...util.traceutils import tracing_is_off
+from opentelemetry.semconv.trace import SpanAttributes
+
+from instana.log import logger
+from instana.propagators.format import Format
+from instana.singletons import agent
+from instana.util.secrets import strip_secrets_from_query
+from instana.util.traceutils import get_tracer_tuple, tracing_is_off
 
 try:
     import aiohttp
-    import asyncio
 
+    if TYPE_CHECKING:
+        from aiohttp.client import ClientSession
+        from instana.span.span import InstanaSpan
 
-    async def stan_request_start(session, trace_config_ctx, params):
+    async def stan_request_start(
+        session: "ClientSession", trace_config_ctx: SimpleNamespace, params
+    ) -> Awaitable[None]:
         try:
             # If we're not tracing, just return
             if tracing_is_off():
-                trace_config_ctx.scope = None
+                trace_config_ctx.span_context = None
                 return
 
-            scope = async_tracer.start_active_span("aiohttp-client", child_of=async_tracer.active_span)
-            trace_config_ctx.scope = scope
+            tracer, parent_span, _ = get_tracer_tuple()
+            parent_context = parent_span.get_span_context() if parent_span else None
+
+            span = tracer.start_span("aiohttp-client", span_context=parent_context)
 
-            async_tracer.inject(scope.span.context, opentracing.Format.HTTP_HEADERS, params.headers)
+            tracer.inject(span.context, Format.HTTP_HEADERS, params.headers)
 
-            parts = str(params.url).split('?')
+            parts = str(params.url).split("?")
             if len(parts) > 1:
-                cleaned_qp = strip_secrets_from_query(parts[1], agent.options.secrets_matcher,
-                                                      agent.options.secrets_list)
-                scope.span.set_tag("http.params", cleaned_qp)
-            scope.span.set_tag("http.url", parts[0])
-            scope.span.set_tag('http.method', params.method)
+                cleaned_qp = strip_secrets_from_query(
+                    parts[1], agent.options.secrets_matcher, agent.options.secrets_list
+                )
+                span.set_attribute("http.params", cleaned_qp)
+            span.set_attribute(SpanAttributes.HTTP_URL, parts[0])
+            span.set_attribute(SpanAttributes.HTTP_METHOD, params.method)
+            trace_config_ctx.span_context = span
         except Exception:
-            logger.debug("stan_request_start", exc_info=True)
+            logger.debug("aiohttp-client stan_request_start error:", exc_info=True)
 
-
-    async def stan_request_end(session, trace_config_ctx, params):
+    async def stan_request_end(
+        session: "ClientSession", trace_config_ctx: SimpleNamespace, params
+    ) -> Awaitable[None]:
         try:
-            scope = trace_config_ctx.scope
-            if scope is not None:
-                scope.span.set_tag('http.status_code', params.response.status)
+            span: "InstanaSpan" = trace_config_ctx.span_context
+            if span:
+                span.set_attribute(
+                    SpanAttributes.HTTP_STATUS_CODE, params.response.status
+                )
 
-                if agent.options.extra_http_headers is not None:
+                if agent.options.extra_http_headers:
                     for custom_header in agent.options.extra_http_headers:
                         if custom_header in params.response.headers:
-                            scope.span.set_tag("http.header.%s" % custom_header, params.response.headers[custom_header])
+                            span.set_attribute(
+                                "http.header.%s" % custom_header,
+                                params.response.headers[custom_header],
+                            )
 
                 if 500 <= params.response.status:
-                    scope.span.mark_as_errored({"http.error": params.response.reason})
+                    span.mark_as_errored({"http.error": params.response.reason})
 
-                scope.close()
+                if span.is_recording():
+                    span.end()
+                trace_config_ctx = None
         except Exception:
-            logger.debug("stan_request_end", exc_info=True)
-
+            logger.debug("aiohttp-client stan_request_end error:", exc_info=True)
 
-    async def stan_request_exception(session, trace_config_ctx, params):
+    async def stan_request_exception(
+        session: "ClientSession", trace_config_ctx: SimpleNamespace, params
+    ) -> Awaitable[None]:
         try:
-            scope = trace_config_ctx.scope
-            if scope is not None:
-                scope.span.log_exception(params.exception)
-                scope.span.set_tag("http.error", str(params.exception))
-                scope.close()
+            span: "InstanaSpan" = trace_config_ctx.span_context
+            if span:
+                span.record_exception(params.exception)
+                span.set_attribute("http.error", str(params.exception))
+                if span.is_recording():
+                    span.end()
+                trace_config_ctx = None
         except Exception:
-            logger.debug("stan_request_exception", exc_info=True)
-
-
-    @wrapt.patch_function_wrapper('aiohttp.client', 'ClientSession.__init__')
-    def init_with_instana(wrapped, instance, argv, kwargs):
+            logger.debug("aiohttp-client stan_request_exception error:", exc_info=True)
+
+    @wrapt.patch_function_wrapper("aiohttp.client", "ClientSession.__init__")
+    def init_with_instana(
+        wrapped: Callable[..., Awaitable["ClientSession"]],
+        instance: aiohttp.client.ClientSession,
+        args: Tuple[int, str, Tuple[object, ...]],
+        kwargs: Dict[str, Any],
+    ) -> object:
         instana_trace_config = aiohttp.TraceConfig()
         instana_trace_config.on_request_start.append(stan_request_start)
         instana_trace_config.on_request_end.append(stan_request_end)
         instana_trace_config.on_request_exception.append(stan_request_exception)
-        if 'trace_configs' in kwargs:
-            kwargs['trace_configs'].append(instana_trace_config)
+        if "trace_configs" in kwargs:
+            kwargs["trace_configs"].append(instana_trace_config)
         else:
-            kwargs['trace_configs'] = [instana_trace_config]
-
-        return wrapped(*argv, **kwargs)
+            kwargs["trace_configs"] = [instana_trace_config]
 
+        return wrapped(*args, **kwargs)
 
     logger.debug("Instrumenting aiohttp client")
 except ImportError:

src/instana/instrumentation/aiohttp/server.py

@@ -2,44 +2,58 @@
 # (c) Copyright Instana Inc. 2019
 
 
-import opentracing
+from typing import TYPE_CHECKING, Any, Awaitable, Callable, Dict, Tuple
+
 import wrapt
+from opentelemetry.semconv.trace import SpanAttributes
+
+from instana.log import logger
+from instana.propagators.format import Format
+from instana.singletons import agent, tracer
+from instana.util.secrets import strip_secrets_from_query
 
-from ...log import logger
-from ...singletons import agent, async_tracer
-from ...util.secrets import strip_secrets_from_query
+if TYPE_CHECKING:
+    from instana.span.span import InstanaSpan
 
 try:
     import aiohttp
-    import asyncio
-
     from aiohttp.web import middleware
 
+    if TYPE_CHECKING:
+        import aiohttp.web
 
     @middleware
-    async def stan_middleware(request, handler):
+    async def stan_middleware(
+        request: "aiohttp.web.Request",
+        handler: Callable[..., object],
+    ) -> Awaitable["aiohttp.web.Response"]:
         try:
-            ctx = async_tracer.extract(opentracing.Format.HTTP_HEADERS, request.headers)
-            request['scope'] = async_tracer.start_active_span('aiohttp-server', child_of=ctx)
-            scope = request['scope']
+            span_context = tracer.extract(Format.HTTP_HEADERS, request.headers)
+            span: "InstanaSpan" = tracer.start_span(
+                "aiohttp-server", span_context=span_context
+            )
+            request["span"] = span
 
             # Query param scrubbing
             url = str(request.url)
-            parts = url.split('?')
+            parts = url.split("?")
             if len(parts) > 1:
-                cleaned_qp = strip_secrets_from_query(parts[1],
-                                                      agent.options.secrets_matcher,
-                                                      agent.options.secrets_list)
-                scope.span.set_tag("http.params", cleaned_qp)
+                cleaned_qp = strip_secrets_from_query(
+                    parts[1], agent.options.secrets_matcher, agent.options.secrets_list
+                )
+                span.set_attribute("http.params", cleaned_qp)
 
-            scope.span.set_tag("http.url", parts[0])
-            scope.span.set_tag("http.method", request.method)
+            span.set_attribute(SpanAttributes.HTTP_URL, parts[0])
+            span.set_attribute(SpanAttributes.HTTP_METHOD, request.method)
 
             # Custom header tracking support
-            if agent.options.extra_http_headers is not None:
+            if agent.options.extra_http_headers:
                 for custom_header in agent.options.extra_http_headers:
                     if custom_header in request.headers:
-                        scope.span.set_tag("http.header.%s" % custom_header, request.headers[custom_header])
+                        span.set_attribute(
+                            "http.header.%s" % custom_header,
+                            request.headers[custom_header],
+                        )
 
             response = None
             try:
@@ -52,33 +66,38 @@ async def stan_middleware(request, handler):
             if response is not None:
                 # Mark 500 responses as errored
                 if 500 <= response.status:
-                    scope.span.mark_as_errored()
+                    span.mark_as_errored()
 
-                scope.span.set_tag("http.status_code", response.status)
-                async_tracer.inject(scope.span.context, opentracing.Format.HTTP_HEADERS, response.headers)
-                response.headers['Server-Timing'] = "intid;desc=%s" % scope.span.context.trace_id
+                span.set_attribute(SpanAttributes.HTTP_STATUS_CODE, response.status)
+                tracer.inject(span.context, Format.HTTP_HEADERS, response.headers)
+                response.headers["Server-Timing"] = (
+                    f"intid;desc={span.context.trace_id}"
+                )
 
             return response
         except Exception as exc:
-            logger.debug("aiohttp stan_middleware", exc_info=True)
-            if scope is not None:
-                scope.span.set_tag("http.status_code", 500)
-                scope.span.log_exception(exc)
+            logger.debug("aiohttp server stan_middleware:", exc_info=True)
+            if span:
+                span.set_attribute(SpanAttributes.HTTP_STATUS_CODE, 500)
+                span.record_exception(exc)
             raise
         finally:
-            if scope is not None:
-                scope.close()
-
-
-    @wrapt.patch_function_wrapper('aiohttp.web', 'Application.__init__')
-    def init_with_instana(wrapped, instance, argv, kwargs):
+            if span and span.is_recording():
+                span.end()
+
+    @wrapt.patch_function_wrapper("aiohttp.web", "Application.__init__")
+    def init_with_instana(
+        wrapped: Callable[..., "aiohttp.web.Application.__init__"],
+        instance: "aiohttp.web.Application",
+        args: Tuple[int, str, Tuple[Any, ...]],
+        kwargs: Dict[str, Any],
+    ) -> object:
         if "middlewares" in kwargs:
             kwargs["middlewares"].insert(0, stan_middleware)
         else:
             kwargs["middlewares"] = [stan_middleware]
 
-        return wrapped(*argv, **kwargs)
-
+        return wrapped(*args, **kwargs)
 
     logger.debug("Instrumenting aiohttp server")
 except ImportError:

src/instana/propagators/http_propagator.py

@@ -53,8 +53,8 @@ def inject_key_value(carrier, key, value):
             if span_context.suppression:
                 return
 
-            inject_key_value(carrier, self.HEADER_KEY_T, trace_id)
-            inject_key_value(carrier, self.HEADER_KEY_S, span_id)
+            inject_key_value(carrier, self.HEADER_KEY_T, str(trace_id))
+            inject_key_value(carrier, self.HEADER_KEY_S, str(span_id))
 
         except Exception:
             logger.debug("inject error:", exc_info=True)