Add support for B3 headers

ocelotl · ocelotl · commit 04eaa423ef44 · 2019-10-21T17:07:29.000-06:00
Closes #78
diff --git a/examples/http/b3_context_in_headers.py b/examples/http/b3_context_in_headers.py
@@ -0,0 +1,184 @@
+"""Demonstrates a Trace distributed across multiple machines.
+
+A SpanContext's text representation is stored in the headers of an HTTP request.
+
+Runs two threads, starts a Trace in the client and passes the SpanContext to the server.
+"""
+
+import argparse
+import errno
+import socket
+import sys
+import threading
+
+try:
+    # For Python 3.0 and later
+    from urllib.request import (
+        Request,
+        urlopen,
+    )
+    from http.server import BaseHTTPRequestHandler, HTTPServer
+except ImportError:
+    # Fall back to Python 2
+    from urllib2 import (
+        Request,
+        urlopen,
+    )
+    from BaseHTTPServer import BaseHTTPRequestHandler, HTTPServer
+
+import opentracing
+import opentracing.ext.tags
+import lightstep
+
+from lightstep.propagation import LightStepFormat, B3
+
+
+class RemoteHandler(BaseHTTPRequestHandler):
+    """This handler receives the request from the client.
+    """
+
+    def do_GET(self):
+        server_span = before_answering_request(self, opentracing.tracer)
+        with opentracing.tracer.scope_manager.activate(server_span, True):
+            server_span.log_event('request received', self.path)
+
+            self.send_response(200)
+            self.send_header('Content-type', 'text/html')
+            self.end_headers()
+            self.wfile.write("Hello World!".encode("utf-8"))
+
+            server_span.log_event('prepared response', self.path)
+
+
+def before_sending_request(request):
+    """Context manager creates Span and encodes the span's SpanContext into request.
+    """
+    span = opentracing.tracer.start_span('Sending request')
+    span.set_tag('server.http.url', request.get_full_url())
+    try:
+        # Python 2
+        host = request.get_host()
+    except:
+        # Python 3
+        host = request.host
+
+    if host:
+        span.set_tag(opentracing.ext.tags.PEER_HOST_IPV4, host)
+
+    carrier_dict = {}
+    span.tracer.inject(span.context, B3, carrier_dict)
+    for k, v in carrier_dict.items():
+        request.add_header(k, v)
+    return span
+
+
+def before_answering_request(handler, tracer):
+    """Context manager creates a Span, using SpanContext encoded in handler if possible.
+    """
+    operation = 'handle_request:' + handler.path
+    carrier_dict = {}
+    for k, v in handler.headers.items():
+        carrier_dict[k] = v
+    extracted_context = tracer.extract(B3, carrier_dict)
+
+    span = None
+    if extracted_context:
+        span = tracer.start_span(
+            operation_name=operation,
+            child_of=extracted_context)
+    else:
+        print('ERROR: Context missing, starting new trace')
+        global _exit_code
+        _exit_code = errno.ENOMSG
+        span = tracer.start_span(operation_name=operation)
+        headers = ', '.join({k + '=' + v for k, v in handler.headers.items()})
+        span.log_event('extract_failed', headers)
+        print('Could not extract context from http headers: ' + headers)
+
+    host, port = handler.client_address
+    if host:
+        span.set_tag(opentracing.ext.tags.PEER_HOST_IPV4, host)
+    if port:
+        span.set_tag(opentracing.ext.tags.PEER_PORT, str(port))
+
+    return span
+
+
+def pick_unused_port():
+    """ Since we don't reserve the port, there's a chance it'll get grabed, but that's unlikely.
+    """
+    s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+    s.bind(('localhost', 0))
+    port = s.getsockname()[1]
+    s.close()
+    return port
+
+
+def lightstep_tracer_from_args():
+    """Initializes lightstep from the commandline args.
+    """
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--token', help='Your LightStep access token.',
+                        default='{your_access_token}')
+    parser.add_argument('--host', help='The LightStep reporting service host to contact.',
+                        default='collector.lightstep.com')
+    parser.add_argument('--port', help='The LightStep reporting service port.',
+                        type=int, default=443)
+    parser.add_argument('--no_tls', help='Disable TLS for reporting',
+                        dest="no_tls", action='store_true')
+    parser.add_argument('--component_name', help='The LightStep component name',
+                        default='TrivialExample')
+    args = parser.parse_args()
+
+    if args.no_tls:
+        collector_encryption = 'none'
+    else:
+        collector_encryption = 'tls'
+
+    return lightstep.Tracer(
+        component_name=args.component_name,
+        access_token=args.token,
+        collector_host=args.host,
+        collector_port=args.port,
+        collector_encryption=collector_encryption,
+    )
+
+
+if __name__ == '__main__':
+    with lightstep_tracer_from_args() as tracer:
+        opentracing.tracer = tracer
+        global _exit_code
+        _exit_code = 0
+
+        # Create a web server and define the handler to manage the incoming request
+        port_number = pick_unused_port()
+        server = HTTPServer(('', port_number), RemoteHandler)
+
+        try:
+            # Run the server in a separate thread.
+            server_thread = threading.Thread(target=server.serve_forever)
+            server_thread.start()
+            print('Started httpserver on port ', port_number)
+
+            # Prepare request in the client
+            url = 'http://localhost:{}'.format(port_number)
+            request = Request(url)
+            client_span = before_sending_request(request)
+            with opentracing.tracer.scope_manager.activate(client_span, True):
+                client_span.log_event('sending request', url)
+
+                # Send request to server
+                response = urlopen(request)
+
+                response_body = response.read()
+                client_span.log_event('server returned', {
+                    "code": response.code,
+                    "body": response_body,
+                })
+
+            print('Server returned ' + str(response.code) + ': ' + str(response_body))
+
+            sys.exit(_exit_code)
+
+        finally:
+            server.shutdown()
diff --git a/lightstep/b3_propagator.py b/lightstep/b3_propagator.py
@@ -0,0 +1,136 @@
+from warnings import warn
+from re import compile as re_compile, IGNORECASE
+from logging import getLogger
+
+from basictracer.propagator import Propagator
+from basictracer.context import SpanContext
+from opentracing import span
+from opentracing import SpanContext as OTSpanContext
+
+_LOG = getLogger(__name__)
+_SINGLE_HEADER = "b3"
+# Lower case is used here as the B3 specification recommends
+_TRACEID = "x-b3-traceid"
+_SPANID = "x-b3-spanid"
+_PARENTSPANID = "x-b3-parentspanid"
+_SAMPLED = "x-b3-sampled"
+_FLAGS = "x-b3-flags"
+
+
+class B3Propagator(Propagator):
+    """
+    Propagator for the B3 HTTP header format.
+
+    See: https://github.com/openzipkin/b3-propagation
+    """
+
+    def inject(self, span_context, carrier):
+
+        traceid = span_context.trace_id
+        spanid = span_context.span_id
+
+        baggage = span_context.baggage
+
+        parentspanid = baggage.pop(_PARENTSPANID, None)
+        if parentspanid is not None:
+            carrier[_PARENTSPANID] = parentspanid
+
+        flags = baggage.pop(_FLAGS, None)
+        if flags is not None:
+            carrier[self._FLAGS] = flags
+
+        sampled = baggage.pop(_SAMPLED, None)
+        if sampled is not None:
+            if flags == 1:
+                _LOG.warning(
+                    "x-b3-flags: 1 implies x-b3-sampled: 1, not sending "
+                    "the value of x-b3-sampled"
+                )
+            else:
+                if sampled in [True, False]:
+                    warn(
+                        "The value of x-b3-sampled should "
+                        "be {} instead of {}".format(
+                            int(sampled), sampled
+                        )
+                    )
+                carrier[_SAMPLED] = int(sampled)
+
+        if (sampled is None and flags is None) and (
+            traceid is None or spanid is None
+        ):
+            warn(
+                "If not propagating only the sampling state, traceid and "
+                "spanid must be defined"
+            )
+
+        carrier.update(baggage)
+
+        if traceid is not None:
+            carrier[_TRACEID] = format(traceid, "032x")
+        if spanid is not None:
+            carrier[_SPANID] = format(spanid, "016x")
+
+    def extract(self, carrier):
+
+        case_insensitive_carrier = {}
+        for key, value in carrier.items():
+            for b3_key in [
+                _SINGLE_HEADER,
+                _TRACEID,
+                _SPANID,
+                _PARENTSPANID,
+                _SAMPLED,
+                _FLAGS,
+            ]:
+                if re_compile(key, IGNORECASE).match(b3_key):
+                    case_insensitive_carrier[b3_key] = value
+                else:
+                    case_insensitive_carrier[key] = value
+
+        carrier = case_insensitive_carrier
+        baggage = {}
+
+        if _SINGLE_HEADER in carrier.keys():
+            fields = carrier[_SINGLE_HEADER].split("-", 4)
+            len_fields = len(fields)
+            if len(fields) == 1:
+                sampled = fields[0]
+            elif len(fields) == 2:
+                traceid, spanid = fields
+            elif len(fields) == 3:
+                traceid, spanid, sampled = fields
+            else:
+                traceid, spanid, sampled, parent_spanid = fields
+                baggage[_PARENTSPANID] = parent_spanid
+            if sampled == "d":
+                baggage[_FLAGS] = 1
+            else:
+                baggage[_SAMPLED] = sampled
+        else:
+            traceid = carrier.pop(_TRACEID, None)
+            spanid = carrier.pop(_SPANID, None)
+            parentspanid = carrier.pop(_PARENTSPANID, None)
+            sampled = carrier.pop(_SAMPLED, None)
+            flags = carrier.pop(_FLAGS, None)
+
+            if parentspanid is not None:
+                baggage[_PARENTSPANID] = parentspanid
+            if flags == 1:
+                baggage[_FLAGS] = flags
+                if sampled is not None:
+                    warn(
+                        "x-b3-flags: 1 implies x-b3-sampled: 1, ignoring "
+                        "the received value of x-b3-sampled"
+                    )
+            elif sampled is not None:
+                baggage[_SAMPLED] = sampled
+            baggage.update(carrier)
+            if baggage == OTSpanContext.EMPTY_BAGGAGE:
+                baggage = None
+        return SpanContext(
+            # traceid and spanid are encoded in hex, so thet must be encoded
+            trace_id=int(traceid, 16),
+            span_id=int(spanid, 16),
+            baggage=baggage
+        )
diff --git a/lightstep/propagation.py b/lightstep/propagation.py
@@ -4,6 +4,10 @@
 
 from __future__ import absolute_import
 
