Add 128bit trace_id support (#230)

* Add 128bit trace_id support

- Creates a new property in Config 'generate_128bit_trace_id' if set
  to True then 128bit trace_ids are generated
  - Adds support for environment variable JAEGER_TRACEID_128BIT. If
  the value is 'true' then 128bit trace_ids are generated
- Tracer takes variable generate_128bit_trace_id, if the value is True
  then max_trace_id_bits is set to constants._max_trace_id_bits
  - Adds new function _random_id(bitsize)
  - Deprecate random_id() in favor of _random_id(bitsize)
- Deprecate MAX_ID_BITS in favor of _max_trace_id_bits and _max_id_bits
- Set low and high bits of traceId correctly
  - Uses _id_to_high() and _id_to_low() to do that

Signed-off-by: Bhavin Gandhi <[email protected]>

* Add tests for 128bit trace_id support

Signed-off-by: Bhavin Gandhi <[email protected]>

* Update ProbabilisticSampler to trim the trace_id to 64bits

Signed-off-by: Bhavin Gandhi <[email protected]>

* Update test_round_trip to test with 128bit trace_id

- Uses itertools.product to create a list with all possible
  combinations for tracer i.e. tracer with 64bit trace_id and tracer
  with 128bit trace_id

Signed-off-by: Bhavin Gandhi <[email protected]>

* Add test_inject_with_128bit_trace_id to test_codecs.py

- This tests the behavior of inject for different codes when it
  receives 64bit/128bit trace_id on wire when tracer is configured to
  genereate 64/128bit trace_ids
- Also tests the behavior of inject when tracer is configured to
  generate 64/128bit trace_ids

Signed-off-by: Bhavin Gandhi <[email protected]>

Author: bhavin192

jaeger_client/config.py

@@ -122,6 +122,7 @@ def _validate_config(self, config):
                         'reporter_flush_interval',
                         'sampling_refresh_interval',
                         'trace_id_header',
+                        'generate_128bit_trace_id',
                         'baggage_header_prefix',
                         'service_name',
                         'throttler']
@@ -166,6 +167,16 @@ def trace_id_header(self):
         """
         return self.config.get('trace_id_header', TRACE_ID_HEADER)
 
+    @property
+    def generate_128bit_trace_id(self):
+        """
+        :return: Returns boolean value to indicate if 128bit trace_id
+        generation is enabled
+        """
+        if 'generate_128bit_trace_id' in self.config:
+            return get_boolean(self.config['generate_128bit_trace_id'], False)
+        return os.getenv('JAEGER_TRACEID_128BIT') == 'true'
+
     @property
     def baggage_header_prefix(self):
         """
@@ -395,6 +406,7 @@ def create_tracer(self, reporter, sampler, throttler=None):
             sampler=sampler,
             metrics_factory=self._metrics_factory,
             trace_id_header=self.trace_id_header,
+            generate_128bit_trace_id=self.generate_128bit_trace_id,
             baggage_header_prefix=self.baggage_header_prefix,
             debug_id_header=self.debug_id_header,
             tags=self.tags,

jaeger_client/constants.py

@@ -18,9 +18,15 @@
 
 import six
 
-# Max number of bits to use when generating random ID
+# DEPRECATED: max number of bits to use when generating random ID
 MAX_ID_BITS = 64
 
+# Max number of bits allowed to use when generating Trace ID
+_max_trace_id_bits = 128
+
+# Max number of bits to use when generating random ID
+_max_id_bits = 64
+
 # How often remotely controlled sampler polls for sampling strategy
 DEFAULT_SAMPLING_INTERVAL = 60
 

jaeger_client/sampler.py

@@ -22,7 +22,7 @@
 from threading import Lock
 from tornado.ioloop import PeriodicCallback
 from .constants import (
-    MAX_ID_BITS,
+    _max_id_bits,
     DEFAULT_SAMPLING_INTERVAL,
     SAMPLER_TYPE_CONST,
     SAMPLER_TYPE_PROBABILISTIC,
@@ -121,10 +121,11 @@ def __init__(self, rate):
         )
         assert 0.0 <= rate <= 1.0, 'Sampling rate must be between 0.0 and 1.0'
         self.rate = rate
-        self.max_number = 1 << MAX_ID_BITS
+        self.max_number = 1 << _max_id_bits
         self.boundary = rate * self.max_number
 
     def is_sampled(self, trace_id, operation=''):
+        trace_id = trace_id & (self.max_number - 1)
         return trace_id < self.boundary, self._tags
 
     def close(self):

jaeger_client/thrift.py

@@ -26,6 +26,24 @@
     long = int
 
 
