support tracestate HTTP propagation (#261)

Author: reyang

opencensus/trace/propagation/trace_context_http_header_format.py

@@ -17,8 +17,11 @@
 
 from opencensus.trace.span_context import SpanContext
 from opencensus.trace.trace_options import TraceOptions
+from opencensus.trace.propagation.tracestate_string_format \
+    import TracestateStringFormatter
 
-_TRACE_PARENT_HEADER_NAME = 'traceparent'
+_TRACEPARENT_HEADER_NAME = 'traceparent'
+_TRACESTATE_HEADER_NAME = 'tracestate'
 _TRACE_CONTEXT_HEADER_FORMAT = \
     '([0-9a-f]{2})(-([0-9a-f]{32}))(-([0-9a-f]{16}))?(-([0-9a-f]{2}))?'
 _TRACE_CONTEXT_HEADER_RE = re.compile(_TRACE_CONTEXT_HEADER_FORMAT)
@@ -88,11 +91,19 @@ def from_headers(self, headers):
         """
         if headers is None:
             return SpanContext()
-        header = headers.get(_TRACE_PARENT_HEADER_NAME)
+        header = headers.get(_TRACEPARENT_HEADER_NAME)
         if header is None:
             return SpanContext()
         header = str(header.encode('utf-8'))
-        return self.from_header(header)
+
+        span_context = self.from_header(header)
+
+        header = headers.get(_TRACESTATE_HEADER_NAME)
+        if header is not None:
+            span_context.tracestate = \
+                TracestateStringFormatter().from_string(header)
+
+        return span_context
 
     def to_header(self, span_context):
         """Convert a SpanContext object to header string, using version 0.
@@ -128,6 +139,11 @@ def to_headers(self, span_context):
         :rtype: dict
         :returns: W3C Distributed Tracing headers.
         """
-        return {
-            _TRACE_PARENT_HEADER_NAME: self.to_header(span_context),
+        headers = {
+            _TRACEPARENT_HEADER_NAME: self.to_header(span_context),
         }
+        tracestate = span_context.tracestate
+        if tracestate:
+            headers[_TRACESTATE_HEADER_NAME] = \
+                TracestateStringFormatter().to_string(tracestate)
+        return headers

tests/unit/trace/ext/httplib/test_httplib_trace.py

@@ -177,6 +177,7 @@ def start_span(self):
         span.context_tracer.span_context = mock.Mock()
         span.context_tracer.span_context.trace_id = '123'
         span.context_tracer.span_context.span_id = '456'
+        span.context_tracer.span_context.tracestate = None
         self.span = span
         return span
 

tests/unit/trace/propagation/test_trace_context_http_header_format.py

@@ -26,7 +26,7 @@ def test_from_header_no_header(self):
             TraceContextPropagator()
         span_context = propagator.from_header(None)
 
-        assert isinstance(span_context, SpanContext)
+        self.assertTrue(isinstance(span_context, SpanContext))
 
     def test_from_headers_none(self):
         from opencensus.trace.span_context import SpanContext
@@ -35,7 +35,7 @@ def test_from_headers_none(self):
             TraceContextPropagator()
         span_context = propagator.from_headers(None)
 
-        assert isinstance(span_context, SpanContext)
+        self.assertTrue(isinstance(span_context, SpanContext))
 
     def test_from_headers_empty(self):
         from opencensus.trace.span_context import SpanContext
@@ -44,7 +44,20 @@ def test_from_headers_empty(self):
             TraceContextPropagator()
         span_context = propagator.from_headers({})
 
-        assert isinstance(span_context, SpanContext)
+        self.assertTrue(isinstance(span_context, SpanContext))
+
+    def test_from_headers_with_tracestate(self):
+        from opencensus.trace.span_context import SpanContext
+
+        propagator = trace_context_http_header_format.\
+            TraceContextPropagator()
+        span_context = propagator.from_headers({
+            'traceparent': '00-a66ee7820d074463aff4c617a63e929f-91e072af6a404137-01',
+            'tracestate': 'foo=1,bar=2,baz=3',
+        })
+
+        self.assertTrue(isinstance(span_context, SpanContext))
+        self.assertTrue(span_context.tracestate)
 
     def test_header_type_error(self):
         header = 1234
@@ -63,7 +76,7 @@ def test_header_version_not_support(self):
             TraceContextPropagator()
         span_context = propagator.from_header(header)
 
-        assert isinstance(span_context, SpanContext)
+        self.assertTrue(isinstance(span_context, SpanContext))
 
     def test_header_match(self):
         # Trace option is not enabled.
@@ -153,7 +166,7 @@ def test_to_header(self):
 
         self.assertEqual(header, expected_header)
 
-    def test_to_headers(self):
+    def test_to_headers_without_tracestate(self):
         from opencensus.trace import span_context
         from opencensus.trace import trace_options
 
@@ -168,8 +181,58 @@ def test_to_headers(self):
             TraceContextPropagator()
 
         headers = propagator.to_headers(span_context)
-        expected_headers = {
-            'traceparent': '00-{}-{}-01'.format(trace_id, span_id_hex),
-        }
 
-        self.assertEqual(headers, expected_headers)
+        self.assertTrue('traceparent' in headers)
+        self.assertEqual(headers['traceparent'],
+            '00-{}-{}-01'.format(trace_id, span_id_hex))
+
+        self.assertFalse('tracestate' in headers)
+
+    def test_to_headers_with_empty_tracestate(self):
+        from opencensus.trace import span_context
+        from opencensus.trace import trace_options
+        from opencensus.trace.tracestate import Tracestate
+
+        trace_id = '6e0c63257de34c92bf9efcd03927272e'
+        span_id_hex = '00f067aa0ba902b7'
+        span_context = span_context.SpanContext(
+            trace_id=trace_id,
+            span_id=span_id_hex,
+            tracestate=Tracestate(),
+            trace_options=trace_options.TraceOptions('1'))
+
+        propagator = trace_context_http_header_format.\
+            TraceContextPropagator()
+
+        headers = propagator.to_headers(span_context)
+
+        self.assertTrue('traceparent' in headers)
+        self.assertEqual(headers['traceparent'],
+            '00-{}-{}-01'.format(trace_id, span_id_hex))
+
+        self.assertFalse('tracestate' in headers)
+
+    def test_to_headers_with_tracestate(self):
+        from opencensus.trace import span_context
+        from opencensus.trace import trace_options
+        from opencensus.trace.tracestate import Tracestate
+
+        trace_id = '6e0c63257de34c92bf9efcd03927272e'
+        span_id_hex = '00f067aa0ba902b7'
+        span_context = span_context.SpanContext(
+            trace_id=trace_id,
+            span_id=span_id_hex,
+            tracestate=Tracestate(foo = "xyz"),
+            trace_options=trace_options.TraceOptions('1'))
+
+        propagator = trace_context_http_header_format.\
+            TraceContextPropagator()
+
+        headers = propagator.to_headers(span_context)
+
+        self.assertTrue('traceparent' in headers)
+        self.assertEqual(headers['traceparent'],
+            '00-{}-{}-01'.format(trace_id, span_id_hex))
+
+        self.assertTrue('tracestate' in headers)
+        self.assertEqual(headers['tracestate'], 'foo=xyz')