Implement Binary codec (#304)

* Implement Binary codec

Implements extract and inject methods of the Binary codec to be able
to inject/extract span context to/from bytearray.

Signed-off-by: George Leman <[email protected]>

* Add test for binary codec compatibility with goclient

Signed-off-by: George Leman <[email protected]>

* Remove extra baggage items from binary codec test

Signed-off-by: George Leman <[email protected]>

Co-authored-by: Yuri Shkuro <[email protected]>

Author: cedy

jaeger_client/codecs.py

@@ -14,6 +14,8 @@
 
 from __future__ import absolute_import
 
+import struct
+
 from opentracing import (
     InvalidCarrierException,
     SpanContextCorruptedException,
@@ -129,19 +131,78 @@ def _parse_baggage_header(self, header, baggage):
 
 class BinaryCodec(Codec):
     """
-    BinaryCodec is a no-op.
-
+    Implements inject/extract of SpanContext to/from binary that compatible with golang
+    implementation
+    https://github.com/jaegertracing/jaeger-client-go/blob/master/propagation.go#L177-L290
+    Supports propagation of trace_id, span_id, flags and baggage
     """
     def inject(self, span_context, carrier):
         if not isinstance(carrier, bytearray):
             raise InvalidCarrierException('carrier not a bytearray')
-        pass  # TODO binary encoding not implemented
+        # check if we have 128 bit trace_id, break it into two 64 units
+        max_int64 = 0xFFFFFFFFFFFFFFFF
+        if span_context.trace_id > max_int64:
+            high = (span_context.trace_id >> 64) & max_int64
+            low = span_context.trace_id & max_int64
+        else:
+            high = 0
+            low = span_context.trace_id
+        carrier += struct.pack('>QQQQBI', high, low, span_context.span_id or 0,
+                               span_context.parent_id or 0, span_context.flags,
+                               len(span_context.baggage))
+
+        for k, v in span_context.baggage.items():
+            carrier += self._pack_baggage_item(k, v)
 
     def extract(self, carrier):
         if not isinstance(carrier, bytearray):
             raise InvalidCarrierException('carrier not a bytearray')
-        # TODO binary encoding not implemented
-        return None
+        baggage = {}
+        high_trace_id, low_trace_id, span_id, parent_id, flags, baggage_count = \
+            struct.unpack('>QQQQBI', carrier[:37])
+        # if high_trace_id isn't 0, then we are dealing with 128bit trace id integer,
+        # therefore unpack into 1 number
+        if high_trace_id:
+            trace_id = (high_trace_id << 64) | low_trace_id
+        else:
+            trace_id = low_trace_id
+
+        if baggage_count != 0:
+            baggage_data = carrier[37:]
+            for _ in range(baggage_count):
+                key, value, bytes_read = self._unpack_baggage_item(baggage_data)
+                baggage[key] = value
+                baggage_data = baggage_data[bytes_read:]
+
+        return SpanContext(trace_id=trace_id, span_id=span_id,
+                           parent_id=parent_id, flags=flags, baggage=baggage)
+
+    def _pack_baggage_item(self, key, value):
+        baggage = bytearray()
+        if not isinstance(key, bytes):
+            key = key.encode('utf-8')
+        baggage += struct.pack('>I', len(key))
+        baggage += key
+
+        if not isinstance(value, bytes):
+            value = value.encode('utf-8')
+        baggage += struct.pack('>I', len(value))
+        baggage += value
+        return baggage
+
+    def _unpack_baggage_item(self, baggage):
+        bytes_read = 0
+        key, b_read = self._read_kv(baggage)
+        bytes_read += b_read
+        value, b_read = self._read_kv(baggage[bytes_read:])
+        bytes_read += b_read
+        return key, value, bytes_read
+
+    def _read_kv(self, data):
+        data_len = struct.unpack('>i', data[:4])[0]
+        data_value = struct.unpack('>' + 'c' * data_len, data[4:4 + data_len])
+        bytes_read = 4 + data_len
+        return b''.join(data_value).decode('utf-8'), bytes_read
 
 
 def span_context_to_string(trace_id, span_id, parent_id, flags):

tests/test_codecs.py

@@ -384,6 +384,97 @@ def test_binary_codec(self):
         with self.assertRaises(InvalidCarrierException):
             codec.extract({})
 
+        tracer = Tracer(
+            service_name='test',
+            reporter=InMemoryReporter(),
+            sampler=ConstSampler(True),
+        )
+        baggage = {'baggage_1': u'data',
+                   u'baggage_2': 'foobar',
+                   'baggage_3': '\x00\x01\x09\xff',
+                   u'baggage_4': u'\U0001F47E'}
+
+        span_context = SpanContext(trace_id=260817200211625699950706086749966912306, span_id=567890,
+                                   parent_id=1234567890, flags=1,
+                                   baggage=baggage)
+
+        carrier = bytearray()
+        tracer.inject(span_context, Format.BINARY, carrier)
+        assert len(carrier) != 0
+
+        extracted_span_context = tracer.extract(Format.BINARY, carrier)
+        assert extracted_span_context.trace_id == span_context.trace_id
+        assert extracted_span_context.span_id == span_context.span_id
+        assert extracted_span_context.parent_id == span_context.parent_id
+        assert extracted_span_context.flags == span_context.flags
+        assert extracted_span_context.baggage == span_context.baggage
+
+    def test_binary_codec_extract_compatibility_with_golang_client(self):
+        tracer = Tracer(
+            service_name='test',
+            reporter=InMemoryReporter(),
+            sampler=ConstSampler(True),
+        )
+        tests = {
+            b'\x00\x00\x00\x00\x00\x00\x00\x00u\x18\xa9\x13\xa0\xd2\xaf4u\x18\xa9\x13\xa0\xd2\xaf4'
+            b'\x00\x00\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00':
+                {'trace_id_high': 0,
+                 'trace_id_low': 8437679803646258996,
+                 'span_id': 8437679803646258996,
+                 'parent_id': None,
+                 'flags': 1,
+                 'baggage_count': 0,
+                 'baggage': {}},
+            b'K2\x88\x8b\x8f\xb5\x96\xe9+\xc6\xe6\xf5\x9d\xed\x8a\xd0+\xc6\xe6\xf5\x9d\xed\x8a\xd0'
+            b'\x00\x00\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00':
+                {'trace_id_high': 5418543434673002217,
+                 'trace_id_low': 3154462531610577616,
+                 'span_id': 3154462531610577616,
+                 'parent_id': None,
+                 'flags': 1,
+                 'baggage_count': 0,
+                 'baggage': {}},
+            b'd\xb7^Y\x1afI\x0bi\xe4lc`\x1e\xbep[\x0fw\xc8\x87\xfd\xb2Ti\xe4lc`\x1e\xbep\x01\x00'
+            b'\x00\x00\x00':
+                {'trace_id_high': 7257373061318854923,
+                 'trace_id_low': 7630342842742652528,
+                 'span_id': 6561594885260816980,
+                 'parent_id': 7630342842742652528,
+                 'flags': 1,
+                 'baggage_count': 0,
+                 'baggage': {}},
+            b'a]\x85\xe0\xe0\x06\xd5[6k\x9d\x86\xaa\xbc\\\x8f#c\x06\x80jV\xdf\x826k\x9d\x86\xaa\xbc'
+            b'\\\x8f\x01\x00\x00\x00\x01\x00\x00\x00\x07key_one\x00\x00\x00\tvalue_one':
+                {'trace_id_high': 7015910995390813531,
+                 'trace_id_low': 3921401102271798415,
+                 'span_id': 2549888962631491458,
+                 'parent_id': 3921401102271798415,
+                 'flags': 1,
+                 'baggage_count': 1,
+                 'baggage': {'key_one': 'value_one'}
+                 },
+        }
+
+        for span_context_serialized, expected in tests.items():
+            span_context = tracer.extract(Format.BINARY, bytearray(span_context_serialized))
+            # because python supports 128bit number as one number and go splits it in two 64 bit
+            # numbers, we need to split python number to compare it properly to go implementation
+            max_int64 = 0xFFFFFFFFFFFFFFFF
+            trace_id_high = (span_context.trace_id >> 64) & max_int64
+            trace_id_low = span_context.trace_id & max_int64
+
+            assert trace_id_high == expected['trace_id_high']
+            assert trace_id_low == expected['trace_id_low']
+            assert span_context.span_id == expected['span_id']
+            assert span_context.parent_id == expected['parent_id']
+            assert span_context.flags == expected['flags']
+            assert len(span_context.baggage) == expected['baggage_count']
+            assert span_context.baggage == expected['baggage']
+
+            carrier = bytearray()
+            tracer.inject(span_context, Format.BINARY, carrier)
+            assert carrier == bytearray(span_context_serialized)
+
 
 def test_default_baggage_without_trace_id(tracer):
     _test_baggage_without_trace_id(