Add support for reading B3 single header (#165)

* Add support for reading B3 single header * Update lib/zipkin-tracer/zipkin_b3_single_header_format.rb Co-Authored-By: Adrian Cole <[email protected]> * Add magic comment about frozen strings * Avoid instantiating
openzipkin · Oct 9, 2019 · 0a8be52 · 0a8be52
1 parent a1b1523
commit 0a8be52
Show file tree

Hide file tree

Showing 7 changed files with 213 additions and 11 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -1,3 +1,6 @@
+# 0.40.0
+* Add support for reading B3 single header.
+
 # 0.39.2
 * Make Faraday 0.13 the minimum requirement.
 

diff --git a/lib/zipkin-tracer.rb b/lib/zipkin-tracer.rb
@@ -5,6 +5,7 @@
 require 'zipkin-tracer/trace_container'
 require 'zipkin-tracer/trace_generator'
 require 'zipkin-tracer/trace_wrapper'
+require 'zipkin-tracer/zipkin_b3_single_header_format'
 
 begin
   require 'faraday'

diff --git a/lib/zipkin-tracer/rack/zipkin_env.rb b/lib/zipkin-tracer/rack/zipkin_env.rb
@@ -10,33 +10,44 @@ def initialize(env, config)
     end
 
     def trace_id(default_flags = Trace::Flags::EMPTY)
-      trace_id, span_id, parent_span_id, shared = retrieve_or_generate_ids
-      sampled = sampled_header_value(@env['HTTP_X_B3_SAMPLED'])
-      flags = (@env['HTTP_X_B3_FLAGS'] || default_flags).to_i
+      trace_id, span_id, parent_span_id, sampled, flags, shared = retrieve_or_generate_ids
+      sampled = sampled_header_value(sampled)
+      flags = (flags || default_flags).to_i
       Trace::TraceId.new(trace_id, parent_span_id, span_id, sampled, flags, shared)
     end
 
+    def called_with_zipkin_b3_single_header?
+      @called_with_zipkin_b3_single_header ||= @env.key?(B3_SINGLE_HEADER)
+    end
+
     def called_with_zipkin_headers?
       @called_with_zipkin_headers ||= B3_REQUIRED_HEADERS.all? { |key| @env.key?(key) }
     end
 
     private
 
-    B3_REQUIRED_HEADERS = %w(HTTP_X_B3_TRACEID HTTP_X_B3_SPANID).freeze
-    B3_OPT_HEADERS = %w(HTTP_X_B3_PARENTSPANID HTTP_X_B3_SAMPLED HTTP_X_B3_FLAGS).freeze
+    B3_SINGLE_HEADER = 'HTTP_B3'.freeze
+    B3_REQUIRED_HEADERS = %w[HTTP_X_B3_TRACEID HTTP_X_B3_SPANID].freeze
+    B3_OPT_HEADERS = %w[HTTP_X_B3_PARENTSPANID HTTP_X_B3_SAMPLED HTTP_X_B3_FLAGS].freeze
 
     def retrieve_or_generate_ids
-      if called_with_zipkin_headers?
-        trace_id, span_id = @env.values_at(*B3_REQUIRED_HEADERS)
-        parent_span_id = @env['HTTP_X_B3_PARENTSPANID']
+      if called_with_zipkin_b3_single_header?
+        trace_id, span_id, parent_span_id, sampled, flags =
+          B3SingleHeaderFormat.parse_from_header(@env[B3_SINGLE_HEADER])
         shared = true
-      else
+      elsif called_with_zipkin_headers?
+        trace_id, span_id, parent_span_id, sampled, flags = @env.values_at(*B3_REQUIRED_HEADERS, *B3_OPT_HEADERS)
+        shared = true
+      end
+
+      unless trace_id
         span_id = TraceGenerator.new.generate_id
         trace_id = TraceGenerator.new.generate_id_from_span_id(span_id)
         parent_span_id = nil
         shared = false
       end
-      [trace_id, span_id, parent_span_id, shared]
+
+      [trace_id, span_id, parent_span_id, sampled, flags, shared]
     end
 
     def new_sampled_header_value(sampled)

diff --git a/lib/zipkin-tracer/version.rb b/lib/zipkin-tracer/version.rb
@@ -1,3 +1,3 @@
 module ZipkinTracer
-  VERSION = '0.39.2'.freeze
+  VERSION = '0.40.0'.freeze
 end
diff --git a/lib/zipkin-tracer/zipkin_b3_single_header_format.rb b/lib/zipkin-tracer/zipkin_b3_single_header_format.rb
@@ -0,0 +1,28 @@
+# frozen_string_literal: true
+
+module ZipkinTracer
+  # This format corresponds to the propagation key "b3" (or "B3").
+  # b3: {x-b3-traceid}-{x-b3-spanid}-{if x-b3-flags 'd' else x-b3-sampled}-{x-b3-parentspanid}
+  # For details, see: https://github.com/openzipkin/b3-propagation
+  class B3SingleHeaderFormat
+    def self.parse_from_header(b3_single_header)
+      if b3_single_header.size == 1
+        flag = b3_single_header
+      else
+        trace_id, span_id, flag, parent_span_id = b3_single_header.split('-')
+      end
+      [trace_id, span_id, parent_span_id, parse_sampled(flag), parse_flags(flag)]
+    end
+
+    def self.parse_sampled(flag)
+      case flag
+      when '1', '0'
+        flag
+      end
+    end
+
+    def self.parse_flags(flag)
+      flag == 'd' ? Trace::Flags::DEBUG : Trace::Flags::EMPTY
+    end
+  end
+end
diff --git a/spec/lib/rack/zipkin_env_spec.rb b/spec/lib/rack/zipkin_env_spec.rb
@@ -213,4 +213,90 @@ def mock_env(params = {}, path = '/')
       end
     end
   end
+
+  context 'with zipkin b3 single header' do
+    let(:id) { "e457b5a2e4d86bd1" }
+    let(:zipkin_headers) { { 'HTTP_B3' => b3_single_header } }
+    let(:env) { mock_env(zipkin_headers) }
+
+    context 'not yet sampled root span' do
+      let(:b3_single_header) { "#{id}-#{id}" }
+
+      it '#called_with_zipkin_b3_single_header? returns true' do
+        expect(zipkin_env.called_with_zipkin_b3_single_header?).to eq(true)
+      end
+
+      it 'shared is true' do
+        expect(zipkin_env.trace_id.shared).to eq(true)
+      end
+
+      it 'sampling information is set' do
+        # Because sample rate == 1
+        expect(zipkin_env.trace_id.sampled?).to eq(true)
+      end
+
+      context 'parent_id is not provided' do
+        it 'uses the trace_id and span_id' do
+          trace_id = zipkin_env.trace_id
+          expect(trace_id.trace_id.to_s).to eq(id)
+          expect(trace_id.span_id.to_s).to eq(id)
+        end
+
+        it 'parent_id is empty' do
+          expect(zipkin_env.trace_id.parent_id).to eq(nil)
+        end
+      end
+    end
+
+    context 'all information is provided' do
+      let(:parent_id) { "05e3ac9a4f6e3b90" }
+      let(:b3_single_header) { "#{id}-#{id}-1-#{parent_id}" }
+
+      it 'uses the trace_id and span_id' do
+        trace_id = zipkin_env.trace_id
+        expect(trace_id.trace_id.to_s).to eq(id)
+        expect(trace_id.span_id.to_s).to eq(id)
+      end
+
+      it 'uses the parent_id' do
+        expect(zipkin_env.trace_id.parent_id.to_s).to eq(parent_id.to_s)
+      end
+
+      it 'uses the sampling information' do
+        expect(zipkin_env.trace_id.sampled?).to eq(true)
+      end
+
+      it 'uses the flags' do
+        expect(zipkin_env.trace_id.flags.to_i).to eq(0)
+      end
+    end
+
+    context 'debug flag only' do
+      let(:b3_single_header) { "d" }
+
+      it 'generates a trace_id and a span_id' do
+        trace_id = zipkin_env.trace_id
+        expect(trace_id.trace_id).not_to eq(nil)
+        expect(trace_id.span_id).not_to eq(nil)
+        expect(trace_id.span_id.to_s).to eq(trace_id.trace_id.to_s)
+      end
+
+      it 'parent_id is nil' do
+        expect(zipkin_env.trace_id.parent_id).to eq(nil)
+      end
+
+      it 'flags is debug' do
+        expect(zipkin_env.trace_id.flags).to eq(Trace::Flags::DEBUG)
+      end
+
+      it 'sampling information is set' do
+        # Because sample rate == 1
+        expect(zipkin_env.trace_id.sampled?).to eq(true)
+      end
+
+      it 'shared is false' do
+        expect(zipkin_env.trace_id.shared).to eq(false)
+      end
+    end
+  end
 end
diff --git a/spec/lib/zipkin_b3_single_header_format_spec.rb b/spec/lib/zipkin_b3_single_header_format_spec.rb
@@ -0,0 +1,73 @@
+require 'spec_helper'
+
+describe ZipkinTracer::B3SingleHeaderFormat do
+  let(:b3_single_header_format) { described_class.parse_from_header(b3_single_header) }
+
+  context 'child span' do
+    let(:b3_single_header) { '80f198ee56343ba864fe8b2a57d3eff7-e457b5a2e4d86bd1-1-05e3ac9a4f6e3b90' }
+
+    it 'has all fields' do
+      expect(b3_single_header_format.to_a).to eq(
+          ['80f198ee56343ba864fe8b2a57d3eff7', 'e457b5a2e4d86bd1', '05e3ac9a4f6e3b90', '1', 0]
+        )
+    end
+  end
+
+  context 'sampled root span' do
+    let(:b3_single_header) { '80f198ee56343ba864fe8b2a57d3eff7-e457b5a2e4d86bd1-1' }
+
+    it 'does not have parent_span_id' do
+      expect(b3_single_header_format.to_a).to eq(['80f198ee56343ba864fe8b2a57d3eff7', 'e457b5a2e4d86bd1', nil, '1', 0])
+    end
+  end
+
+  context 'not yet sampled root span' do
+    let(:b3_single_header) { '80f198ee56343ba864fe8b2a57d3eff7-e457b5a2e4d86bd1' }
+
+    it 'does not have parent_span_id and sampled' do
+      expect(b3_single_header_format.to_a).to eq(['80f198ee56343ba864fe8b2a57d3eff7', 'e457b5a2e4d86bd1', nil, nil, 0])
+    end
+  end
+
+  context 'debug RPC child span' do
+    let(:b3_single_header) { '80f198ee56343ba864fe8b2a57d3eff7-e457b5a2e4d86bd1-d-05e3ac9a4f6e3b90' }
+
+    it 'has debug flag' do
+      expect(b3_single_header_format.to_a).to eq(
+          ['80f198ee56343ba864fe8b2a57d3eff7', 'e457b5a2e4d86bd1', '05e3ac9a4f6e3b90', nil, 1]
+        )
+    end
+  end
+
+  context 'do not sample flag only' do
+    let(:b3_single_header) { '0' }
+
+    it 'has do not sample flag only' do
+      expect(b3_single_header_format.to_a).to eq([nil, nil, nil, '0', 0])
+    end
+  end
+
+  context 'sampled flag only' do
+    let(:b3_single_header) { '1' }
+
+    it 'has sampled flag only' do
+      expect(b3_single_header_format.to_a).to eq([nil, nil, nil, '1', 0])
+    end
+  end
+
+  context 'debug flag only' do
+    let(:b3_single_header) { 'd' }
+
+    it 'has debug flag only' do
+      expect(b3_single_header_format.to_a).to eq([nil, nil, nil, nil, 1])
+    end
+  end
+
+  context 'unknown flag only' do
+    let(:b3_single_header) { 'u' }
+
+    it 'has nothing' do
+      expect(b3_single_header_format.to_a).to eq([nil, nil, nil, nil, 0])
+    end
+  end
+end