+# The B3 format represents SpanContexts in B3 format.
+# https://github.com/openzipkin/b3-propagation
+B3 = "b3"
+
 
 class LightStepFormat(object):
     """A namespace for lightstep supported carrier formats.
diff --git a/lightstep/tracer.py b/lightstep/tracer.py
@@ -13,7 +13,8 @@
 from opentracing import Format
 
 from lightstep.lightstep_binary_propagator import LightStepBinaryPropagator
-from lightstep.propagation import LightStepFormat
+from lightstep.b3_propagator import B3Propagator
+from lightstep.propagation import LightStepFormat, B3
 from .recorder import Recorder
 
 
@@ -75,6 +76,7 @@ def __init__(self, enable_binary_format, recorder, scope_manager):
         super(_LightstepTracer, self).__init__(recorder, scope_manager=scope_manager)
         self.register_propagator(Format.TEXT_MAP, TextPropagator())
         self.register_propagator(Format.HTTP_HEADERS, TextPropagator())
+        self.register_propagator(B3, B3Propagator())
         if enable_binary_format:
             # We do this import lazily because protobuf versioning issues
             # can cause process-level failure at import time.
diff --git a/tests/b3_propagator_test.py b/tests/b3_propagator_test.py
@@ -0,0 +1,45 @@
+import unittest
+
+import lightstep
+from lightstep.propagation import B3
+
+
+class B3PropagatorTest(unittest.TestCase):
+    def setUp(self):
+        self._tracer = lightstep.Tracer(
+                periodic_flush_seconds=0,
+                collector_host='localhost')
+
+    def tracer(self):
+        return self._tracer
+
+    def tearDown(self):
+        self._tracer.flush()
+
+    def testInjectExtract(self):
+        carrier = {}
+        span = self.tracer().start_span('Sending request')
+
+        span.set_baggage_item('checked', 'baggage')
+
+        self.tracer().inject(span.context, B3, carrier)
+
+        result = self.tracer().extract(B3, carrier)
+        self.assertEqual(span.context.span_id, result.span_id)
+        self.assertEqual(span.context.trace_id, result.trace_id)
+        self.assertEqual(span.context.baggage, result.baggage)
+        self.assertEqual(span.context.sampled, result.sampled)
+
+    def single_header_extraction(self):
+        # FIXME complete this test case
+        input_ = 1
+
+        result = self.tracer().extract(B3, bytearray(input, 'utf-8'))
+        self.assertEqual(6397081719746291766, result.span_id)
+        self.assertEqual(506100417967962170, result.trace_id)
+        self.assertEqual(True, result.sampled)
+        self.assertDictEqual(result.baggage, {"checked" : "baggage"})
+
+
+if __name__ == '__main__':
+    unittest.main()
