New ASGI Instrumentation for FastAPI & Starlette

instana/__init__.py

@@ -125,6 +125,10 @@ def boot_agent():
         # Import & initialize instrumentation
         from .instrumentation.aws import lambda_inst
 
+        if sys.version_info >= (3, 6, 0):
+            from .instrumentation import fastapi_inst
+            from .instrumentation import starlette_inst
+
         if sys.version_info >= (3, 5, 3):
             from .instrumentation import asyncio
             from .instrumentation.aiohttp import client

instana/agent/host.py

@@ -102,6 +102,9 @@ def can_send(self):
         Are we in a state where we can send data?
         @return: Boolean
         """
+        if "INSTANA_TEST" in os.environ:
+            return True
+
         # Watch for pid change (fork)
         self.last_fork_check = datetime.now()
         current_pid = os.getpid()

instana/collector/aws_fargate.py

@@ -1,5 +1,5 @@
 """
-Snapshot & metrics collection for AWS Fargate
+AWS Fargate Collector: Manages the periodic collection of metrics & snapshot data
 """
 import os
 import json

instana/collector/aws_lambda.py

@@ -1,5 +1,5 @@
 """
-Snapshot & metrics collection for AWS Lambda
+AWS Lambda Collector: Manages the periodic collection of metrics & snapshot data
 """
 from ..log import logger
 from .base import BaseCollector

instana/collector/base.py

@@ -29,7 +29,15 @@ def __init__(self, agent):
         self.THREAD_NAME = "Instana Collector"
 
         # The Queue where we store finished spans before they are sent
-        self.span_queue = queue.Queue()
+        if env_is_test:
+            # Override span queue with a multiprocessing version
+            # The test suite runs background applications - some in background threads,
+            # others in background processes.  This multiprocess queue allows us to collect
+            # up spans from all sources.
+            import multiprocessing
+            self.span_queue = multiprocessing.Queue()
+        else:
+            self.span_queue = queue.Queue()
 
         # The Queue where we store finished profiles before they are sent
         self.profile_queue = queue.Queue()

instana/collector/host.py

@@ -1,11 +1,10 @@
 """
