Remove span_stack and store parent_span in span object (#52)

Author: liyanhui1228

trace/opencensus/trace/execution_context.py

@@ -47,6 +47,14 @@ def get_opencensus_attr(attr_key):
     return None
 
 
+def get_current_span():
+    return getattr(_thread_local, 'current_span', None)
+
+
+def set_current_span(current_span):
+    setattr(_thread_local, 'current_span', current_span)
+
+
 def clear():
     """Clear the thread local, used in test."""
     _thread_local.__dict__.clear()

trace/opencensus/trace/span.py

@@ -17,6 +17,7 @@
 
 from opencensus.trace.enums import Enum
 from opencensus.trace.span_context import generate_span_id
+from opencensus.trace.tracer import base
 
 
 class Span(object):
@@ -37,8 +38,8 @@ class Span(object):
                  distinguished using RPC_CLIENT and RPC_SERVER to identify
                  queueing latency associated with the span.
 
-    :type parent_span_id: int
-    :param parent_span_id: (Optional) ID of the parent span.
+    :type parent_span: :class:`~opencensus.trace.span.Span`
+    :param parent_span: (Optional) Parent span.
 
     :type labels: dict
     :param labels: Collection of labels associated with the span.
@@ -70,15 +71,15 @@ def __init__(
             self,
             name,
             kind=Enum.SpanKind.SPAN_KIND_UNSPECIFIED,
-            parent_span_id=None,
+            parent_span=None,
             labels=None,
             start_time=None,
             end_time=None,
             span_id=None,
             context_tracer=None):
         self.name = name
         self.kind = kind
-        self.parent_span_id = parent_span_id
+        self.parent_span = parent_span
         self.start_time = start_time
         self.end_time = end_time
 
@@ -88,6 +89,11 @@ def __init__(
         if labels is None:
             labels = {}
 
+        # Do not manipulate spans directly using the methods in Span Class,
+        # make sure to use the RequestTracer.
+        if parent_span is None:
+            parent_span = base.NullContextManager()
+
         self.labels = labels
         self.span_id = span_id
         self._child_spans = []
@@ -108,7 +114,7 @@ def span(self, name='child_span'):
         :rtype: :class: `~opencensus.trace.span.Span`
         :returns: A child Span to be added to the current span.
         """
-        child_span = Span(name, parent_span_id=self.span_id)
+        child_span = Span(name, parent_span=self)
         self._child_spans.append(child_span)
         return child_span
 
@@ -168,8 +174,8 @@ def format_span_json(span):
         'endTime': span.end_time,
     }
 
-    if span.parent_span_id is not None:
-        span_json['parentSpanId'] = span.parent_span_id
+    if span.parent_span is not None:
+        span_json['parentSpanId'] = span.parent_span.span_id
 
     if span.labels is not None:
         span_json['labels'] = span.labels

trace/opencensus/trace/tracer/base.py

@@ -78,8 +78,9 @@ class NullContextManager(object):
     """Empty object as a helper for faking Trace and Span when tracing is
     disabled.
     """
-    def __init__(self):
+    def __init__(self, span_id=None):
         self.name = None
+        self.span_id = span_id
 
     def __enter__(self):
         pass  # pragma: NO COVER

trace/opencensus/trace/tracer/context_tracer.py

@@ -14,6 +14,7 @@
 
 import logging
 
+from opencensus.trace import execution_context
 from opencensus.trace.span_context import SpanContext
 from opencensus.trace import labels_helper
 from opencensus.trace import span as trace_span
@@ -35,9 +36,6 @@ def __init__(self, span_context=None):
         self.trace_id = span_context.trace_id
         self.root_span_id = span_context.span_id
 
-        # Stack which maintains nested spans
-        self._span_stack = []
-
         # List of spans to report
         self._spans_list = []
 
@@ -76,14 +74,21 @@ def start_span(self, name='span'):
         :rtype: :class:`~opencensus.trace.span.Span`
         :returns: The Span object.
         """
-        parent_span_id = self.span_context.span_id
+        parent_span = self.current_span()
+
+        # If a span has remote parent span, then the parent_span.span_id
+        # should be the span_id from the request header.
+        if parent_span is None:
+            parent_span = base.NullContextManager(
+                span_id=self.span_context.span_id)
+
         span = trace_span.Span(
             name,
-            parent_span_id=parent_span_id,
+            parent_span=parent_span,
             context_tracer=self)
         self._spans_list.append(span)
-        self._span_stack.append(span)
         self.span_context.span_id = span.span_id
+        execution_context.set_current_span(span)
         span.start()
         return span
 
@@ -94,23 +99,23 @@ def end_span(self):
         """
         cur_span = self.current_span()
 
-        self._span_stack.pop()
+        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
 
-        if not self._span_stack:
-            self.span_context.span_id = self.root_span_id
+        if isinstance(cur_span.parent_span, trace_span.Span):
+            execution_context.set_current_span(cur_span.parent_span)
         else:
-            self.span_context.span_id = self._span_stack[-1].span_id
+            execution_context.set_current_span(None)
 
     def current_span(self):
         """Return the current span."""
-        try:
-            cur_span = self._span_stack[-1]
-        except IndexError:
-            logging.error('No span in the span stack.')
-            cur_span = base.NullContextManager()
+        current_span = execution_context.get_current_span()
 
-        return cur_span
+        return current_span
 
     def list_collected_spans(self):
         return self._spans_list

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

@@ -187,7 +187,7 @@ def test_process_request(self):
             '/http/method': 'GET',
         }
         self.assertEqual(span.labels, expected_labels)
-        self.assertEqual(span.parent_span_id, span_id)
+        self.assertEqual(span.parent_span.span_id, span_id)
 
         span_context = tracer.span_context
         self.assertEqual(span_context.trace_id, trace_id)

trace/tests/unit/ext/flask/test_flask_middleware.py

@@ -104,7 +104,7 @@ def test__before_request(self):
             }
 
             self.assertEqual(span.labels, expected_labels)
-            self.assertEqual(span.parent_span_id, span_id)
+            self.assertEqual(span.parent_span.span_id, span_id)
 
             span_context = tracer.span_context
             self.assertEqual(span_context.trace_id, trace_id)
@@ -116,6 +116,7 @@ def test_header_encoding(self):
         # in SpanContext because it cannot match the pattern for trace_id,
         # And a new trace_id will generate for the context.
         from opencensus.trace import execution_context
+        from opencensus.trace.tracer import base
 
         flask_trace_header = 'X_CLOUD_TRACE_CONTEXT'
         trace_id = "你好"
@@ -141,13 +142,14 @@ def test_header_encoding(self):
             }
 
             self.assertEqual(span.labels, expected_labels)
-            self.assertIsNone(span.parent_span_id)
+            assert isinstance(span.parent_span, base.NullContextManager)
 
             span_context = tracer.span_context
             self.assertNotEqual(span_context.trace_id, trace_id)
 
     def test_header_is_none(self):
         from opencensus.trace import execution_context
+        from opencensus.trace.tracer import base
 
         app = self.create_app()
         flask_middleware.FlaskMiddleware(app=app)
@@ -167,7 +169,7 @@ def test_header_is_none(self):
             }
 
             self.assertEqual(span.labels, expected_labels)
-            self.assertIsNone(span.parent_span_id)
+            assert isinstance(span.parent_span, base.NullContextManager)
 
     def test__after_request_not_sampled(self):
         from opencensus.trace.samplers import always_off

trace/tests/unit/test_request_tracer.py

@@ -188,11 +188,13 @@ def test_end_span_not_sampled(self):
         self.assertFalse(span_context.span_id.called)
 
     def test_end_span_sampled(self):
+        from opencensus.trace import execution_context
+
         sampler = mock.Mock()
         sampler.should_sample.return_value = True
         tracer = request_tracer.RequestTracer(sampler=sampler)
         span = mock.Mock()
-        tracer.tracer._span_stack.append(span)
+        execution_context.set_current_span(span)
         tracer.end_span()
 
         self.assertTrue(span.finish.called)
@@ -209,11 +211,13 @@ def test_current_span_not_sampled(self):
         assert isinstance(span, base.NullContextManager)
 
     def test_current_span_sampled(self):
+        from opencensus.trace import execution_context
+
         sampler = mock.Mock()
         sampler.should_sample.return_value = True
         tracer = request_tracer.RequestTracer(sampler=sampler)
         span = mock.Mock()
-        tracer.tracer._span_stack.append(span)
+        execution_context.set_current_span(span)
 
         result = tracer.current_span()
 

trace/tests/unit/test_span.py

@@ -46,7 +46,7 @@ def test_constructor_defaults(self):
         self.assertEqual(span.name, span_name)
         self.assertEqual(span.span_id, span_id)
         self.assertEqual(span.kind, Enum.SpanKind.SPAN_KIND_UNSPECIFIED)
-        self.assertIsNone(span.parent_span_id)
+        self.assertIsNone(span.parent_span)
         self.assertEqual(span.labels, {})
         self.assertIsNone(span.start_time)
         self.assertIsNone(span.end_time)
@@ -61,7 +61,7 @@ def test_constructor_explicit(self):
         span_id = 'test_span_id'
         span_name = 'test_span_name'
         kind = Enum.SpanKind.RPC_CLIENT
-        parent_span_id = 1234
+        parent_span = mock.Mock()
         start_time = datetime.utcnow().isoformat() + 'Z'
         end_time = datetime.utcnow().isoformat() + 'Z'
         labels = {
@@ -73,7 +73,7 @@ def test_constructor_explicit(self):
         span = self._make_one(
             name=span_name,
             kind=kind,
-            parent_span_id=parent_span_id,
+            parent_span=parent_span,
             labels=labels,
             start_time=start_time,
             end_time=end_time,
@@ -83,7 +83,7 @@ def test_constructor_explicit(self):
         self.assertEqual(span.name, span_name)
         self.assertEqual(span.span_id, span_id)
         self.assertEqual(span.kind, kind)
-        self.assertEqual(span.parent_span_id, parent_span_id)
+        self.assertEqual(span.parent_span, parent_span)
         self.assertEqual(span.labels, labels)
         self.assertEqual(span.start_time, start_time)
         self.assertEqual(span.end_time, end_time)
@@ -115,7 +115,7 @@ def test_span(self):
         self.assertEqual(result_child_span.name, child_span_name)
         self.assertEqual(result_child_span.span_id, span_id)
         self.assertEqual(result_child_span.kind, kind)
-        self.assertEqual(result_child_span.parent_span_id, root_span.span_id)
+        self.assertEqual(result_child_span.parent_span, root_span)
         self.assertEqual(result_child_span.labels, {})
         self.assertIsNone(result_child_span.start_time)
         self.assertIsNone(result_child_span.end_time)