Merge pull request #36 from lightstep/ngauthier/rack-http-inject-extract

Rack HTTP Inject + Extract

Makefile

@@ -7,6 +7,7 @@ test:
 	bundle exec rake spec
 	ruby example.rb
 	ruby examples/fork_children/main.rb
+	ruby examples/rack/inject_extract.rb
 
 benchmark:
 	ruby benchmark/bench.rb

examples/rack/inject_extract.rb

@@ -0,0 +1,66 @@
+require 'bundler/setup'
+require 'lightstep'
+
+require 'rack'
+require 'rack/server'
+
+$token = '{your_access_token}'
+$request_id = 'abc123'
+
+class Router
+  def initialize
+    @tracer = LightStep::Tracer.new(component_name: 'router', access_token: $token)
+  end
+
+  def call(env)
+    span = @tracer.start_span("router_call").set_baggage_item("request-id", $request_id)
+    span.log(event: "router_request", env: env)
+    puts "parent #{span.span_context.trace_id}"
+
+    client = Net::HTTP.new("localhost", "9002")
+    req = Net::HTTP::Post.new("/")
+    @tracer.inject(span, LightStep::Tracer::FORMAT_RACK, req)
+    res = client.request(req)
+
+    span.log(event: "application_response", response: res.to_s)
+    span.finish
+    @tracer.flush
+    puts "----> https://app.lightstep.com/#{$token}/trace?span_guid=#{span.span_context.id}&at_micros=#{span.start_micros} <----"
+    [200, {}, [res.body]]
+  end
+end
+
+class App
+  def initialize
+    @tracer = LightStep::Tracer.new(component_name: 'app', access_token: $token)
+  end
+
+  def call(env)
+    span = @tracer.extract("app_call", LightStep::Tracer::FORMAT_RACK, env)
+    puts "child  #{span.to_h[:trace_guid]}"
+    span.log(event: "application", env: env)
+    sleep 0.05
+    span.finish
+    @tracer.flush
+    [200, {}, ["application"]]
+  end
+end
+
+router_thread = Thread.new do
+  Thread.abort_on_exception = true
+  Rack::Server.start(app: Router.new, Port: 9001)
+end
+
+app_thread = Thread.new do
+  Thread.abort_on_exception = true
+  Rack::Server.start(app: App.new, Port: 9002)
+end
+
+loop do
+  begin
+    p Net::HTTP.get(URI("http://localhost:9001/"))
+    break
+  rescue Errno::ECONNREFUSED
+    sleep 0.05
+  end
+end

lib/lightstep/tracer.rb

@@ -11,6 +11,7 @@ module LightStep
   class Tracer
     FORMAT_TEXT_MAP = 1
     FORMAT_BINARY = 2
+    FORMAT_RACK = 3
 
     class Error < LightStep::Error; end
     class ConfigurationError < LightStep::Tracer::Error; end
@@ -89,13 +90,15 @@ module LightStep
     # Inject a span into the given carrier
     # @param span [Span]
     # @param format [LightStep::Tracer::FORMAT_TEXT_MAP, LightStep::Tracer::FORMAT_BINARY]
-    # @param carrier [Hash-like]
+    # @param carrier [Hash]
     def inject(span, format, carrier)
       case format
       when LightStep::Tracer::FORMAT_TEXT_MAP
         inject_to_text_map(span, carrier)
       when LightStep::Tracer::FORMAT_BINARY
         warn 'Binary inject format not yet implemented'
+      when LightStep::Tracer::FORMAT_RACK
+        inject_to_rack(span, carrier)
       else
         warn 'Unknown inject format'
       end
@@ -104,7 +107,7 @@ module LightStep
     # Extract a span from a carrier
     # @param operation_name [String]
     # @param format [LightStep::Tracer::FORMAT_TEXT_MAP, LightStep::Tracer::FORMAT_BINARY]
-    # @param carrier [Hash-like]
+    # @param carrier [Hash]
     # @return [Span]
     def extract(operation_name, format, carrier)
       case format
