Export span once it's finished (#68)

Author: liyanhui1228

trace/opencensus/trace/ext/django/middleware.py

@@ -65,10 +65,10 @@ def _set_django_labels(tracer, request):
 
     # User id is the django autofield for User model as the primary key
     if user_id is not None:
-        tracer.add_label_to_spans('/django/user/id', str(user_id))
+        tracer.add_label_to_current_span('/django/user/id', str(user_id))
 
     if user_name is not None:
-        tracer.add_label_to_spans('/django/user/name', user_name)
+        tracer.add_label_to_current_span('/django/user/name', user_name)
 
 
 def get_django_header():
@@ -148,10 +148,10 @@ def process_request(self, request):
 
             # Span name is being set at process_view
             tracer.start_span()
-            tracer.add_label_to_spans(
+            tracer.add_label_to_current_span(
                 label_key=HTTP_METHOD,
                 label_value=request.method)
-            tracer.add_label_to_spans(
+            tracer.add_label_to_current_span(
                 label_key=HTTP_URL,
                 label_value=request.path)
         except Exception:  # pragma: NO COVER
@@ -173,7 +173,7 @@ def process_view(self, request, view_func, *args, **kwargs):
     def process_response(self, request, response):
         try:
             tracer = _get_current_request_tracer()
-            tracer.add_label_to_spans(
+            tracer.add_label_to_current_span(
                 label_key=HTTP_STATUS_CODE,
                 label_value=str(response.status_code))
 

trace/opencensus/trace/ext/flask/flask_middleware.py

@@ -90,8 +90,8 @@ def _before_request(self):
             span.name = '[{}]{}'.format(
                 flask.request.method,
                 flask.request.url)
-            tracer.add_label_to_spans(HTTP_METHOD, flask.request.method)
-            tracer.add_label_to_spans(HTTP_URL, flask.request.url)
+            tracer.add_label_to_current_span(HTTP_METHOD, flask.request.method)
+            tracer.add_label_to_current_span(HTTP_URL, flask.request.url)
         except Exception:  # pragma: NO COVER
             log.error('Failed to trace request', exc_info=True)
 
@@ -102,7 +102,7 @@ def _after_request(self, response):
         """
         try:
             tracer = execution_context.get_opencensus_tracer()
-            tracer.add_label_to_spans(
+            tracer.add_label_to_current_span(
                 HTTP_STATUS_CODE,
                 str(response.status_code))
 

trace/opencensus/trace/request_tracer.py

@@ -81,7 +81,9 @@ def get_tracer(self):
         sampled = self.should_sample()
 
         if sampled:
-            return context_tracer.ContextTracer(self.span_context)
+            return context_tracer.ContextTracer(
+                exporter=self.exporter,
+                span_context=self.span_context)
         else:
             return noop_tracer.NoopTracer()
 
@@ -90,11 +92,8 @@ def store_tracer(self):
         execution_context.set_opencensus_tracer(self)
 
     def finish(self):
-        """End all spans and send spans using reporter."""
-        trace = self.tracer.finish()
-
-        if trace is not None:
-            self.exporter.export(trace)
+        """End all spans."""
+        self.tracer.finish()
 
     def span(self, name=None):
         """Create a new span with the trace using the context information.
@@ -111,9 +110,8 @@ def start_span(self, name=None):
         return self.tracer.start_span(name)
 
     def end_span(self):
-        """End a span. Remove the span from the span stack, and update the
-        span_id in TraceContext as the current span_id which is the peek
-        element in the span stack.
+        """End a span. Update the span_id in SpanContext to the current span's
+        parent span id; Update the current span; Send the span to exporter.
         """
         self.tracer.end_span()
 
@@ -132,9 +130,6 @@ def add_label_to_current_span(self, label_key, label_value):
         """
         self.tracer.add_label_to_current_span(label_key, label_value)
 
-    def add_label_to_spans(self, label_key, label_value):
-        self.tracer.add_label_to_spans(label_key, label_value)
-
     def trace_decorator(self):
         """Decorator to trace a function."""
 

trace/opencensus/trace/tracer/base.py

@@ -58,17 +58,6 @@ def current_span(self):
     def add_label_to_current_span(self, label_key, label_value):
         raise NotImplementedError
 
-    def add_label_to_spans(self, label_key, label_value):
-        """Add label to the spans in current trace.
-
-        :type label_key: str
-        :param label_key: Label key.
-
-        :type label_value:str
-        :param label_value: Label value.
-        """
-        raise NotImplementedError
-
     def list_collected_spans(self):
         """List collected spans."""
         raise NotImplementedError

trace/opencensus/trace/tracer/context_tracer.py

@@ -17,6 +17,7 @@
 from opencensus.trace import execution_context
 from opencensus.trace.span_context import SpanContext
 from opencensus.trace import span as trace_span
+from opencensus.trace.exporters import print_exporter
 from opencensus.trace.tracer import base
 
 
@@ -27,10 +28,14 @@ class ContextTracer(base.Tracer):
     :param span_context: SpanContext encapsulates the current context within
                          the request's trace.
     """
-    def __init__(self, span_context=None):
+    def __init__(self, exporter=None, span_context=None):
+        if exporter is None:
+            exporter = print_exporter.PrintExporter()
+
         if span_context is None:
             span_context = SpanContext()
 
+        self.exporter = exporter
         self.span_context = span_context
         self.trace_id = span_context.trace_id
         self.root_span_id = span_context.span_id
@@ -44,11 +49,8 @@ def finish(self):
         :rtype: dict
         :returns: JSON format trace.
         """
-        trace = self._get_trace_json()
         self._spans_list = []
 
-        return trace
-
     def span(self, name='span'):
         """Create a new span with the trace using the context information.
 
@@ -88,10 +90,9 @@ def start_span(self, name='span'):
         span.start()
         return span
 
-    def end_span(self):
-        """End a span. Remove the span from the span stack, and update the
-        span_id in TraceContext as the current span_id which is the peek
-        element in the span stack.
+    def end_span(self, *args, **kwargs):
+        """End a span. Update the span_id in SpanContext to the current span's
+        parent span id; Update the current span.
         """
         cur_span = self.current_span()
 
@@ -107,6 +108,11 @@ def end_span(self):
         else:
             execution_context.set_current_span(None)
 
+        span_json = self.get_trace_json(cur_span)
+        self.exporter.export(span_json)
+
+        return cur_span
+
     def current_span(self):
         """Return the current span."""
         current_span = execution_context.get_current_span()
@@ -128,33 +134,13 @@ def add_label_to_current_span(self, label_key, label_value):
         current_span = self.current_span()
         current_span.add_label(label_key, label_value)
 
-    def add_label_to_spans(self, label_key, label_value):
-        """Add label to the spans in current trace.
-
-        :type label_key: str
-        :param label_key: Label key.
-
-        :type label_value:str
-        :param label_value: Label value.
-        """
-        for span in self._spans_list:
-            span.add_label(label_key, label_value)
-
-    def _get_trace_json(self):
+    def get_trace_json(self, span):
         """Get the JSON format trace."""
-        spans_list = []
-        for root_span in self._spans_list:
-            span_tree = list(iter(root_span))
-            span_tree_json = [trace_span.format_span_json(span)
-                              for span in span_tree]
-            spans_list.extend(span_tree_json)
-
-        if len(spans_list) == 0:
-            return
+        span_json = trace_span.format_span_json(span)
 
         trace = {
             'traceId': self.trace_id,
-            'spans': spans_list,
+            'spans': [span_json],
         }
 
         return trace

trace/opencensus/trace/tracer/noop_tracer.py

@@ -50,7 +50,7 @@ def end_span(self):
         span_id in TraceContext as the current span_id which is the peek
         element in the span stack.
         """
-        return
+        return base.NullContextManager()
 
     def current_span(self):
         """Return the current span."""
@@ -67,17 +67,6 @@ def add_label_to_current_span(self, label_key, label_value):
         """
         return
 
-    def add_label_to_spans(self, label_key, label_value):
-        """Add label to the spans in current trace.
-
-        :type label_key: str
-        :param label_key: Label key.
-
-        :type label_value:str
-        :param label_value: Label value.
-        """
-        return
-
     def list_collected_spans(self):
         """List collected spans."""
         return None

trace/tests/system/basic_trace/basic_trace_system_test.py

@@ -17,7 +17,6 @@
 
 def func_to_trace():
     import time
-
     print('Test simple tracing...')
     time.sleep(2)
 
@@ -64,7 +63,7 @@ def test_request_tracer(self):
         spans = trace_json.get('spans')
 
         self.assertEqual(trace_json.get('traceId'), trace_id)
-        self.assertEqual(len(spans), 2)
+        self.assertEqual(len(spans), 1)
 
         for span in spans:
             if span.get('name') == 'root_span':

trace/tests/unit/ext/django/test_middleware.py

@@ -238,7 +238,6 @@ def test_process_response(self):
         middleware_obj.process_response(django_request, django_response)
 
         self.assertEqual(span.labels, expected_labels)
-        self.assertTrue(exporter_mock.export.called)
 
 
 class Test__set_django_labels(unittest.TestCase):
@@ -247,7 +246,7 @@ class Tracer(object):
         def __init__(self):
             self.labels = {}
 
-        def add_label_to_spans(self, key, value):
+        def add_label_to_current_span(self, key, value):
             self.labels[key] = value
 
     def test__set_django_labels_no_user(self):

trace/tests/unit/test_span.py

@@ -156,6 +156,14 @@ def test_finish_with_context_tracer(self):
 
         self.assertTrue(context_tracer.end_span.called)
 
+    def test_finish(self):
+        span_name = 'root_span'
+        span = self._make_one(span_name)
+        self.assertIsNone(span.end_time)
+
+        span.finish()
+        self.assertIsNotNone(span.end_time)
+
     def test___iter__(self):
         root_span_name = 'root_span_name'
         child1_span_name = 'child1_span_name'
@@ -175,3 +183,75 @@ def test___iter__(self):
         self.assertEqual(
             span_iter_list,
             [child1_child1_span, child1_span, child2_span, root_span])
+
+
+class Test_format_span_json(unittest.TestCase):
+
+    def test_format_span_json_no_parent_span(self):
+        from opencensus.trace.span import format_span_json
+        from opencensus.trace.enums import Enum
+
+        name = 'test span'
+        kind = Enum.SpanKind.SPAN_KIND_UNSPECIFIED
+        span_id = 1234
+        start_time = '2017-06-25'
+        end_time = '2017-06-26'
+
+        span = mock.Mock()
+        span.name = name
+        span.kind = kind
+        span.span_id = span_id
+        span.start_time = start_time
+        span.end_time = end_time
+        span.parent_span = None
+        span.labels = None
+
+        expected_span_json = {
+            'name': name,
+            'kind': kind,
+            'spanId': span_id,
+            'startTime': start_time,
+            'endTime': end_time,
+        }
+
+        span_json = format_span_json(span)
+        self.assertEqual(span_json, expected_span_json)
+
+    def test_format_span_json_with_parent_span(self):
+        from opencensus.trace.span import format_span_json
+        from opencensus.trace.enums import Enum
+
+        name = 'test span'
+        kind = Enum.SpanKind.SPAN_KIND_UNSPECIFIED
+        span_id = 1234
+        labels = {
+            '/http/status_code': '200',
+            '/component': 'HTTP load balancer',
+        }
+        start_time = '2017-06-25'
+        end_time = '2017-06-26'
+        parent_span = mock.Mock()
+        parent_span_id = 5678
+        parent_span.span_id = parent_span_id
+
+        span = mock.Mock()
+        span.parent_span = parent_span
+        span.name = name
+        span.kind = kind
+        span.labels = labels
+        span.span_id = span_id
+        span.start_time = start_time
+        span.end_time = end_time
+
+        expected_span_json = {
+            'name': name,
+            'kind': kind,
+            'spanId': span_id,
+            'parentSpanId': parent_span_id,
+            'startTime': start_time,
+            'endTime': end_time,
+            'labels': labels,
+        }
+
+        span_json = format_span_json(span)
+        self.assertEqual(span_json, expected_span_json)

trace/tests/unit/tracer/test_base_tracer.py

@@ -57,14 +57,6 @@ def test_add_label_to_current_span(self):
         with self.assertRaises(NotImplementedError):
             tracer.add_label_to_current_span(label_key, label_value)
 
-    def test_add_label_to_spans_abstract(self):
-        tracer = base.Tracer()
-        label_key = 'key'
-        label_value = 'value'
-
-        with self.assertRaises(NotImplementedError):
-            tracer.add_label_to_spans(label_key, label_value)
-
     def test_list_collected_spans_abstract(self):
         tracer = base.Tracer()