From de1ca7dfa976d336418395e5bfc2e63d0319b647 Mon Sep 17 00:00:00 2001 From: Reiley Yang Date: Fri, 17 Aug 2018 14:00:36 -0700 Subject: [PATCH] support tracestate HTTP propagation --- .../trace_context_http_header_format.py | 26 ++++-- .../trace/ext/httplib/test_httplib_trace.py | 1 + .../test_trace_context_http_header_format.py | 81 ++++++++++++++++--- 3 files changed, 94 insertions(+), 14 deletions(-) diff --git a/opencensus/trace/propagation/trace_context_http_header_format.py b/opencensus/trace/propagation/trace_context_http_header_format.py index 4d3d4f609..c3b07f4de 100644 --- a/opencensus/trace/propagation/trace_context_http_header_format.py +++ b/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 diff --git a/tests/unit/trace/ext/httplib/test_httplib_trace.py b/tests/unit/trace/ext/httplib/test_httplib_trace.py index a50976a94..60bd4a23f 100644 --- a/tests/unit/trace/ext/httplib/test_httplib_trace.py +++ b/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 diff --git a/tests/unit/trace/propagation/test_trace_context_http_header_format.py b/tests/unit/trace/propagation/test_trace_context_http_header_format.py index 653d7836e..8e0f512e8 100644 --- a/tests/unit/trace/propagation/test_trace_context_http_header_format.py +++ b/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')