+def _id_to_low(big_id):
+    """
+    :param big_id: id in integer
+    :return: Returns the right most 64 bits of big_id
+    """
+    if big_id is not None:
+        return big_id & (_max_unsigned_id - 1)
+
+
+def _id_to_high(big_id):
+    """
+    :param big_id: id in integer
+    :return: Returns the left most 64 bits of big_id
+    """
+    if big_id is not None:
+        return (big_id >> 64) & (_max_unsigned_id - 1)
+
+
 def id_to_int(big_id):
     if big_id is None:
         return None
@@ -150,8 +168,8 @@ def make_jaeger_batch(spans, process):
     for span in spans:
         with span.update_lock:
             jaeger_span = ttypes.Span(
-                traceIdLow=id_to_int(span.trace_id),
-                traceIdHigh=0,
+                traceIdLow=id_to_int(_id_to_low(span.trace_id)),
+                traceIdHigh=id_to_int(_id_to_high(span.trace_id)),
                 spanId=id_to_int(span.span_id),
                 parentSpanId=id_to_int(span.parent_id) or 0,
                 operationName=span.operation_name,

jaeger_client/tracer.py

@@ -44,6 +44,7 @@ def __init__(
         self, service_name, reporter, sampler, metrics=None,
         metrics_factory=None,
         trace_id_header=constants.TRACE_ID_HEADER,
+        generate_128bit_trace_id=False,
         baggage_header_prefix=constants.BAGGAGE_HEADER_PREFIX,
         debug_id_header=constants.DEBUG_ID_HEADER_KEY,
         one_span_per_rpc=False, extra_codecs=None,
@@ -60,6 +61,8 @@ def __init__(
         self.debug_id_header = debug_id_header
         self.one_span_per_rpc = one_span_per_rpc
         self.max_tag_value_length = max_tag_value_length
+        self.max_trace_id_bits = constants._max_trace_id_bits if generate_128bit_trace_id \
+            else constants._max_id_bits
         self.codecs = {
             Format.TEXT_MAP: TextCodec(
                 url_encoding=False,
@@ -144,8 +147,8 @@ def start_span(self,
             tags.get(ext_tags.SPAN_KIND) == ext_tags.SPAN_KIND_RPC_SERVER
 
         if parent is None or not parent.has_trace:
-            trace_id = self.random_id()
-            span_id = trace_id
+            trace_id = self._random_id(self.max_trace_id_bits)
+            span_id = self._random_id(constants._max_id_bits)
             parent_id = None
             flags = 0
             baggage = None
@@ -170,7 +173,7 @@ def start_span(self,
                 span_id = parent.span_id
                 parent_id = parent.parent_id
             else:
-                span_id = self.random_id()
+                span_id = self._random_id(constants._max_id_bits)
                 parent_id = parent.span_id
             flags = parent.flags
             baggage = dict(parent.baggage)  # TODO do we need to clone?
@@ -238,8 +241,14 @@ def report_span(self, span):
         self.metrics.spans_finished(1)
 
     def random_id(self):
+        """
+        DEPRECATED: use _random_id() instead
+        """
         return self.random.getrandbits(constants.MAX_ID_BITS)
 
+    def _random_id(self, bitsize):
+        return self.random.getrandbits(bitsize)
+
     def is_debug_allowed(self, *args, **kwargs):
         if not self.throttler:
             return True

tests/test_codecs.py

@@ -16,6 +16,7 @@
 
 import unittest
 from collections import namedtuple
+from itertools import product
 import six
 
 import mock
@@ -28,6 +29,7 @@
 )
 from jaeger_client.config import Config
 from jaeger_client.reporter import InMemoryReporter
+from jaeger_client import constants
 from opentracing import Format
 from opentracing.propagation import (
     InvalidCarrierException,
@@ -82,6 +84,8 @@ def test_trace_context_from_to_string(self):
         tests = [
             [(256, 127, None, 1), '100:7f:0:1'],
             [(256, 127, 256, 0), '100:7f:100:0'],
+            [(0xffffffffffffffffffffffffffffffff, 127, 256, 0),
+             'ffffffffffffffffffffffffffffffff:7f:100:0'],
         ]
         for test in tests:
             ctx = test[0]
@@ -211,11 +215,11 @@ def test_context_from_large_ids(self):
         codec = TextCodec(trace_id_header='Trace_ID',
                           baggage_header_prefix='Trace-Attr-')
         headers = {
-            'Trace-ID': 'FFFFFFFFFFFFFFFF:FFFFFFFFFFFFFFFF:FFFFFFFFFFFFFFFF:1',
+            'Trace-ID': 'FFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFF:FFFFFFFFFFFFFFFF:FFFFFFFFFFFFFFFF:1',
         }
         context = codec.extract(headers)
-        assert context.trace_id == 0xFFFFFFFFFFFFFFFF
-        assert context.trace_id == (1 << 64) - 1
+        assert context.trace_id == 0xFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFF
+        assert context.trace_id == (1 << 128) - 1
         assert context.trace_id > 0
         assert context.span_id == 0xFFFFFFFFFFFFFFFF
         assert context.span_id == (1 << 64) - 1
@@ -437,11 +441,64 @@ def _test_baggage_without_trace_id(tracer, trace_id_header, baggage_header_prefi
     (ZipkinSpanFormat, {}),
 ])
 def test_round_trip(tracer, fmt, carrier):
-    span = tracer.start_span('test-%s' % fmt)
-    tracer.inject(span, fmt, carrier)
-    context = tracer.extract(fmt, carrier)
-    span2 = tracer.start_span('test-%s' % fmt, child_of=context)
-    assert span.trace_id == span2.trace_id
+    tracer_128bit = Tracer(
+        service_name='test',
+        reporter=InMemoryReporter(),
+        sampler=ConstSampler(True),
+        generate_128bit_trace_id=True)
+
+    for tracer1, tracer2 in product([tracer, tracer_128bit], repeat=2):
+        span = tracer1.start_span('test-%s' % fmt)
+        tracer1.inject(span, fmt, carrier)
+        context = tracer2.extract(fmt, carrier)
+        span2 = tracer2.start_span('test-%s' % fmt, child_of=context)
+        assert span.trace_id == span2.trace_id
+
+
+def _text_codec_to_trace_id_string(carrier):
+    return carrier[constants.TRACE_ID_HEADER].split(':')[0]
+
+
+def _zipkin_codec_to_trace_id_string(carrier):
+    return '{:x}'.format(carrier['trace_id'])
+
+
+@pytest.mark.parametrize('fmt,carrier,get_trace_id', [
+    (Format.TEXT_MAP, {}, _text_codec_to_trace_id_string),
+    (Format.HTTP_HEADERS, {}, _text_codec_to_trace_id_string),
+    (ZipkinSpanFormat, {}, _zipkin_codec_to_trace_id_string),
+])
+def test_inject_with_128bit_trace_id(tracer, fmt, carrier, get_trace_id):
+    tracer_128bit = Tracer(
+        service_name='test',
+        reporter=InMemoryReporter(),
+        sampler=ConstSampler(True),
+        generate_128bit_trace_id=True)
+
+    for tracer in [tracer, tracer_128bit]:
+        length = tracer.max_trace_id_bits / 4
+        trace_id = (1 << 64) - 1 if length == 16 else (1 << 128) - 1
+        ctx = SpanContext(trace_id=trace_id, span_id=127, parent_id=None,
+                          flags=1)
+        span = Span(ctx, operation_name='test-%s' % fmt, tracer=None, start_time=1)
+        tracer.inject(span, fmt, carrier)
+        assert len(get_trace_id(carrier)) == length
+
+        # test if the trace_id arrived on wire remains same even if
+        # the tracer is configured for 64bit ids or 128bit ids
+        ctx = SpanContext(trace_id=(1 << 128) - 1, span_id=127, parent_id=None,
+                          flags=0)
+        span = tracer.start_span('test-%s' % fmt, child_of=ctx)
+        carrier = dict()
+        tracer.inject(span, fmt, carrier)
+        assert len(get_trace_id(carrier)) == 32
+
+        ctx = SpanContext(trace_id=(1 << 64) - 1, span_id=127, parent_id=None,
+                          flags=0)
+        span = tracer.start_span('test-%s' % fmt, child_of=ctx)
+        carrier = dict()
+        tracer.inject(span, fmt, carrier)
+        assert len(get_trace_id(carrier)) == 16
 
 
 def test_debug_id():

tests/test_config.py

@@ -156,3 +156,20 @@ def test_initialize_tracer(self):
     def test_default_local_agent_reporting_port(self):
         c = Config({}, service_name='x')
         assert c.local_agent_reporting_port == 6831
+
+    def test_generate_128bit_trace_id(self):
+        c = Config({}, service_name='x')
+        assert c.generate_128bit_trace_id is False
+
+        c = Config({'generate_128bit_trace_id': True}, service_name='x')
+        assert c.generate_128bit_trace_id is True
+
+        os.environ['JAEGER_TRACEID_128BIT'] = 'true'
+        c = Config({'generate_128bit_trace_id': False}, service_name='x')
+        assert c.generate_128bit_trace_id is False
+
+        c = Config({}, service_name='x')
+        assert c.generate_128bit_trace_id is True
+
+        os.environ.pop('JAEGER_TRACEID_128BIT')
+        assert os.getenv('JAEGER_TRACEID_128BIT', None) is None

tests/test_sampler.py

@@ -195,7 +195,7 @@ def test_guaranteed_throughput_probabilistic_sampler():
     sampled, tags = sampler.is_sampled(MAX_INT - 10)
     assert sampled
     assert tags == get_tags('probabilistic', 0.51)
-    sampled, tags = sampler.is_sampled(MAX_INT + (MAX_INT / 4))
+    sampled, tags = sampler.is_sampled(int(MAX_INT + (MAX_INT / 4)))
     assert sampled
     assert tags == get_tags('lowerbound', 0.51)
 
@@ -234,7 +234,7 @@ def test_adaptive_sampler():
         # Move time forward by a second to guarantee the rate limiter has enough credits
         mock_time.side_effect = lambda: ts + 1
 
-        sampled, tags = sampler.is_sampled(MAX_INT + (MAX_INT / 4), 'new_op')
+        sampled, tags = sampler.is_sampled(int(MAX_INT + (MAX_INT / 4)), 'new_op')
         assert sampled
         assert tags == get_tags('lowerbound', 0.51)
 
@@ -243,7 +243,7 @@ def test_adaptive_sampler():
     sampled, tags = sampler.is_sampled(MAX_INT - 10, 'new_op_2')
     assert sampled
     assert tags == get_tags('probabilistic', 0.51)
-    sampled, _ = sampler.is_sampled(MAX_INT + (MAX_INT / 4), 'new_op_2')
+    sampled, _ = sampler.is_sampled(int(MAX_INT + (MAX_INT / 4)), 'new_op_2')
     assert not sampled
     assert '%s' % sampler == 'AdaptiveSampler(0.510000, 3.000000, 2)'
 

tests/test_thrift.py

@@ -60,9 +60,9 @@ def now_reading(self):
 
 def test_large_ids(tracer):
 
-    def serialize(span_id):
+    def serialize(trace_id, span_id):
         """Checks that there are no exceptions during marshalling."""
-        parent_ctx = SpanContext(trace_id=span_id, span_id=span_id,
+        parent_ctx = SpanContext(trace_id=trace_id, span_id=span_id,
                                  parent_id=0, flags=1)
         parent = Span(context=parent_ctx, operation_name='x', tracer=tracer)
         span = tracer.start_span(operation_name='x',
@@ -71,25 +71,44 @@ def serialize(span_id):
         _marshall_span(span)
 
     trace_id = 0x77fd53dc6b437681
-    serialize(trace_id)
+    serialize(trace_id, trace_id)
     assert thrift.id_to_int(trace_id) == 0x77fd53dc6b437681
 
     trace_id = 0x7fffffffffffffff
-    serialize(trace_id)
+    serialize(trace_id, trace_id)
     assert thrift.id_to_int(trace_id) == 0x7fffffffffffffff
 
     trace_id = 0x8000000000000000
-    serialize(trace_id)
+    serialize(trace_id, trace_id)
     assert thrift.id_to_int(trace_id) == -0x8000000000000000
 
     trace_id = 0x97fd53dc6b437681
-    serialize(trace_id)
+    serialize(trace_id, trace_id)
 
     trace_id = (1 << 64) - 1
     assert trace_id == 0xffffffffffffffff
-    serialize(trace_id)
+    serialize(trace_id, trace_id)
     assert thrift.id_to_int(trace_id) == -1
 
+    trace_id = (1 << 128) - 1
+    span_id = 0xffffffffffffffff
+    assert trace_id == 0xffffffffffffffffffffffffffffffff
+    serialize(trace_id, span_id)
+    assert thrift._id_to_low(trace_id) == 0xffffffffffffffff
+    assert thrift._id_to_high(trace_id) == 0xffffffffffffffff
+
+    trace_id = 0xfb34678b8864f051e5c8c603484e57
+    span_id = 0x77fd53dc6b437681
+    serialize(trace_id, span_id)
+    assert thrift._id_to_low(trace_id) == 0x51e5c8c603484e57
+    assert thrift._id_to_high(trace_id) == 0xfb34678b8864f0
+
+
+def test_none_ids():
+    assert thrift.id_to_int(None) is None
+    assert thrift._id_to_low(None) is None
+    assert thrift._id_to_high(None) is None
+
 
 def test_large_tags():
     tag = thrift.make_tag('x', 'y' * 300, max_length=256)