Update ContextTracer to flush leftover spans on finish (#150)

Author: wkiser

opencensus/trace/tracers/context_tracer.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import logging
+import threading
 
 from opencensus.trace import execution_context
 from opencensus.trace.span_context import SpanContext
@@ -42,6 +43,7 @@ def __init__(self, exporter=None, span_context=None):
         self.trace_id = span_context.trace_id
         self.root_span_id = span_context.span_id
 
+        self._spans_list_condition = threading.Condition()
         # List of spans to report
         self._spans_list = []
 
@@ -51,7 +53,8 @@ def finish(self):
         :rtype: dict
         :returns: JSON format trace.
         """
-        self._spans_list = []
+        while self._spans_list:
+            self.end_span()
 
     def span(self, name='span'):
         """Create a new span with the trace using the context information.
@@ -86,7 +89,8 @@ def start_span(self, name='span'):
             name,
             parent_span=parent_span,
             context_tracer=self)
-        self._spans_list.append(span)
+        with self._spans_list_condition:
+            self._spans_list.append(span)
         self.span_context.span_id = span.span_id
         execution_context.set_current_span(span)
         span.start()
@@ -97,21 +101,27 @@ def end_span(self, *args, **kwargs):
         parent span id; Update the current span.
         """
         cur_span = self.current_span()
+        if cur_span is None and self._spans_list:
+            cur_span = self._spans_list[-1]
 
         if cur_span is None:
             logging.warning('No active span, cannot do end_span.')
             return
 
         cur_span.finish()
-        self.span_context.span_id = cur_span.parent_span.span_id
+        self.span_context.span_id = cur_span.parent_span.span_id if \
+            cur_span.parent_span else None
 
         if isinstance(cur_span.parent_span, trace_span.Span):
             execution_context.set_current_span(cur_span.parent_span)
         else:
             execution_context.set_current_span(None)
 
-        span_datas = self.get_span_datas(cur_span)
-        self.exporter.export(span_datas)
+        with self._spans_list_condition:
+            if cur_span in self._spans_list:
+                span_datas = self.get_span_datas(cur_span)
+                self.exporter.export(span_datas)
+                self._spans_list.remove(cur_span)
 
         return cur_span
 
@@ -148,7 +158,8 @@ def get_span_datas(self, span):
                 name=span.name,
                 context=self.span_context,
                 span_id=span.span_id,
-                parent_span_id=span.parent_span.span_id,
+                parent_span_id=span.parent_span.span_id if
+                span.parent_span else None,
                 attributes=span.attributes,
                 start_time=span.start_time,
                 end_time=span.end_time,

tests/unit/trace/test_tracer.py

@@ -17,6 +17,7 @@
 import mock
 
 from opencensus.trace import tracer as tracer_module
+from opencensus.trace import span_data
 
 
 class TestTracer(unittest.TestCase):
@@ -246,7 +247,8 @@ def test_add_attribute_to_current_span_not_sampled(self):
         assert isinstance(span, base.NullContextManager)
 
     def test_trace_decorator(self):
-        tracer = tracer_module.Tracer()
+        mock_exporter = mock.MagicMock()
+        tracer = tracer_module.Tracer(exporter=mock_exporter)
 
         return_value = "test"
 
@@ -256,6 +258,8 @@ def test_decorator():
 
         returned = test_decorator()
 
-        self.assertEqual(len(tracer.tracer._spans_list), 1)
-        self.assertEqual(tracer.tracer._spans_list[0].name, 'test_decorator')
         self.assertEqual(returned, return_value)
+        self.assertEqual(mock_exporter.export.call_count, 1)
+        exported_spandata = mock_exporter.export.call_args[0][0][0]
+        self.assertIsInstance(exported_spandata, span_data.SpanData)
+        self.assertEqual(exported_spandata.name, 'test_decorator')

tests/unit/trace/tracers/test_context_tracer.py

@@ -17,9 +17,15 @@
 import mock
 
 from opencensus.trace.tracers import context_tracer
+from opencensus.trace import span
+from opencensus.trace import execution_context
 
 
 class TestContextTracer(unittest.TestCase):
+
+    def tearDown(self):
+        execution_context.clear()
+
     def test_constructor_defaults(self):
         from opencensus.trace import span_context
         from opencensus.trace.exporters import print_exporter
@@ -58,7 +64,14 @@ def test_finish_without_spans(self):
 
     def test_finish_with_spans(self):
         tracer = context_tracer.ContextTracer()
-        tracer._spans_list = [mock.Mock()]
+        tracer.start_span('span')
+        tracer.finish()
+
+        self.assertEqual(tracer._spans_list, [])
+
+    def test_end_leftover_spans(self):
+        tracer = context_tracer.ContextTracer()
+        tracer._spans_list = [span.Span(name='span')]
         tracer.finish()
 
         self.assertEqual(tracer._spans_list, [])
@@ -132,7 +145,7 @@ def test_end_span_active(self, mock_current_span):
 
         self.assertTrue(mock_span.finish.called)
         self.assertEqual(tracer.span_context.span_id, parent_span_id)
-        self.assertTrue(tracer.exporter.export.called)
+        self.assertFalse(tracer.exporter.export.called)
 
     @mock.patch.object(context_tracer.ContextTracer, 'current_span')
     def test_end_span_without_parent(self, mock_current_span):
@@ -155,30 +168,16 @@ def test_end_span_without_parent(self, mock_current_span):
         cur_span = get_current_span()
         self.assertIsNone(cur_span)
 
-    @mock.patch.object(context_tracer.ContextTracer, 'current_span')
-    def test_end_span_batch_export(self, mock_current_span):
-        from opencensus.trace import span
-
-        span = span.Span(name='test')
+    def test_end_span_batch_export(self):
         exporter = mock.Mock()
         tracer = context_tracer.ContextTracer(exporter=exporter)
-        tracer._spans_list = [span]
-        mock_span = mock.Mock()
-        mock_span.name = 'span'
-        mock_span.children = []
-        mock_span.status = None
-        mock_span.links = None
-        mock_span.stack_trace = None
-        mock_span.time_events = None
-        mock_span.attributes = {}
-        mock_span.__iter__ = mock.Mock(
-            return_value=iter([mock_span]))
+        span = tracer.start_span('test')
         parent_span_id = '6e0c63257de34c92'
-        mock_span.parent_span.span_id = parent_span_id
-        mock_current_span.return_value = mock_span
+        span.parent_span.span_id = parent_span_id
+        span.finish = mock.Mock()
         tracer.end_span()
 
-        self.assertTrue(mock_span.finish.called)
+        self.assertTrue(span.finish.called)
         self.assertEqual(tracer.span_context.span_id, parent_span_id)
         self.assertTrue(tracer.exporter.export.called)
 
@@ -208,4 +207,5 @@ def test_add_attribute_to_current_span(self):
 
         tracer.add_attribute_to_current_span(attribute_key, attribute_value)
 
-        span1.add_attribute.assert_called_once_with(attribute_key, attribute_value)
+        span1.add_attribute.assert_called_once_with(attribute_key,
+                                                    attribute_value)