Fix handling of missing headers in the b3 codec (#215)

* Fix handling of missing headers in the b3 codec

The b3 codec uses .get() to extract headers and so does not throw a key
error if headers are missing.  However many of these headers are decoded
via "header_to_hex" which throws an exception on a None value.  Other
codecs handle the missing header case by returning None, update the b3
codec to follow this convention.

Signed-off-by: Caleb Howe <[email protected]>

Author: cshowe

jaeger_client/codecs.py

@@ -298,19 +298,24 @@ def inject(self, span_context, carrier):
     def extract(self, carrier):
         if not isinstance(carrier, dict):
             raise InvalidCarrierException('carrier not a dictionary')
-        lowercase_keys = dict([(k.lower(), k) for k in carrier])
-        trace_id = header_to_hex(carrier.get(lowercase_keys.get(self._trace_header_lc)))
-        span_id = header_to_hex(carrier.get(lowercase_keys.get(self._span_header_lc)))
-        parent_id = carrier.get(lowercase_keys.get(self._parent_span_header_lc))
-        if parent_id:
-            parent_id = header_to_hex(parent_id)
+        trace_id = span_id = parent_id = None
         flags = 0x00
-        sampled = carrier.get(lowercase_keys.get(self._sampled_header_lc))
-        if sampled == '1':
-            flags |= SAMPLED_FLAG
-        debug = carrier.get(lowercase_keys.get(self._flags_header_lc))
-        if debug == '1':
-            flags |= DEBUG_FLAG
+        for header_key, header_value in six.iteritems(carrier):
+            if header_value is None:
+                continue
+            lower_key = header_key.lower()
+            if lower_key == self._trace_header_lc:
+                trace_id = header_to_hex(header_value)
+            elif lower_key == self._span_header_lc:
+                span_id = header_to_hex(header_value)
+            elif lower_key == self._parent_span_header_lc:
+                parent_id = header_to_hex(header_value)
+            elif lower_key == self._sampled_header_lc and header_value == '1':
+                flags |= SAMPLED_FLAG
+            elif lower_key == self._flags_header_lc and header_value == '1':
+                flags |= DEBUG_FLAG
+        if not trace_id or not span_id:
+            return None
         return SpanContext(trace_id=trace_id, span_id=span_id,
                            parent_id=parent_id, flags=flags,
                            baggage=None)

tests/test_codecs.py

@@ -339,6 +339,20 @@ def test_b3_extract(self):
         span_context = codec.extract(carrier)
         assert span_context.flags == 0x01
 
+        # validate present debug header with falsy value
+        carrier = {'X-b3-SpanId': 'a2fb4a1d1a96d312', 'X-B3-flags': '0',
+                   'X-B3-traceId': '463ac35c9f6413ad48485a3953bb6124'}
+        span_context = codec.extract(carrier)
+        assert span_context.flags == 0x00
+
+        # validate missing context
+        assert codec.extract({}) is None
+
+        # validate explicit none in context
+        carrier = {'X-b3-SpanId': None,
+                   'X-B3-traceId': '463ac35c9f6413ad48485a3953bb6124'}
+        assert codec.extract(carrier) is None
+
         # validate invalid hex string
         with self.assertRaises(SpanContextCorruptedException):
             codec.extract({'x-B3-TraceId': 'a2fb4a1d1a96d312z'})