Add support for references (#241)

* Add support for references

- Checks if all the given references are valid
- If the parent is not set then sets the first reference as parent of
  the span
- Allows to pass one Reference object or list of them

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

* Add tests for the references support

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

* Avoid allocation if references are None

Signed-off-by: Yuri Shkuro <[email protected]>

Author: bhavin192

jaeger_client/span.py

@@ -32,12 +32,13 @@ class Span(opentracing.Span):
                  'logs', 'tags', 'update_lock']
 
     def __init__(self, context, tracer, operation_name,
-                 tags=None, start_time=None):
+                 tags=None, start_time=None, references=None):
         super(Span, self).__init__(context=context, tracer=tracer)
         self.operation_name = operation_name
         self.start_time = start_time or time.time()
         self.end_time = None
         self.update_lock = threading.Lock()
+        self.references = references
         # we store tags and logs as Thrift objects to avoid extra allocations
         self.tags = []
         self.logs = []

jaeger_client/thrift.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import six
+from opentracing.tracer import ReferenceType
 
 import jaeger_client.thrift_gen.jaeger.ttypes as ttypes
 import jaeger_client.thrift_gen.sampling.SamplingManager as sampling_manager
@@ -160,6 +161,26 @@ def make_process(service_name, tags, max_length):
     )
 
 
+def make_ref_type(span_ref_type):
+    if span_ref_type == ReferenceType.FOLLOWS_FROM:
+        return ttypes.SpanRefType.FOLLOWS_FROM
+    return ttypes.SpanRefType.CHILD_OF
+
+
+def make_references(references):
+    if not references:
+        return None
+    list_of_span_refs = list()
+    for span_ref in references:
+        list_of_span_refs.append(ttypes.SpanRef(
+            refType=make_ref_type(span_ref.type),
+            traceIdLow=id_to_int(_id_to_low(span_ref.referenced_context.trace_id)),
+            traceIdHigh=id_to_int(_id_to_high(span_ref.referenced_context.trace_id)),
+            spanId=id_to_int(span_ref.referenced_context.span_id),
+        ))
+    return list_of_span_refs
+
+
 def make_jaeger_batch(spans, process):
     batch = ttypes.Batch(
         spans=[],
@@ -173,7 +194,7 @@ def make_jaeger_batch(spans, process):
                 spanId=id_to_int(span.span_id),
                 parentSpanId=id_to_int(span.parent_id) or 0,
                 operationName=span.operation_name,
-                # references = references, # TODO
+                references=make_references(span.references),
                 flags=span.context.flags,
                 startTime=timestamp_micros(span.start_time),
                 duration=timestamp_micros(span.end_time - span.start_time),

jaeger_client/tracer.py

@@ -123,7 +123,7 @@ def start_span(self,
         :param child_of: shortcut for 'child_of' reference
         :param references: (optional) either a single Reference object or a
             list of Reference objects that identify one or more parent
-            SpanContexts. (See the Reference documentation for detail)
+            SpanContexts. (See the opentracing.Reference documentation for detail)
         :param tags: optional dictionary of Span Tags. The caller gives up
             ownership of that dictionary, because the Tracer may use it as-is
             to avoid extra data copying.
@@ -133,11 +133,6 @@ def start_span(self,
         :return: Returns an already-started Span instance.
         """
         parent = child_of
-        if references:
-            if isinstance(references, list):
-                # TODO only the first reference is currently used
-                references = references[0]
-            parent = references.referenced_context
 
         # allow Span to be passed as reference, not just SpanContext
         if isinstance(parent, Span):
@@ -146,6 +141,19 @@ def start_span(self,
         rpc_server = tags and \
             tags.get(ext_tags.SPAN_KIND) == ext_tags.SPAN_KIND_RPC_SERVER
 
+        valid_references = None
+        if references:
+            valid_references = list()
+            if not isinstance(references, list):
+                references = [references]
+            for reference in references:
+                if reference.referenced_context is not None:
+                    valid_references.append(reference)
+
+        # setting first reference as parent
+        if valid_references and (parent is None or not parent.has_trace):
+            parent = valid_references[0].referenced_context
+
         if parent is None or not parent.has_trace:
             trace_id = self._random_id(self.max_trace_id_bits)
             span_id = self._random_id(constants._max_id_bits)
@@ -183,7 +191,7 @@ def start_span(self,
                                baggage=baggage)
         span = Span(context=span_ctx, tracer=self,
                     operation_name=operation_name,
-                    tags=tags, start_time=start_time)
+                    tags=tags, start_time=start_time, references=valid_references)
 
         self._emit_span_metrics(span=span, join=rpc_server)
 

tests/test_thrift.py

@@ -16,7 +16,7 @@
 
 import jaeger_client.thrift_gen.jaeger.ttypes as ttypes
 import jaeger_client.thrift_gen.sampling.SamplingManager as sampling_manager
-from opentracing import child_of
+from opentracing import child_of, follows_from
 from jaeger_client import ProbabilisticSampler, RateLimitingSampler
 from jaeger_client import thrift, Span, SpanContext
 from jaeger_client.thrift_gen.agent import Agent as Agent
@@ -175,3 +175,14 @@ def test_parse_sampling_strategy():
     resp.strategyType = 'x'
     s, e = thrift.parse_sampling_strategy(response=resp)
     assert s is None and e is not None
+
+
+def test_parse_span_references(tracer):
+    span = tracer.start_span('test')
+    span2 = tracer.start_span('test2')
+    follow_span = tracer.start_span('test-follow', references=[follows_from(span.context),
+                                                               child_of(span2.context)])
+    span.finish()
+    span2.finish()
+    follow_span.finish()
+    _marshall_span(follow_span)

tests/test_tracer.py

@@ -20,7 +20,7 @@
 import pytest
 import tornado.httputil
 
-from opentracing import Format, child_of
+from opentracing import Format, child_of, follows_from
 from opentracing.ext import tags as ext_tags
 from jaeger_client import ConstSampler, SpanContext, Tracer
 from jaeger_client import constants as c
@@ -129,6 +129,41 @@ def test_child_span(tracer):
     tracer.close()
 
 
+def test_follows_from(tracer):
+    span = tracer.start_span('test')
+    span1 = tracer.start_span('test2')
+    follow_span = tracer.start_span('follow-span', references=[follows_from(span.context),
+                                                               follows_from(span1.context),
+                                                               follows_from(None)])
+    span.finish()
+    span1.finish()
+    follow_span.finish()
+    tracer.reporter.report_span.assert_called_once()
+    assert len(follow_span.references) == 2
+    assert follow_span.context.parent_id == span.context.span_id
+    for reference in follow_span.references:
+        assert reference.referenced_context is not None
+
+    span = tracer.start_span('test')
+    follow_span = tracer.start_span(references=follows_from(span.context))
+    span.finish()
+    follow_span.finish()
+    tracer.reporter.report_span.assert_called_once()
+    assert isinstance(follow_span.references, list)
+
+    span = tracer.start_span('test')
+    parent_span = tracer.start_span('test-parent')
+    child_span = tracer.start_span('test-child', child_of=parent_span,
+                                   references=follows_from(span.context))
+    span.finish()
+    parent_span.finish()
+    child_span.finish()
+    tracer.reporter.report_span.assert_called_once()
+    assert child_span.context.parent_id == parent_span.context.span_id
+    assert len(child_span.references) == 1
+    tracer.close()
+
+
 def test_sampler_effects(tracer):
     tracer.sampler = ConstSampler(True)
     span = tracer.start_span('test')