trace_headers.rb in faraday-zipkin-0.3.0

- old
+ new

@@ -8,11 +8,12 @@
     class TraceHeaders < ::Faraday::Middleware
       B3_HEADERS = {
         :trace_id => "X-B3-TraceId",
         :parent_id => "X-B3-ParentSpanId",
         :span_id => "X-B3-SpanId",
-        :sampled => "X-B3-Sampled"
+        :sampled => "X-B3-Sampled",
+        :flags => "X-B3-Flags"
       }.freeze
 
       def initialize(app, service_name=nil)
         @app = app
         @service_name = service_name
@@ -22,11 +23,12 @@
         # handle either a URI object (passed by Faraday v0.8.x in testing), or something string-izable
         url = env[:url].respond_to?(:host) ? env[:url] : URI.parse(env[:url].to_s)
         service_name = @service_name || url.host.split('.').first # default to url-derived service name
         endpoint = ::Trace::Endpoint.new(host_ip_for(url.host), url.port, service_name)
 
-        trace_id = ::Trace.id
-        ::Trace.push(trace_id.next_id) do
+        ::Trace.unwind do
+          trace_id = ::Trace.id
+          ::Trace.push(trace_id.next_id)
           B3_HEADERS.each do |method, header|
             env[:request_headers][header] = ::Trace.id.send(method).to_s
           end
 
           # annotate with method (GET/POST/etc.) and uri path