Merge pull request #141 from swiftstack/adding-annotations

Add a zipkin_span.add_annotation() method

py_zipkin/logging_helper.py

@@ -37,6 +37,7 @@ def __init__(
         max_span_batch_size=None,
         firehose_handler=None,
         encoding=None,
+        annotations=None,
     ):
         self.zipkin_attrs = zipkin_attrs
         self.endpoint = endpoint
@@ -51,6 +52,7 @@ def __init__(
         self.client_context = client_context
         self.max_span_batch_size = max_span_batch_size
         self.firehose_handler = firehose_handler
+        self.annotations = annotations or {}
 
         self.remote_endpoint = None
         self.encoder = get_encoder(encoding)
@@ -108,10 +110,8 @@ def _emit_spans_with_span_sender(self, span_sender):
 
                 span_sender.add_span(span)
 
-            annotations = {}
-
             if self.add_logging_annotation:
-                annotations[LOGGING_END_KEY] = time.time()
+                self.annotations[LOGGING_END_KEY] = time.time()
 
             span_sender.add_span(Span(
                 trace_id=self.zipkin_attrs.trace_id,
@@ -124,7 +124,7 @@ def _emit_spans_with_span_sender(self, span_sender):
                 local_endpoint=self.endpoint,
                 remote_endpoint=self.remote_endpoint,
                 shared=not self.report_root_timestamp,
-                annotations=annotations,
+                annotations=self.annotations,
                 tags=self.tags,
             ))
 

py_zipkin/zipkin.py

@@ -457,6 +457,7 @@ def start(self):
                 max_span_batch_size=self.max_span_batch_size,
                 firehose_handler=self.firehose_handler,
                 encoding=self.encoding,
+                annotations=self.annotations,
             )
             self.logging_context.start()
             self.get_tracer().set_transport_configured(configured=True)
@@ -542,6 +543,26 @@ def update_binary_annotations(self, extra_annotations):
             # the binary annotations for the logging context directly.
             self.logging_context.tags.update(extra_annotations)
 
+    def add_annotation(self, value, timestamp=None):
+        """Add an annotation for the current span
+
+        The timestamp defaults to "now", but may be specified.
+
+        :param value: The annotation string
+        :type value: str
+        :param timestamp: Timestamp for the annotation
+        :type timestamp: float
+        """
+        timestamp = timestamp or time.time()
+        if not self.logging_context:
+            # This is not the root span, so annotations will be added
+            # to the log handler when this span context exits.
+            self.annotations[value] = timestamp
+        else:
+            # Otherwise, we're in the context of the root span, so just update
+            # the annotations for the logging context directly.
+            self.logging_context.annotations[value] = timestamp
+
     def add_sa_binary_annotation(
         self,
         port=0,

tests/zipkin_test.py

@@ -470,6 +470,7 @@ def test_start_root_span(self, mock_time, mock_log_ctx):
             max_span_batch_size=50,
             firehose_handler=firehose,
             encoding=Encoding.V2_JSON,
+            annotations={},
         )
         assert mock_log_ctx.return_value.start.call_count == 1
         assert tracer.is_transport_configured() is True
@@ -729,6 +730,69 @@ def test_update_binary_annotations_non_root_not_traced(self):
                 'status': '200',
             }
 
+    def test_add_annotation_root(self):
+        with zipkin.zipkin_span(
+            service_name='test_service',
+            span_name='test_span',
+            transport_handler=MockTransportHandler(),
+            sample_rate=100.0,
+            annotations={'abc': 123},
+            add_logging_annotation=True,
+        ) as span:
+            span.add_annotation('def', 345)
+            span.add_annotation('ghi', timestamp=678)
+            with mock.patch('py_zipkin.zipkin.time.time') as mock_time:
+                mock_time.return_value = 91011
+                span.add_annotation('jkl')
+
+            assert span.logging_context.annotations == {
+                'abc': 123,
+                'def': 345,
+                'ghi': 678,
+                'jkl': 91011,
+            }
+
+    def test_add_annotation_non_root(self):
+        context = zipkin.zipkin_span(
+            service_name='test_service',
+            span_name='test_span',
+            annotations={'abc': 123},
+        )
+        context.get_tracer()._context_stack.push(zipkin.create_attrs_for_span())
+        with context as span:
+            span.add_annotation('def', 345)
+            span.add_annotation('ghi', timestamp=678)
+            with mock.patch('py_zipkin.zipkin.time.time') as mock_time:
+                mock_time.return_value = 91011
+                span.add_annotation('jkl')
+
+            assert span.annotations == {
+                'abc': 123,
+                'def': 345,
+                'ghi': 678,
+                'jkl': 91011,
+            }
+
+    def test_add_annotation_non_root_not_traced(self):
+        # nothing happens if the request is not traced
+        with zipkin.zipkin_span(
+            service_name='test_service',
+            span_name='test_span',
+            annotations={'abc': 123},
+        ) as span:
+            span.add_annotation('def', 345)
+            span.add_annotation('ghi', timestamp=678)
+            with mock.patch('py_zipkin.zipkin.time.time') as mock_time:
+                mock_time.return_value = 91011
+                span.add_annotation('jkl')
+
+            assert span.annotations == {
+                'abc': 123,
+                'def': 345,
+                'ghi': 678,
+                'jkl': 91011,
+            }
+
     def test_add_sa_binary_annotation_non_client(self):
         # Nothing happens if this is not a client span
         context = zipkin.zipkin_span('test_service', 'test_span')