Support Custom Header collection (#99)

* Mr. Linter

* Store and collect custom headers in WSGI and Django middleware

* Assure params are there before reporting.

* Tests to validate custom header capture

* Use dict methods instead of our own

* a more precise function name

Author: pglombardo

instana/__init__.py

@@ -55,6 +55,7 @@ def load(module):
 
 import instana.singletons #noqa
 
+
 def load_instrumentation():
     if "INSTANA_DISABLE_AUTO_INSTR" not in os.environ:
         # Import & initialize instrumentation
@@ -63,6 +64,7 @@ def load_instrumentation():
         from .instrumentation import mysqlpython  # noqa
         from .instrumentation.django import middleware  # noqa
 
+
 if "INSTANA_MAGIC" in os.environ:
     # If we're being loaded into an already running process, then delay
     # instrumentation load.

instana/agent.py

@@ -47,6 +47,7 @@ class Agent(object):
     last_seen = None
     last_fork_check = None
     _boot_pid = os.getpid()
+    extra_headers = None
 
     def __init__(self):
         logger.debug("initializing agent")
@@ -165,7 +166,13 @@ def set_from(self, json_string):
         else:
             raw_json = json_string
 
-        self.from_ = From(**json.loads(raw_json))
+        res_data = json.loads(raw_json)
+
+        if "extraHeaders" in res_data:
+            self.extra_headers = res_data['extraHeaders']
+            logger.debug("Will also capture these custom headers: %s", self.extra_headers)
+
+        self.from_ = From(pid=res_data['pid'], agentUuid=res_data['agentUuid'])
 
     def reset(self):
         self.last_seen = None

instana/instrumentation/django/middleware.py

@@ -33,10 +33,17 @@ def process_request(self, request):
 
             request.iscope = tracer.start_active_span('django', child_of=ctx)
 
+            if agent.extra_headers is not None:
+                for custom_header in agent.extra_headers:
+                    # Headers are available in this format: HTTP_X_CAPTURE_THIS
+                    django_header = ('HTTP_' + custom_header.upper()).replace('-', '_')
+                    if django_header in env:
+                        request.iscope.span.set_tag("http.%s" % custom_header, env[django_header])
+
             request.iscope.span.set_tag(ext.HTTP_METHOD, request.method)
             if 'PATH_INFO' in env:
                 request.iscope.span.set_tag(ext.HTTP_URL, env['PATH_INFO'])
-            if 'QUERY_STRING' in env:
+            if 'QUERY_STRING' in env and len(env['QUERY_STRING']):
                 request.iscope.span.set_tag("http.params", env['QUERY_STRING'])
             if 'HTTP_HOST' in env:
                 request.iscope.span.set_tag("http.host", env['HTTP_HOST'])

instana/recorder.py

@@ -44,9 +44,11 @@ def report_spans(self):
         """ Periodically report the queued spans """
         logger.debug("Span reporting thread is now alive")
         while 1:
-            if self.queue.qsize() > 0 and instana.singletons.agent.can_send():
+            queue_size = self.queue.qsize()
+            if queue_size > 0 and instana.singletons.agent.can_send():
                 url = instana.singletons.agent.make_url(AGENT_TRACES_URL)
                 instana.singletons.agent.request(url, "POST", self.queued_spans())
+                logger.debug("reported %d spans" % queue_size)
             time.sleep(1)
 
     def queue_size(self):
@@ -91,20 +93,20 @@ def build_registered_span(self, span):
                                       logs=self.collect_logs(span)))
 
         if span.operation_name in self.http_spans:
-            data.http = HttpData(host=self.get_host_name(span),
-                                 url=self.get_string_tag(span, ext.HTTP_URL),
-                                 method=self.get_string_tag(span, ext.HTTP_METHOD),
-                                 status=self.get_tag(span, ext.HTTP_STATUS_CODE),
-                                 error=self.get_tag(span, 'http.error'))
+            data.http = HttpData(host=self.get_http_host_name(span),
+                                 url=span.tags.pop(ext.HTTP_URL, ""),
+                                 method=span.tags.pop(ext.HTTP_METHOD, ""),
+                                 status=span.tags.pop(ext.HTTP_STATUS_CODE, None),
+                                 error=span.tags.pop('http.error', None))
 
         if span.operation_name == "soap":
