Implement B3 propagation (census-instrumentation#265)

kornholi · reyang · commit f029bcdda643 · 2019-02-07T15:44:34.000-08:00
* Implement B3 propagation

* Update copyright year

* Clarify comments a little bit

* Fix to_headers

* Implement deserialization for single-header b3 format

* Adjust how TraceOptions is constructed

* Formatting
diff --git a/opencensus/trace/propagation/b3_format.py b/opencensus/trace/propagation/b3_format.py
@@ -0,0 +1,117 @@
+# Copyright 2018, OpenCensus Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+from opencensus.trace.span_context import SpanContext, INVALID_SPAN_ID
+from opencensus.trace.trace_options import TraceOptions
+
+_STATE_HEADER_KEY = 'b3'
+_TRACE_ID_KEY = 'x-b3-traceid'
+_SPAN_ID_KEY = 'x-b3-spanid'
+_SAMPLED_KEY = 'x-b3-sampled'
+
+
+class B3FormatPropagator(object):
+    """Propagator for the B3 HTTP header format.
+
+    See: https://github.com/openzipkin/b3-propagation
+    """
+
+    def from_headers(self, headers):
+        """Generate a SpanContext object from B3 propagation headers.
+
+        :type headers: dict
+        :param headers: HTTP request headers.
+
+        :rtype: :class:`~opencensus.trace.span_context.SpanContext`
+        :returns: SpanContext generated from B3 propagation headers.
+        """
+        if headers is None:
+            return SpanContext(from_header=False)
+
+        trace_id, span_id, sampled = None, None, None
+
+        state = headers.get(_STATE_HEADER_KEY)
+        if state:
+            fields = state.split('-', 4)
+
+            if len(fields) == 1:
+                sampled = fields[0]
+            elif len(fields) == 2:
+                trace_id, span_id = fields
+            elif len(fields) == 3:
+                trace_id, span_id, sampled = fields
+            elif len(fields) == 4:
+                trace_id, span_id, sampled, _parent_span_id = fields
+            else:
+                return SpanContext(from_header=False)
+        else:
+            trace_id = headers.get(_TRACE_ID_KEY)
+            span_id = headers.get(_SPAN_ID_KEY)
+            sampled = headers.get(_SAMPLED_KEY)
+
+        if sampled is not None:
+            if len(sampled) != 1:
+                return SpanContext(from_header=False)
+
+            sampled = sampled in ('1', 'd')
+        else:
+            # If there's no incoming sampling decision, it was deferred to us.
+            # Even though we set it to False here, we might still sample
+            # depending on the tracer configuration.
+            sampled = False
+
+        trace_options = TraceOptions()
+        trace_options.set_enabled(sampled)
+
+        # TraceId and SpanId headers both have to exist
+        if not trace_id or not span_id:
+            return SpanContext(trace_options=trace_options)
+
+        # Convert 64-bit trace ids to 128-bit
+        if len(trace_id) == 16:
+            trace_id = '0'*16 + trace_id
+
+        span_context = SpanContext(
+            trace_id=trace_id,
+            span_id=span_id,
+            trace_options=trace_options,
+            from_header=True
+        )
+
+        return span_context
+
+    def to_headers(self, span_context):
+        """Convert a SpanContext object to B3 propagation headers.
+
+        :type span_context:
+            :class:`~opencensus.trace.span_context.SpanContext`
+        :param span_context: SpanContext object.
+
+        :rtype: dict
+        :returns: B3 propagation headers.
+        """
+
+        if not span_context.span_id:
+            span_id = INVALID_SPAN_ID
+        else:
+            span_id = span_context.span_id
+
+        sampled = span_context.trace_options.enabled
+
+        return {
+            _TRACE_ID_KEY: span_context.trace_id,
+            _SPAN_ID_KEY: span_id,
+            _SAMPLED_KEY: '1' if sampled else '0'
+        }
diff --git a/opencensus/trace/span_context.py b/opencensus/trace/span_context.py
@@ -16,6 +16,7 @@
 
 import logging
 import re
+import six
 import uuid
 
 from opencensus.trace import trace_options
@@ -100,7 +101,7 @@ def _check_span_id(self, span_id):
         """
         if span_id is None:
             return None
-        assert isinstance(span_id, str)
+        assert isinstance(span_id, six.string_types)
 
         if span_id is INVALID_SPAN_ID:
             logging.warning(
@@ -130,7 +131,7 @@ def _check_trace_id(self, trace_id):
         :rtype: str
         :returns: Trace_id for the current context.
         """
-        assert isinstance(trace_id, str)
+        assert isinstance(trace_id, six.string_types)
 
         if trace_id is _INVALID_TRACE_ID:
             logging.warning(
diff --git a/tests/unit/trace/propagation/test_b3_format.py b/tests/unit/trace/propagation/test_b3_format.py
@@ -0,0 +1,176 @@
+# Copyright 2018, OpenCensus Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+import mock
+
+from opencensus.trace.span_context import INVALID_SPAN_ID
+from opencensus.trace.propagation import b3_format
+
+
+class TestB3FormatPropagator(unittest.TestCase):
+
+    def test_from_headers_no_headers(self):
+        propagator = b3_format.B3FormatPropagator()
+        span_context = propagator.from_headers(None)
+
+        self.assertFalse(span_context.from_header)
+
+    def test_from_headers_keys_exist(self):
+        test_trace_id = '6e0c63257de34c92bf9efcd03927272e'
+        test_span_id = '00f067aa0ba902b7'
+        test_sampled = '1'
+
+        headers = {
+            b3_format._TRACE_ID_KEY: test_trace_id,
+            b3_format._SPAN_ID_KEY: test_span_id,
+            b3_format._SAMPLED_KEY: test_sampled,
+        }
+
+        propagator = b3_format.B3FormatPropagator()
+        span_context = propagator.from_headers(headers)
+
+        self.assertEqual(span_context.trace_id, test_trace_id)
+        self.assertEqual(span_context.span_id, test_span_id)
+        self.assertEqual(
+            span_context.trace_options.enabled,
+            bool(test_sampled)
+        )
+
+    def test_from_headers_keys_not_exist(self):
+        propagator = b3_format.B3FormatPropagator()
+        span_context = propagator.from_headers({})
+
+        self.assertIsNotNone(span_context.trace_id)
+        self.assertIsNone(span_context.span_id)
+        self.assertFalse(span_context.trace_options.enabled)
+
+    def test_from_headers_64bit_traceid(self):
+        test_trace_id = 'bf9efcd03927272e'
+        test_span_id = '00f067aa0ba902b7'
+
+        headers = {
+            b3_format._TRACE_ID_KEY: test_trace_id,
+            b3_format._SPAN_ID_KEY: test_span_id,
+        }
+
+        propagator = b3_format.B3FormatPropagator()
+        span_context = propagator.from_headers(headers)
+
+        converted_trace_id = "0"*16 + test_trace_id
+
+        self.assertEqual(span_context.trace_id, converted_trace_id)
+        self.assertEqual(span_context.span_id, test_span_id)
+
+    def test_to_headers_has_span_id(self):
+        test_trace_id = '6e0c63257de34c92bf9efcd03927272e'
+        test_span_id = '00f067aa0ba902b7'
+        test_options = '1'
+
+        span_context = mock.Mock()
+        span_context.trace_id = test_trace_id
+        span_context.span_id = test_span_id
+        span_context.trace_options.trace_options_byte = test_options
+
+        propagator = b3_format.B3FormatPropagator()
+        headers = propagator.to_headers(span_context)
+
+        self.assertEqual(headers[b3_format._TRACE_ID_KEY], test_trace_id)
+        self.assertEqual(headers[b3_format._SPAN_ID_KEY], test_span_id)
+        self.assertEqual(headers[b3_format._SAMPLED_KEY], test_options)
+
+    def test_to_headers_no_span_id(self):
+        test_trace_id = '6e0c63257de34c92bf9efcd03927272e'
+        test_options = '1'
+
+        span_context = mock.Mock()
+        span_context.trace_id = test_trace_id
+        span_context.span_id = None
+        span_context.trace_options.trace_options_byte = test_options
+
+        propagator = b3_format.B3FormatPropagator()
+        headers = propagator.to_headers(span_context)
+
+        self.assertEqual(headers[b3_format._TRACE_ID_KEY], test_trace_id)
+        self.assertEqual(headers.get(b3_format._SPAN_ID_KEY), INVALID_SPAN_ID)
+        self.assertEqual(headers[b3_format._SAMPLED_KEY], test_options)
+
+    def test_from_single_header_keys_exist(self):
+        trace_id = "80f198ee56343ba864fe8b2a57d3eff7"
+        span_id = "e457b5a2e4d86bd1"
+
+        headers = {
+            'b3': "{}-{}-d-05e3ac9a4f6e3b90".format(trace_id, span_id)
+        }
+        propagator = b3_format.B3FormatPropagator()
+        span_context = propagator.from_headers(headers)
+
+        self.assertEqual(span_context.trace_id, trace_id)
+        self.assertEqual(span_context.span_id, span_id)
+        self.assertEqual(span_context.trace_options.enabled, True)
+
+    def test_from_headers_invalid_single_header(self):
+        headers = {
+            'b3': "01234567890123456789012345678901;o=1"
+        }
+        propagator = b3_format.B3FormatPropagator()
+        span_context = propagator.from_headers(headers)
+
+        self.assertFalse(span_context.from_header)
+
+    def test_from_headers_invalid_single_header_fields(self):
+        headers = {
+            'b3': "a-b-c-d-e-f-g"
+        }
+        propagator = b3_format.B3FormatPropagator()
+        span_context = propagator.from_headers(headers)
+
+        self.assertFalse(span_context.from_header)
+
+    def test_from_single_header_deny_sampling(self):
+        headers = {
+            'b3': "0"
+        }
+        propagator = b3_format.B3FormatPropagator()
+        span_context = propagator.from_headers(headers)
+
+        self.assertEqual(span_context.trace_options.enabled, False)
+
+    def test_from_single_header_defer_sampling(self):
+        trace_id = "80f198ee56343ba864fe8b2a57d3eff7"
+        span_id = "e457b5a2e4d86bd1"
+        headers = {
+            'b3': "{}-{}".format(trace_id, span_id)
+        }
+        propagator = b3_format.B3FormatPropagator()
+        span_context = propagator.from_headers(headers)
+
+        self.assertEqual(span_context.trace_id, trace_id)
+        self.assertEqual(span_context.span_id, span_id)
+
+    def test_from_single_header_precedence(self):
+        headers = {
+            'b3': "80f198ee56343ba864fe8b2a57d3eff7-e457b5a2e4d86bd1-1",
+            'X-B3-TraceId': '6e0c63257de34c92bf9efcd03927272e',
+            'X-B3-SpanId': '00f067aa0ba902b7'
+        }
+        propagator = b3_format.B3FormatPropagator()
+        span_context = propagator.from_headers(headers)
+
+        self.assertEqual(
+            span_context.trace_id,
+            "80f198ee56343ba864fe8b2a57d3eff7"
+        )
+        self.assertEqual(span_context.span_id, "e457b5a2e4d86bd1")
+        self.assertEqual(span_context.trace_options.enabled, True)
