Tracing.continue_from! with a block maintains active trace until the block ends

Author: marcotc

docs/GettingStarted.md

@@ -55,6 +55,12 @@ class TraceOperation
         :sampled,
         :service
 
+      # Creates a new TraceOperation.
+      #
+      # @param auto_finish [Boolean] when true, automatically finishes the trace when the local root span finishes.
+      #   When false, the trace remains unfinished until {#finish!} is called.
+      #   This is useful when this {TraceOperation} represents the continuation of a remote {TraceDigest},
+      #   in which case local root spans in this {TraceOperation} are children of the {TraceDigest}'s last active span.
       def initialize(
         logger: Datadog.logger,
         agent_sample_rate: nil,
@@ -80,7 +86,8 @@ def initialize(
         trace_state_unknown_fields: nil,
         remote_parent: false,
         tracer: nil, # DEV-3.0: deprecated, remove in 3.0
-        baggage: nil
+        baggage: nil,
+        auto_finish: true
       )
         @logger = logger
 
@@ -119,6 +126,7 @@ def initialize(
         @events = events || Events.new
         @finished = false
         @spans = []
+        @auto_finish = auto_finish
       end
 
       def full?
@@ -318,6 +326,29 @@ def flush!
         build_trace(spans, !finished)
       end
 
+      # When automatic context management is disabled (@auto_finish is false),
+      # this method finishes the trace, marking it as completed.
+      #
+      # The trace will **not** automatically finish when its local root span
+      # when @auto_finish is false, thus calling this method is mandatory
+      # in such scenario.
+      #
+      # Unfinished spans are discarded.
+      #
+      # This method is idempotent and safe to call after the trace is finished.
+      # It is also a no-op when @auto_finish is true, to prevent misuse.
+      #
+      # @!visibility private
+      def finish!
+        return if @auto_finish || finished?
+
+        @finished = true
+        @active_span = nil
+        @active_span_count = 0
+
+        events.trace_finished.publish(self)
+      end
+
       # Returns a set of trace headers used for continuing traces.
       # Used for propagation across execution contexts.
       # Data should reflect the active state of the trace.
@@ -349,7 +380,7 @@ def to_digest
           trace_state: @trace_state,
           trace_state_unknown_fields: @trace_state_unknown_fields,
           span_remote: @remote_parent && @active_span.nil?,
-          baggage: (@baggage.nil? || @baggage.empty?) ? nil : @baggage
+          baggage: @baggage.nil? || @baggage.empty? ? nil : @baggage
         ).freeze
       end
 
@@ -460,7 +491,7 @@ def activate_span!(span_op)
 
         @active_span = span_op
 
-        set_root_span!(span_op) unless root_span
+        set_local_root_span!(span_op)
       end
 
       def deactivate_span!(span_op)
@@ -483,15 +514,22 @@ def start_span(span_op)
         logger.debug { "Error starting span on trace: #{e} Backtrace: #{e.backtrace.first(3)}" }
       end
 
+      # For traces with automatic context management (auto_finish),
+      # when the local root span finishes, the trace also finishes.
+      # The trace cannot receive new spans after finished.
+      #
+      # Without auto_finish, the trace can still receive spans
+      # until explicitly finished.
       def finish_span(span, span_op, parent)
         # Save finished span & root span
         @spans << span unless span.nil?
 
         # Deactivate the span, re-activate parent.
         deactivate_span!(span_op)
 
-        # Set finished, to signal root span has completed.
-        @finished = true if span_op == root_span
+        # Finish if the local root span is finished and automatic
+        # context management is enabled.
+        @finished = true if span_op == root_span && @auto_finish
 
         # Update active span count
         @active_span_count -= 1
@@ -505,8 +543,8 @@ def finish_span(span, span_op, parent)
         logger.debug { "Error finishing span on trace: #{e} Backtrace: #{e.backtrace.first(3)}" }
       end
 
-      # Track the root span
-      def set_root_span!(span)
+      # Track the root {SpanOperation} object from the current execution context.
+      def set_local_root_span!(span)
         return if span.nil? || root_span
 
         @root_span = span
@@ -531,7 +569,7 @@ def build_trace(spans, partial = false)
           service: service,
           tags: meta,
           metrics: metrics,
-          root_span_id: (!partial) ? root_span&.id : nil,
+          root_span_id: !partial ? root_span&.id : nil,
           profiling_enabled: @profiling_enabled,
           apm_tracing_enabled: @apm_tracing_enabled
         )

lib/datadog/tracing/trace_operation.rb