-            data.soap = SoapData(action=self.get_tag(span, 'soap.action'))
+            data.soap = SoapData(action=span.tags.pop('soap.action', None))
 
         if span.operation_name == "mysql":
-            data.mysql = MySQLData(host=self.get_tag(span, 'host'),
-                                   db=self.get_tag(span, ext.DATABASE_INSTANCE),
-                                   user=self.get_tag(span, ext.DATABASE_USER),
-                                   stmt=self.get_tag(span, ext.DATABASE_STATEMENT))
+            data.mysql = MySQLData(host=span.tags.pop('host', None),
+                                   db=span.tags.pop(ext.DATABASE_INSTANCE, None),
+                                   user=span.tags.pop(ext.DATABASE_USER, None),
+                                   stmt=span.tags.pop(ext.DATABASE_STATEMENT, None))
             if len(data.custom.logs.keys()):
                 tskey = list(data.custom.logs.keys())[0]
                 data.mysql.error = data.custom.logs[tskey]['message']
@@ -121,8 +123,8 @@ def build_registered_span(self, span):
                              f=entityFrom,
                              data=data)
 
-        error = self.get_tag(span, "error", False)
-        ec = self.get_tag(span, "ec", None)
+        error = span.tags.pop("error", False)
+        ec = span.tags.pop("ec", None)
 
         if error and ec:
             json_span.error = error
@@ -154,30 +156,17 @@ def build_sdk_span(self, span):
                              f=entityFrom,
                              data=data)
 
-        error = self.get_tag(span, "error", False)
-        ec = self.get_tag(span, "ec", None)
+        error = span.tags.pop("error", False)
+        ec = span.tags.pop("ec", None)
 
         if error and ec:
             json_span.error = error
             json_span.ec = ec
 
         return json_span
 
-    def get_tag(self, span, tag, default=None):
-        if tag in span.tags:
-            return span.tags[tag]
-
-        return default
-
-    def get_string_tag(self, span, tag):
-        ret = self.get_tag(span, tag)
-        if not ret:
-            return ""
-
-        return ret
-
-    def get_host_name(self, span):
-        h = self.get_string_tag(span, "http.host")
+    def get_http_host_name(self, span):
+        h = span.tags.pop("http.host", "")
         if len(h) > 0:
             return h
 

instana/wsgi.py

@@ -3,7 +3,7 @@
 import opentracing as ot
 import opentracing.ext.tags as tags
 
-from .singletons import tracer
+from .singletons import agent, tracer
 
 
 class iWSGIMiddleware(object):
@@ -37,9 +37,17 @@ def new_start_response(status, headers, exc_info=None):
 
         self.scope = tracer.start_active_span("wsgi", child_of=ctx)
 
+        if agent.extra_headers is not None:
+            for custom_header in agent.extra_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(tags.HTTP_URL, env['PATH_INFO'])
-        if 'QUERY_STRING' in env:
+        if 'QUERY_STRING' in env and len(env['QUERY_STRING']):
             self.scope.span.set_tag("http.params", env['QUERY_STRING'])
         if 'REQUEST_METHOD' in env:
             self.scope.span.set_tag(tags.HTTP_METHOD, env['REQUEST_METHOD'])

tests/test_django.py

@@ -5,7 +5,7 @@
 from django.contrib.staticfiles.testing import StaticLiveServerTestCase
 from nose.tools import assert_equals
 
-from instana.singletons import tracer
+from instana.singletons import agent, tracer
 
 from .apps.app_django import INSTALLED_APPS
 
@@ -123,3 +123,48 @@ def test_complex_request(self):
         assert_equals('/complex', django_span.data.http.url)
         assert_equals('GET', django_span.data.http.method)
         assert_equals(200, django_span.data.http.status)