-Snapshot & metrics collection for AWS Fargate
+Host Collector: Manages the periodic collection of metrics & snapshot data
 """
 from time import time
 from ..log import logger
 from .base import BaseCollector
 from ..util import DictionaryOfStan
-from ..singletons import env_is_test
 from .helpers.runtime import RuntimeHelper
 
 

instana/instrumentation/asgi.py

@@ -0,0 +1,100 @@
+"""
+Instana ASGI Middleware
+"""
+import opentracing
+
+from ..log import logger
+from ..singletons import async_tracer, agent
+from ..util import strip_secrets_from_query
+
+class InstanaASGIMiddleware:
+    """
+    Instana ASGI Middleware
+    """
+    def __init__(self, app):
+        self.app = app
+
+    def _extract_custom_headers(self, span, headers):
+        try:
+            for custom_header in agent.options.extra_http_headers:
+                # Headers are in the following format: b'x-header-1'
+                for header_pair in headers:
+                    if header_pair[0].decode('utf-8').lower() == custom_header.lower():
+                        span.set_tag("http.%s" % custom_header, header_pair[1].decode('utf-8'))
+        except Exception:
+            logger.debug("extract_custom_headers: ", exc_info=True)
+
+    def _collect_kvs(self, scope, span):
+        try:
+            span.set_tag('http.path', scope.get('path'))
+            span.set_tag('http.method', scope.get('method'))
+
+            server = scope.get('server')
+            if isinstance(server, tuple):
+                span.set_tag('http.host', server[0])
+
+            query = scope.get('query_string')
+            if isinstance(query, (str, bytes)) and len(query):
+                if isinstance(query, bytes):
+                    query = query.decode('utf-8')
+                scrubbed_params = strip_secrets_from_query(query, agent.options.secrets_matcher, agent.options.secrets_list)
+                span.set_tag("http.params", scrubbed_params)
+
+            app = scope.get('app')
+            if app is not None and hasattr(app, 'routes'):
+                # Attempt to detect the Starlette routes registered.
+                # If Starlette isn't present, we harmlessly dump out.
+                from starlette.routing import Match
+                for route in scope['app'].routes:
+                    if route.matches(scope)[0] == Match.FULL:
+                        span.set_tag("http.path_tpl", route.path)
+        except Exception:
+            logger.debug("ASGI collect_kvs: ", exc_info=True)
+
+
+    async def __call__(self, scope, receive, send):
+        request_context = None
+
+        if scope["type"] not in ("http", "websocket"):
+            await self.app(scope, receive, send)
+            return
+
+        request_headers = scope.get('headers')
+        if isinstance(request_headers, list):
+            request_context = async_tracer.extract(opentracing.Format.BINARY, request_headers)
+
+        async def send_wrapper(response):
+            span = async_tracer.active_span
+            if span is None:
+                await send(response)
+            else:
+                if response['type'] == 'http.response.start':
+                    try:
+                        status_code = response.get('status')
+                        if status_code is not None:
+                            if 500 <= int(status_code) <= 511:
+                                span.mark_as_errored()
+                            span.set_tag('http.status_code', status_code)
+
+                        headers = response.get('headers')
+                        if headers is not None:
+                            async_tracer.inject(span.context, opentracing.Format.BINARY, headers)
+                    except Exception:
+                        logger.debug("send_wrapper: ", exc_info=True)
+
+                try:
+                    await send(response)
+                except Exception as exc:
+                    span.log_exception(exc)
+                    raise
+
+        with async_tracer.start_active_span("asgi", child_of=request_context) as tracing_scope:
+            self._collect_kvs(scope, tracing_scope.span)
+            if 'headers' in scope and agent.options.extra_http_headers is not None:
+                self._extract_custom_headers(tracing_scope.span, scope['headers'])
+
+            try:
+                await self.app(scope, receive, send_wrapper)
+            except Exception as exc:
+                tracing_scope.span.log_exception(exc)
+                raise exc

instana/instrumentation/fastapi_inst.py

@@ -0,0 +1,33 @@
+"""
+Instrumentation for FastAPI
+https://fastapi.tiangolo.com/
+"""
+try:
+    import fastapi
+    import wrapt
+    import signal
+    import os
+
+    from ..log import logger
+    from ..util import running_in_gunicorn
+    from .asgi import InstanaASGIMiddleware
+    from starlette.middleware import Middleware
+
+    @wrapt.patch_function_wrapper('fastapi.applications', 'FastAPI.__init__')
+    def init_with_instana(wrapped, instance, args, kwargs):
+        middleware = kwargs.get('middleware')
+        if middleware is None:
+            kwargs['middleware'] = [Middleware(InstanaASGIMiddleware)]
+        elif isinstance(middleware, list):
+            middleware.append(Middleware(InstanaASGIMiddleware))
+
+        return wrapped(*args, **kwargs)
+
+    logger.debug("Instrumenting FastAPI")
+
+    # Reload GUnicorn when we are instrumenting an already running application
+    if "INSTANA_MAGIC" in os.environ and running_in_gunicorn():
+        os.kill(os.getpid(), signal.SIGHUP)
+
+except ImportError:
+    pass

instana/instrumentation/pyramid/tweens.py

@@ -12,12 +12,12 @@
 
 class InstanaTweenFactory(object):
     """A factory that provides Instana instrumentation tween for Pyramid apps"""
-    
+
     def __init__(self, handler, registry):
         self.handler = handler
 
     def __call__(self, request):
-        ctx = tracer.extract(ot.Format.HTTP_HEADERS, request.headers)        
+        ctx = tracer.extract(ot.Format.HTTP_HEADERS, dict(request.headers))
         scope = tracer.start_active_span('http', child_of=ctx)
 
         scope.span.set_tag(ext.SPAN_KIND, ext.SPAN_KIND_RPC_SERVER)
@@ -42,7 +42,7 @@ def __call__(self, request):
         response = None
         try:
             response = self.handler(request)
-            
+
             tracer.inject(scope.span.context, ot.Format.HTTP_HEADERS, response.headers)
             response.headers['Server-Timing'] = "intid;desc=%s" % scope.span.context.trace_id
         except HTTPException as e:
@@ -53,21 +53,21 @@ def __call__(self, request):
 
             # we need to explicitly populate the `message` tag with an error here
             # so that it's picked up from an SDK span
-            scope.span.set_tag("message", str(e))            
+            scope.span.set_tag("message", str(e))
             scope.span.log_exception(e)
-            
+
             logger.debug("Pyramid Instana tween", exc_info=True)
         finally:
             if response:
                 scope.span.set_tag("http.status", response.status_int)
-                
+
                 if 500 <= response.status_int <= 511:
                     if response.exception is not None:
                         message = str(response.exception)
                         scope.span.log_exception(response.exception)
                     else:
                         message = response.status
-                        
+
                     scope.span.set_tag("message", message)
                     scope.span.assure_errored()
 

instana/instrumentation/starlette_inst.py

@@ -0,0 +1,24 @@
+"""
+Instrumentation for Starlette
+https://www.starlette.io/
+"""
+try:
+    import starlette
+    import wrapt
+    from ..log import logger
+    from .asgi import InstanaASGIMiddleware
+    from starlette.middleware import Middleware
+
+    @wrapt.patch_function_wrapper('starlette.applications', 'Starlette.__init__')
+    def init_with_instana(wrapped, instance, args, kwargs):
+        middleware = kwargs.get('middleware')
+        if middleware is None:
+            kwargs['middleware'] = [Middleware(InstanaASGIMiddleware)]
+        elif isinstance(middleware, list):
+            middleware.append(Middleware(InstanaASGIMiddleware))
+
+        return wrapped(*args, **kwargs)
+
+    logger.debug("Instrumenting Starlette")
+except ImportError:
+    pass

instana/instrumentation/tornado/server.py

@@ -24,7 +24,9 @@
         def execute_with_instana(wrapped, instance, argv, kwargs):
             try:
                 with tracer_stack_context():
-                    ctx = tornado_tracer.extract(opentracing.Format.HTTP_HEADERS, instance.request.headers)
+                    ctx = None
+                    if hasattr(instance.request.headers, '__dict__') and '_dict' in instance.request.headers.__dict__:
+                        ctx = tornado_tracer.extract(opentracing.Format.HTTP_HEADERS, instance.request.headers.__dict__['_dict'])
                     scope = tornado_tracer.start_active_span('tornado-server', child_of=ctx)
 
                     # Query param scrubbing

instana/instrumentation/wsgi.py

@@ -0,0 +1,55 @@
+"""
+Instana WSGI Middleware
+"""
+import opentracing as ot
+import opentracing.ext.tags as tags
+
+from ..singletons import agent, tracer
+from ..util import strip_secrets_from_query
+
+
+class InstanaWSGIMiddleware(object):
+    """ Instana WSGI middleware """
+
+    def __init__(self, app):
+        self.app = app
+
+    def __call__(self, environ, start_response):
+        env = environ
+
+        def new_start_response(status, headers, exc_info=None):
+            """Modified start response with additional headers."""
+            tracer.inject(self.scope.span.context, ot.Format.HTTP_HEADERS, headers)
+            headers.append(('Server-Timing', "intid;desc=%s" % self.scope.span.context.trace_id))
+
+            res = start_response(status, headers, exc_info)
+
+            sc = status.split(' ')[0]
+            if 500 <= int(sc) <= 511:
+                self.scope.span.mark_as_errored()
+
+            self.scope.span.set_tag(tags.HTTP_STATUS_CODE, sc)
+            self.scope.close()
+            return res
+
+        ctx = tracer.extract(ot.Format.HTTP_HEADERS, env)
+        self.scope = tracer.start_active_span("wsgi", child_of=ctx)
+
+        if agent.options.extra_http_headers is not None:
+            for custom_header in agent.options.extra_http_headers:
+                # Headers are available in this format: HTTP_X_CAPTURE_THIS
+                wsgi_header = ('HTTP_' + custom_header.upper()).replace('-', '_')
+                if wsgi_header in env:
+                    self.scope.span.set_tag("http.%s" % custom_header, env[wsgi_header])
+
+        if 'PATH_INFO' in env:
+            self.scope.span.set_tag('http.path', env['PATH_INFO'])
+        if 'QUERY_STRING' in env and len(env['QUERY_STRING']):
+            scrubbed_params = strip_secrets_from_query(env['QUERY_STRING'], agent.options.secrets_matcher, agent.options.secrets_list)
+            self.scope.span.set_tag("http.params", scrubbed_params)
+        if 'REQUEST_METHOD' in env:
+            self.scope.span.set_tag(tags.HTTP_METHOD, env['REQUEST_METHOD'])
+        if 'HTTP_HOST' in env:
+            self.scope.span.set_tag("http.host", env['HTTP_HOST'])
+
+        return self.app(environ, new_start_response)