@@ -149,10 +149,10 @@ def trace(
           context = call_context
           active_trace = context.active_trace
           trace = if continue_from || active_trace.nil?
-            start_trace(continue_from: continue_from)
-          else
-            active_trace
-          end
+                    start_trace(continue_from: continue_from)
+                  else
+                    active_trace
+                  end
         rescue => e
           logger.debug { "Failed to trace: #{e}" }
 
@@ -241,10 +241,21 @@ def active_correlation(key = nil)
         trace.to_correlation
       end
 
-      # Setup a new trace to continue from where another
+      # Setup a new trace execution context to continue from where another
       # trace left off.
+      # This is useful to continue distributed or async traces.
+      #
+      # The first span created in the restored context is a direct child of the
+      # active span from when the {Datadog::Tracing::TraceDigest} was created.
+      #
+      # When no block is given, the trace context is restored in the current thread.
+      # It remains active until the first span created in this restored context is finished.
+      # After that, if a new span is created, it start a new, unrelated trace.
       #
-      # Used to continue distributed or async traces.
+      # When a block is given, the trace context is restored inside the block execution.
+      # It remains active until the block ends, even when the first span created inside
+      # the block finishes. This means that multiple spans can be direct children of the
+      # active span from when the {Datadog::Tracing::TraceDigest} was created.
       #
       # @param [Datadog::Tracing::TraceDigest] digest continue from the {Datadog::Tracing::TraceDigest}.
       # @param [Thread] key Thread to retrieve trace from. Defaults to current thread. For internal use only.
@@ -260,13 +271,32 @@ def continue_trace!(digest, key = nil, &block)
         # Start a new trace from the digest
         context = call_context(key)
         original_trace = active_trace(key)
-        trace = start_trace(continue_from: digest)
+        # When we want the trace to be bound to a block, we cannot let
+        # it auto finish when the local root span finishes. This would
+        # create mutiple traces inside the block. Instead, we'll
+        # expliclity finish the trace after the block finishes.
+        auto_finish = !block
+
+        trace = start_trace(continue_from: digest, auto_finish: auto_finish)
 
         # If block hasn't been given; we need to manually deactivate
         # this trace. Subscribe to the trace finished event to do this.
         subscribe_trace_deactivation!(context, trace, original_trace) unless block
 
-        context.activate!(trace, &block)
+        if block
+          # When a block is given, the trace will be active until the block finishes.
+          context.activate!(trace) do
+            yield
+          ensure # We have to flush even when an error occurs
+            # On block completion, force the trace to finish and flush its finished spans.
+            # Unfinished spans are lost as the {TraceOperation} has ended.
+            trace.finish!
+            flush_trace(trace)
+          end
+        else
+          # Otherwise, the trace will be bound to the current thread after this point
+          context.activate!(trace)
+        end
       end
 
       # Sample a span, tagging the trace as appropriate.
@@ -329,15 +359,15 @@ def call_context(key = nil)
         @provider.context(key)
       end
 
-      def build_trace(digest = nil)
+      def build_trace(digest, auto_finish)
         # Resolve hostname if configured
         hostname = Core::Environment::Socket.hostname if Datadog.configuration.tracing.report_hostname
-        hostname = (hostname && !hostname.empty?) ? hostname : nil
+        hostname = hostname && !hostname.empty? ? hostname : nil
 
         if digest
           sampling_priority = if propagate_sampling_priority?(upstream_tags: digest.trace_distributed_tags)
-            digest.trace_sampling_priority
-          end
+                                digest.trace_sampling_priority
+                              end
           TraceOperation.new(
             logger: logger,
             hostname: hostname,
@@ -353,7 +383,8 @@ def build_trace(digest = nil)
             trace_state_unknown_fields: digest.trace_state_unknown_fields,
             remote_parent: digest.span_remote,
             tracer: self,
-            baggage: digest.baggage
+            baggage: digest.baggage,
+            auto_finish: auto_finish
           )
         else
           TraceOperation.new(
@@ -362,13 +393,13 @@ def build_trace(digest = nil)
             profiling_enabled: profiling_enabled,
             apm_tracing_enabled: apm_tracing_enabled,
             remote_parent: false,
-            tracer: self
+            tracer: self,
+            auto_finish: auto_finish
           )
         end
       end
       # rubocop:enable Metrics/MethodLength
 
-      # rubocop:disable Metrics/MethodLength
       def bind_trace_events!(trace_op)
         events = trace_op.send(:events)
 
@@ -387,13 +418,12 @@ def bind_trace_events!(trace_op)
           flush_trace(event_trace_op)
         end
       end
-      # rubocop:enable Metrics/MethodLength
 
       # Creates a new TraceOperation, with events bounds to this Tracer instance.
       # @return [TraceOperation]
-      def start_trace(continue_from: nil)
+      def start_trace(continue_from: nil, auto_finish: true)
         # Build a new trace using digest if provided.
-        trace = build_trace(continue_from)
+        trace = build_trace(continue_from, auto_finish)
 
         # Bind trace events: sample trace, set default service, flush spans.
         bind_trace_events!(trace)
@@ -402,7 +432,6 @@ def start_trace(continue_from: nil)
       end
 
       # rubocop:disable Lint/UnderscorePrefixedVariableName
-      # rubocop:disable Metrics/MethodLength
       def start_span(
         name,
         continue_from: nil,
@@ -454,18 +483,17 @@ def start_span(
           span
         end
       end
-      # rubocop:enable Lint/UnderscorePrefixedVariableName
-      # rubocop:enable Metrics/MethodLength
 
+      # rubocop:enable Lint/UnderscorePrefixedVariableName
       def resolve_tags(tags, service)
         merged_tags = if @tags.any? && tags
-          # Combine default tags with provided tags,
-          # preferring provided tags.
-          @tags.merge(tags)
-        else
-          # Use provided tags or default tags if none.
-          tags || @tags.dup
-        end
+                        # Combine default tags with provided tags,
+                        # preferring provided tags.
+                        @tags.merge(tags)
+                      else
+                        # Use provided tags or default tags if none.
+                        tags || @tags.dup
+                      end
         # Remove version tag if service is not the default service
         if merged_tags.key?(Core::Environment::Ext::TAG_VERSION) && service && service != @default_service
           merged_tags.delete(Core::Environment::Ext::TAG_VERSION)

lib/datadog/tracing/tracer.rb

@@ -4,9 +4,9 @@ module Datadog
       include Metadata::Tagging
 
       DEFAULT_MAX_LENGTH: ::Integer
-      
+
       @logger: Core::Logger
-      
+
       attr_reader logger: Core::Logger
 
       attr_accessor agent_sample_rate: untyped
@@ -17,17 +17,23 @@ module Datadog
       attr_accessor sample_rate: untyped
       attr_accessor remote_parent: untyped
       attr_accessor sampling_priority: untyped
+      attr_accessor baggage: untyped
       attr_reader active_span_count: untyped
       attr_reader active_span: untyped
       attr_reader id: untyped
       attr_reader max_length: untyped
       attr_reader parent_span_id: untyped
+      attr_reader trace_state: untyped
+      attr_reader trace_state_unknown_fields: untyped
       attr_writer name: untyped
       attr_writer resource: untyped
       attr_writer sampled: untyped
       attr_writer service: untyped
 
-      def initialize: (?agent_sample_rate: untyped?, ?events: untyped?, ?hostname: untyped?, ?id: untyped?, ?max_length: untyped, ?name: untyped?, ?origin: untyped?, ?parent_span_id: untyped?, ?rate_limiter_rate: untyped?, ?resource: untyped?, ?rule_sample_rate: untyped?, ?sample_rate: untyped?, ?sampled: untyped?, ?sampling_priority: untyped?, ?service: untyped?, ?profiling_enabled: untyped?, ?apm_tracing_enabled: untyped?, ?tags: untyped?, ?metrics: untyped?, ?remote_parent: untyped?) -> void
+      def initialize: (?logger: untyped, ?agent_sample_rate: untyped?, ?events: untyped?, ?hostname: untyped?, ?id: untyped?, ?max_length: untyped, ?name: untyped?, ?origin: untyped?, ?parent_span_id: untyped?, ?rate_limiter_rate: untyped?, ?resource: untyped?, ?rule_sample_rate: untyped?, ?sample_rate: untyped?, ?sampled: untyped?, ?sampling_priority: untyped?, ?service: untyped?, ?profiling_enabled: untyped?, ?apm_tracing_enabled: untyped?, ?tags: untyped?, ?metrics: untyped?, ?trace_state: untyped?, ?trace_state_unknown_fields: untyped?, ?remote_parent: untyped?, ?tracer: untyped?, ?baggage: untyped?, ?auto_finish: bool) -> void
+
+      def finish!: -> void
+
       def full?: () -> untyped
       def finished_span_count: () -> untyped
       def finished?: () -> untyped