+
+    def test_custom_header_capture(self):
+        # Hack together a manual custom headers list
+        agent.extra_headers = [u'X-Capture-This', u'X-Capture-That']
+
+        request_headers = {}
+        request_headers['X-Capture-This'] = 'this'
+        request_headers['X-Capture-That'] = 'that'
+
+        with tracer.start_active_span('test'):
+            response = self.http.request('GET', self.live_server_url + '/', headers=request_headers)
+            # response = self.client.get('/')
+
+        assert_equals(response.status, 200)
+
+        spans = self.recorder.queued_spans()
+        assert_equals(3, len(spans))
+
+        test_span = spans[2]
+        urllib3_span = spans[1]
+        django_span = spans[0]
+
+        # import ipdb; ipdb.set_trace()
+
+        assert_equals("test", test_span.data.sdk.name)
+        assert_equals("urllib3", urllib3_span.n)
+        assert_equals("django", django_span.n)
+
+        assert_equals(test_span.t, urllib3_span.t)
+        assert_equals(urllib3_span.t, django_span.t)
+
+        assert_equals(urllib3_span.p, test_span.s)
+        assert_equals(django_span.p, urllib3_span.s)
+
+        assert_equals(None, django_span.error)
+        assert_equals(None, django_span.ec)
+
+        assert_equals('/', django_span.data.http.url)
+        assert_equals('GET', django_span.data.http.method)
+        assert_equals(200, django_span.data.http.status)
+
+        assert_equals(True, "http.X-Capture-This" in django_span.data.custom.__dict__['tags'])
+        assert_equals("this", django_span.data.custom.__dict__['tags']["http.X-Capture-This"])
+        assert_equals(True, "http.X-Capture-That" in django_span.data.custom.__dict__['tags'])
+        assert_equals("that", django_span.data.custom.__dict__['tags']["http.X-Capture-That"])

tests/test_wsgi.py

@@ -4,7 +4,7 @@
 import unittest
 
 import urllib3
-from instana.singletons import tracer
+from instana.singletons import agent, tracer
 
 
 class TestWSGI(unittest.TestCase):
@@ -116,3 +116,55 @@ def test_complex_request(self):
         self.assertEqual('GET', wsgi_span.data.http.method)
         self.assertEqual('200', wsgi_span.data.http.status)
         self.assertIsNone(wsgi_span.data.http.error)
+
+    def test_custom_header_capture(self):
+        # Hack together a manual custom headers list
+        agent.extra_headers = [u'X-Capture-This', u'X-Capture-That']
+
+        request_headers = {}
+        request_headers['X-Capture-This'] = 'this'
+        request_headers['X-Capture-That'] = 'that'
+
+        with tracer.start_active_span('test'):
+            response = self.http.request('GET', 'http://127.0.0.1:5000/', headers=request_headers)
+
+        spans = self.recorder.queued_spans()
+
+        self.assertEqual(3, len(spans))
+        self.assertIsNone(tracer.active_span)
+
+        wsgi_span = spans[0]
+        urllib3_span = spans[1]
+        test_span = spans[2]
+
+        assert(response)
+        self.assertEqual(200, response.status)
+
+        # Same traceId
+        self.assertEqual(test_span.t, urllib3_span.t)
+        self.assertEqual(urllib3_span.t, wsgi_span.t)
+
+        # Parent relationships
+        self.assertEqual(urllib3_span.p, test_span.s)
+        self.assertEqual(wsgi_span.p, urllib3_span.s)
+
+        # Error logging
+        self.assertFalse(test_span.error)
+        self.assertIsNone(test_span.ec)
+        self.assertFalse(urllib3_span.error)
+        self.assertIsNone(urllib3_span.ec)
+        self.assertFalse(wsgi_span.error)
+        self.assertIsNone(wsgi_span.ec)
+
+        # wsgi
+        self.assertEqual("wsgi", wsgi_span.n)
+        self.assertEqual('127.0.0.1:5000', wsgi_span.data.http.host)
+        self.assertEqual('/', wsgi_span.data.http.url)
+        self.assertEqual('GET', wsgi_span.data.http.method)
+        self.assertEqual('200', wsgi_span.data.http.status)
+        self.assertIsNone(wsgi_span.data.http.error)
+
+        self.assertEqual(True, "http.X-Capture-This" in wsgi_span.data.custom.__dict__['tags'])
+        self.assertEqual("this", wsgi_span.data.custom.__dict__['tags']["http.X-Capture-This"])
+        self.assertEqual(True, "http.X-Capture-That" in wsgi_span.data.custom.__dict__['tags'])
+        self.assertEqual("that", wsgi_span.data.custom.__dict__['tags']["http.X-Capture-That"])