New ASGI Instrumentation for FastAPI & Starlette (#282)

* Load GRPC server on-demand only

* Tests: Support launching threads with args

* Python 2 compatibility

* Avoid using lambda for test suite: multiprocessing & pickle

* New test suite dependencies

* Add support & tests for byte based context keys

* Set test env vars

* Reorg middleware exports

* When in test, we can always send

* Check for empty first

* Add support for byte based context headers

* New ASGI middleware for FastAPI & Starlette

* Add ASGI as registered span

* Tests: FastAPI background server & tests

* Check potential byte based value

* Better custom header capture

* Set custom header options before launching background process

* Vanilla, synthetic and custom header capture tests

* Tests: Secret scrubbing

* Path Templates support & tests

* FastAPI Test server cleanup

* Test path templates always

* Starlette tests: background server and tests

* Version limiters for CircleCI

* Fix version string

* Version limit uvicorn

* Fix Python 2 support

* Dont use __getitem__: Final answer

* Update opentracing version

* Update basictracer version

* When in test, use a multiprocess queue

* Code comments

* Simplify multiprocess launching

* In tests, pause for spans to land

* Break up aiohttp tests

* More robust extraction

* Assure package loaded

* Remove redundant sleeps

* Pause to let spans settle

* Starlette requires aiofiles

* Better conversion for Tornado headers class

* Unify, cleanup & fix context propagators

* Fix binary propagator tests

* Safeties, maturities and updated tests

* Cleanup: remove debug & pydoc

* Cleanup; Remove debug logs

* Maturity Refactoring

* Breakout gunicorn detection

* Make log package independent to avoid circular import issues

* Reload gunicorn on AutoTrace

* New Test Helper: launch_traced_request

Author: Peter Giacomo Lombardo

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