Added processing for client root spans

py_zipkin/logging_helper.py

@@ -30,9 +30,8 @@ def emit(self, record):
 class ZipkinLoggingContext(object):
     """A logging context specific to a Zipkin trace. If the trace is sampled,
     the logging context sends serialized Zipkin spans to a transport_handler.
-    This is meant to only be used server-side, so the logging context sends
-    a single "server" span, as well as all local child spans collected within
-    the server context.
+    The logging context sends root "server" or "client" span, as well as all
+    local child spans collected within this context.
 
     This class should only be used by the main `zipkin_span` entrypoint.
     """
@@ -47,6 +46,7 @@ def __init__(
         report_root_timestamp,
         binary_annotations=None,
         add_logging_annotation=False,
+        client_context=False
     ):
         self.zipkin_attrs = zipkin_attrs
         self.thrift_endpoint = thrift_endpoint
@@ -58,6 +58,7 @@ def __init__(
         self.binary_annotations_dict = binary_annotations or {}
         self.sa_binary_annotations = []
         self.add_logging_annotation = add_logging_annotation
+        self.client_context = client_context
 
     def start(self):
         """Actions to be taken before request is handled.
@@ -81,10 +82,11 @@ def stop(self):
     def log_spans(self):
         """Main function to log all the annotations stored during the entire
         request. This is done if the request is sampled and the response was
-        a success. It also logs the service `ss` and `sr` annotations.
+        a success. It also logs the service (`ss` and `sr`) or the client
+        ('cs' and 'cr') annotations.
         """
         if self.zipkin_attrs.is_sampled:
-            server_end_time = time.time()
+            end_timestamp = time.time()
             # Collect additional annotations from the logging handler
             annotations_by_span_id = defaultdict(dict)
             binary_annotations_by_span_id = defaultdict(dict)
@@ -142,20 +144,20 @@ def log_spans(self):
                     duration_s=duration,
                 )
 
-            # Collect extra annotations for server span
             extra_annotations = annotations_by_span_id[
                 self.zipkin_attrs.span_id]
             extra_binary_annotations = binary_annotations_by_span_id[
                 self.zipkin_attrs.span_id
             ]
-            logging_end = time.time()
-            annotations = dict(
-                sr=self.start_timestamp,
-                ss=server_end_time,
-                **extra_annotations
-            )
+
+            k1, k2 = ('sr', 'ss')
+            if self.client_context:
+                k1, k2 = ('cs', 'cr')
+            annotations = {k1: self.start_timestamp, k2: end_timestamp}
+            annotations.update(extra_annotations)
+
             if self.add_logging_annotation:
-                annotations[LOGGING_END_KEY] = logging_end
+                annotations[LOGGING_END_KEY] = time.time()
 
             thrift_annotations = annotation_list_builder(
                 annotations,
@@ -174,7 +176,7 @@ def log_spans(self):
 
             if self.report_root_timestamp:
                 timestamp = self.start_timestamp
-                duration = server_end_time - self.start_timestamp
+                duration = end_timestamp - self.start_timestamp
             else:
                 timestamp = duration = None
 

py_zipkin/zipkin.py

@@ -281,6 +281,7 @@ def start(self):
             if not self.zipkin_attrs.is_sampled:
                 return self
             endpoint = create_endpoint(self.port, self.service_name, self.host)
+            client_context = set(self.include) == {'client'}
             self.log_handler = ZipkinLoggerHandler(self.zipkin_attrs)
             self.logging_context = ZipkinLoggingContext(
                 self.zipkin_attrs,
@@ -291,6 +292,7 @@ def start(self):
                 report_root_timestamp or self.report_root_timestamp_override,
                 binary_annotations=self.binary_annotations,
                 add_logging_annotation=self.add_logging_annotation,
+                client_context=client_context
             )
             self.logging_context.start()
             return self

setup.py

@@ -3,7 +3,7 @@
 from setuptools import find_packages
 from setuptools import setup
 
-__version__ = '0.8.1'
+__version__ = '0.8.2'
 
 setup(
     name='py_zipkin',

tests/logging_helper_test.py

@@ -64,7 +64,7 @@ def test_zipkin_logging_context(time_mock, mock_logger, context):
             autospec=True)
 @mock.patch('py_zipkin.logging_helper.copy_endpoint_with_new_service_name',
             autospec=True)
-def test_zipkin_logging_context_log_spans(
+def test_zipkin_logging_server_context_log_spans(
     copy_endpoint_mock, bin_ann_list_builder, ann_list_builder,
     log_span_mock, time_mock
 ):
@@ -162,6 +162,73 @@ def test_zipkin_logging_context_log_spans(
     }
 
 
+@mock.patch('py_zipkin.logging_helper.time.time', autospec=True)
+@mock.patch('py_zipkin.logging_helper.log_span', autospec=True)
+@mock.patch('py_zipkin.logging_helper.annotation_list_builder',
+            autospec=True)
+@mock.patch('py_zipkin.logging_helper.binary_annotation_list_builder',
+            autospec=True)
+@mock.patch('py_zipkin.logging_helper.copy_endpoint_with_new_service_name',
+            autospec=True)
+def test_zipkin_logging_client_context_log_spans(
+    copy_endpoint_mock, bin_ann_list_builder, ann_list_builder,
+    log_span_mock, time_mock
+):
+    # This lengthy function tests that the logging context properly
+    # logs root client span
+    trace_id = '000000000000000f'
+    client_span_id = '0000000000000003'
+    attr = ZipkinAttrs(
+        trace_id=trace_id,
+        span_id=client_span_id,
+        parent_span_id=None,
+        flags=None,
+        is_sampled=True,
+    )
+    handler = logging_helper.ZipkinLoggerHandler(attr)
+    handler.client_spans = []
+
+    # Each of the thrift annotation helpers just reflects its first arg
+    # so the annotation dicts can be checked.
+    ann_list_builder.side_effect = lambda x, y: x
+    bin_ann_list_builder.side_effect = lambda x, y: x
+
+    transport_handler = mock.Mock()
+
+    context = logging_helper.ZipkinLoggingContext(
+        zipkin_attrs=attr,
+        thrift_endpoint='thrift_endpoint',
+        log_handler=handler,
+        span_name='GET /foo',
+        transport_handler=transport_handler,
+        report_root_timestamp=True,
+        client_context=True
+    )
+
+    context.start_timestamp = 24
+    context.response_status_code = 200
+
+    context.binary_annotations_dict = {'k': 'v'}
+    time_mock.return_value = 42
+
+    expected_server_annotations = {'cs': 24, 'cr': 42}
+    expected_server_bin_annotations = {'k': 'v'}
+
+    context.log_spans()
+    log_call = log_span_mock.call_args_list[0]
+    assert log_call[1] == {
+        'span_id': client_span_id,
+        'parent_span_id': None,
+        'trace_id': trace_id,
+        'span_name': 'GET /foo',
+        'annotations': expected_server_annotations,
+        'binary_annotations': expected_server_bin_annotations,
+        'transport_handler': transport_handler,
+        'duration_s': 18,
+        'timestamp_s': 24,
+    }
+
+
 @mock.patch('py_zipkin.logging_helper.log_span', autospec=True)
 def test_log_span_not_called_if_not_sampled(log_span_mock):
     attr = ZipkinAttrs(

tests/zipkin_test.py

@@ -55,6 +55,7 @@ def test_zipkin_span_for_new_trace(
         report_root_timestamp=True,
         binary_annotations={},
         add_logging_annotation=False,
+        client_context=False,
     )
     pop_zipkin_attrs_mock.assert_called_once_with()
 
@@ -107,6 +108,7 @@ def test_zipkin_span_passed_sampled_attrs(
         report_root_timestamp=False,
         binary_annotations={},
         add_logging_annotation=False,
+        client_context=False,
     )
     pop_zipkin_attrs_mock.assert_called_once_with()
 
@@ -440,7 +442,7 @@ def test_span_context(
 @mock.patch('py_zipkin.zipkin.create_endpoint')
 @mock.patch('py_zipkin.zipkin.ZipkinLoggerHandler', autospec=True)
 @mock.patch('py_zipkin.zipkin.ZipkinLoggingContext', autospec=True)
-def test_zipkin_span_decorator(
+def test_zipkin_server_span_decorator(
     logging_context_cls_mock,
     logger_handler_cls_mock,
     create_endpoint_mock,
@@ -483,6 +485,62 @@ def test_func(a, b):
         report_root_timestamp=True,
         binary_annotations={},
         add_logging_annotation=False,
+        client_context=False,
+    )
+    pop_zipkin_attrs_mock.assert_called_once_with()
+
+
+@mock.patch('py_zipkin.zipkin.pop_zipkin_attrs', autospec=True)
+@mock.patch('py_zipkin.zipkin.push_zipkin_attrs', autospec=True)
+@mock.patch('py_zipkin.zipkin.create_attrs_for_span', autospec=True)
+@mock.patch('py_zipkin.zipkin.create_endpoint')
+@mock.patch('py_zipkin.zipkin.ZipkinLoggerHandler', autospec=True)
+@mock.patch('py_zipkin.zipkin.ZipkinLoggingContext', autospec=True)
+def test_zipkin_client_span_decorator(
+    logging_context_cls_mock,
+    logger_handler_cls_mock,
+    create_endpoint_mock,
+    create_attrs_for_span_mock,
+    push_zipkin_attrs_mock,
+    pop_zipkin_attrs_mock,
+):
+    transport_handler = mock.Mock()
+
+    @zipkin.zipkin_span(
+        service_name='some_service_name',
+        span_name='span_name',
+        transport_handler=transport_handler,
+        port=5,
+        sample_rate=100.0,
+        include=('client',),
+        host='1.5.1.2',
+    )
+    def test_func(a, b):
+        return a + b
+
+    assert test_func(1, 2) == 3
+
+    create_attrs_for_span_mock.assert_called_once_with(
+        sample_rate=100.0,
+        use_128bit_trace_id=False,
+    )
+    push_zipkin_attrs_mock.assert_called_once_with(
+        create_attrs_for_span_mock.return_value)
+    create_endpoint_mock.assert_called_once_with(5, 'some_service_name', '1.5.1.2')
+    logger_handler_cls_mock.assert_called_once_with(
+        create_attrs_for_span_mock.return_value)
+    # The decorator was passed a sample rate and no Zipkin attrs, so it's
+    # assumed to be the root of a trace and it should report timestamp/duration
+    logging_context_cls_mock.assert_called_once_with(
+        create_attrs_for_span_mock.return_value,
+        create_endpoint_mock.return_value,
+        logger_handler_cls_mock.return_value,
+        'span_name',
+        transport_handler,
+        report_root_timestamp=True,
+        binary_annotations={},
+        add_logging_annotation=False,
+        client_context=True,
     )
     pop_zipkin_attrs_mock.assert_called_once_with()