lib/request_tracer/trace.rb in request-tracer-0.5.4 vs lib/request_tracer/trace.rb in request-tracer-0.5.5

@@ -46,10 +46,13 @@
       attr_reader :trace_id, :parent_id, :span_id
       def self.spawn_from_hash(h)
         span_id = Trace.generate_id
         self.new(h["trace_id"] || span_id, h["span_id"], Trace.generate_id)
       end
+      def self.create(h)
+        self.new(h["trace_id"], h["parent_span_id"], h["span_id"])
+      end
       def initialize(trace_id, parent_id, span_id)
         @trace_id = SpanId.from_value(trace_id)
         @parent_id = parent_id && SpanId.from_value(parent_id)
         @span_id = SpanId.from_value(span_id)
       end
@@ -71,35 +74,49 @@
       def to_json
         to_h.to_json
       end
     end
 
+    def create
+      span_id = generate_id
+      trace_id = TraceId.new(span_id, nil, span_id)
+      stack.push(trace_id)
+      trace_id
+    end
 
     def latest
-      if stack.empty?
-        span_id = generate_id
-        trace_id = TraceId.new(span_id, nil, span_id)
-        stack.push(trace_id)
-      end
       stack.last
     end
 
+    def latest_or_create
+      latest || create
+    end
+
     def push(trace_info)
-      stack.push(TraceId.spawn_from_hash(trace_info))
+      trace = if trace_info.include?("trace_id") && trace_info.include?("span_id")
+        TraceId.create(trace_info)
+      else
+        TraceId.spawn_from_hash(trace_info)
+      end
+      stack.push(trace)
       if block_given?
         begin
-          yield
+          yield trace
         ensure
           pop
         end
       end
     end
 
     def pop
       stack.pop
     end
 
+    def clear
+      stack.clear
+    end
+
     def unwind
       if block_given?
         begin
           saved_stack = stack.dup
           yield
@@ -107,18 +124,14 @@
           stack = saved_stack
         end
       end
     end
 
-    def record(annotation, &block)
-      tracer.record(latest, annotation, &block)
+    def record(annotation = nil, &block)
+      tracer.record(latest.next_id, annotation, &block)
     end
 
-    def set_rpc_name(name)
-      tracer.set_rpc_name(latest, name) unless stack.empty?
-    end
-
     def tracer=(tracer)
       @tracer = tracer
     end
 
     def generate_id
@@ -148,12 +161,12 @@
       block.call
     end
   end
 
   class DefaultTracer
-    def record(*args, &block)
-      Trace.push(args[1].to_h) do
-        block.call
+    def record(latest_trace, annotation, &block)
+      Trace.push(latest_trace.to_h) do |trace|
+        block.call(trace)
       end
     end
   end
 end