@@ -113,6 +116,8 @@ module LightStep
       when LightStep::Tracer::FORMAT_BINARY
         warn 'Binary join format not yet implemented'
         nil
+      when LightStep::Tracer::FORMAT_RACK
+        extract_from_rack(operation_name, carrier)
       else
         warn 'Unknown join format'
         nil
@@ -213,5 +218,29 @@ module LightStep
 
       span
     end
+
+    def inject_to_rack(span, carrier)
+      carrier[CARRIER_TRACER_STATE_PREFIX + 'spanid'] = span.span_context.id
+      carrier[CARRIER_TRACER_STATE_PREFIX + 'traceid'] = span.span_context.trace_id unless span.span_context.trace_id.nil?
+      carrier[CARRIER_TRACER_STATE_PREFIX + 'sampled'] = 'true'
+
+      span.span_context.baggage.each do |key, value|
+        if key =~ /[^A-Za-z0-9\-_]/
+          # TODO: log the error internally
+          next
+        end
+        carrier[CARRIER_BAGGAGE_PREFIX + key] = value
+      end
+    end
+
+    def extract_from_rack(operation_name, env)
+      extract_from_text_map(operation_name, env.reduce({}){|memo, tuple|
+        raw_header, value = tuple
+        header = raw_header.gsub(/^HTTP_/, '').gsub("_", "-").downcase
+
+        memo[header] = value if header.start_with?(CARRIER_TRACER_STATE_PREFIX, CARRIER_BAGGAGE_PREFIX)
+        memo
+      })
+    end
   end
 end

spec/lightstep_spec.rb

@@ -284,6 +284,46 @@ describe LightStep do
     span2.finish
   end
 
+  it 'should handle inject/extract for http requests and rack' do
+    tracer = init_test_tracer
+    span1 = tracer.start_span('test_span')
+    span1.set_baggage_item('footwear', 'cleats')
+    span1.set_baggage_item('umbrella', 'golf')
+    span1.set_baggage_item('unsafe!@#$%$^&header', 'value')
+    span1.set_baggage_item('CASE-Sensitivity_Underscores', 'value')
+
+    carrier = {}
+    tracer.inject(span1, LightStep::Tracer::FORMAT_RACK, carrier)
+    expect(carrier['ot-tracer-traceid']).to eq(span1.span_context.trace_id)
+    expect(carrier['ot-tracer-spanid']).to eq(span1.span_context.id)
+    expect(carrier['ot-baggage-footwear']).to eq('cleats')
+    expect(carrier['ot-baggage-umbrella']).to eq('golf')
+    expect(carrier['ot-baggage-unsafeheader']).to be_nil
+    expect(carrier['ot-baggage-CASE-Sensitivity_Underscores']).to eq('value')
+
+    carrier = carrier.reduce({}) do |memo, tuple|
+      key, value = tuple
+      memo["HTTP_#{key.gsub("-", "_").upcase}"] = value
+      memo
+    end
+
+    span2 = tracer.extract('test_span_2', LightStep::Tracer::FORMAT_RACK, carrier)
+    expect(span2.span_context.trace_id).to eq(span1.span_context.trace_id)
+    expect(span2.tags[:parent_span_guid]).to eq(span1.span_context.id)
+    expect(span2.get_baggage_item('footwear')).to eq('cleats')
+    expect(span2.get_baggage_item('umbrella')).to eq('golf')
+    expect(span2.get_baggage_item('unsafe!@#$%$^&header')).to be_nil
+    expect(span2.get_baggage_item('unsafeheader')).to be_nil
+    expect(span2.get_baggage_item('case-sensitivity-underscores')).to eq('value')
+
+    span3 = tracer.extract('test_span_3', LightStep::Tracer::FORMAT_RACK, {'HTTP_OT_TRACER_TRACEID' => 'abc123'})
+    expect(span3.span_context.trace_id).to eq('abc123')
+
+    span1.finish
+    span2.finish
+    span3.finish
+  end
+
   it 'should handle concurrent spans' do
     result = nil
     tracer = init_callback_tracer(proc { |obj|; result = obj; })