bitfab 0.51.9 → 0.57.0

data/lib/bitfab/subtree.rb

@@ -6,9 +6,12 @@ require "time"
 
 module Bitfab
   # Captures Ruby method calls beneath one traced root via TracePoint.
+  class MixedTracingError < StandardError; end
+
   module Subtree
     DEFAULT_MAX_DEPTH = 30
-    DEFAULT_MAX_SPANS = 500
+    DEFAULT_MAX_SPANS = 2048
+    DEFAULT_MAX_CAPTURED_SUBTREE_SPANS = 512
     SESSION_KEY = :bitfab_subtree_session
     DISPATCHER_KEY = :bitfab_subtree_dispatcher
     FIBER_CONTEXT_KEY = :bitfab_subtree_fiber_context
@@ -46,11 +49,119 @@ module Bitfab
       :kwargs,
       :started_at,
       :error,
+      :span_type,
+      :test_run_id,
+      :span_data,
       keyword_init: true
     )
 
     module_function
 
+    def assert_surface!(surface)
+      enclosing = current_session ? :trace : SpanContext.current&.dig(:surface)
+      return unless enclosing && enclosing != surface
+
+      raise MixedTracingError,
+        "opt-in and opt-out tracing cannot be mixed: #{surface} entered inside #{enclosing}. " \
+        "Use bitfab_node inside bitfab_trace, or use bitfab_span throughout the workflow."
+    rescue MixedTracingError
+      raise
+    rescue
+      nil
+    end
+
+    def nested_root_suppressed?
+      current_session && (ReplayContext.current || (defined?(SeedContext) && SeedContext.current))
+    end
+
+    def qualified_name(receiver, method_name)
+      owner = receiver.is_a?(Module) ? receiver.name : receiver.class.name
+      (owner && !owner.empty? && owner != "Object") ? "#{owner}.#{method_name}" : method_name.to_s
+    end
+
+    class ConfiguredCall
+      attr_reader :span_data, :owner
+
+      def initialize(receiver:, method_name:, name:, type:, args:, kwargs:, test_run_id: nil,
+        nested_trace_key: nil)
+        @receiver = receiver
+        @method_name = method_name
+        @nested_trace_key = nested_trace_key
+        @span_data = {}
+        @completed = false
+        @records = Subtree.active_sessions.map do |session|
+          omitted = session.excluded?(method_name)
+          entry = unless omitted
+            session.configured_entry(name:, type:, args:, kwargs:, test_run_id:)
+          end
+          [session, entry, !omitted]
+        end
+        @owner = @records.find { |_, entry, _| entry }&.first unless nested_trace_key
+      end
+
+      def recorded_elsewhere?
+        @records.any? { |session, entry, _| entry && !session.equal?(@owner) }
+      end
+
+      def start(trace_id:, span_id:)
+        @records.each do |session, entry, _|
+          next unless entry
+
+          if session.equal?(@owner)
+            entry.span_id = span_id
+            session.discard_configured(entry)
+          elsif @nested_trace_key
+            entry.span_data = {
+              "nested_trace_id" => trace_id,
+              "nested_trace_function_key" => @nested_trace_key,
+              "nested_root_span_id" => span_id
+            }
+          end
+        end
+        if @nested_trace_key && (enclosing = @records.find { |_, entry, _| entry })
+          session, entry = enclosing
+          @span_data.merge!(
+            "enclosing_trace_id" => session.trace_id,
+            "enclosing_trace_function_key" => session.trace_function_key,
+            "enclosing_span_id" => entry.span_id
+          )
+        end
+      end
+
+      def run
+        boundaries = @records.filter_map do |session, entry, counts_toward_depth|
+          next if session.equal?(@owner)
+
+          [session, session.push_boundary(
+            receiver: @receiver, method_name: @method_name,
+            span_id: entry&.span_id, counts_toward_depth:
+          )]
+        end
+        yield
+      ensure
+        boundaries&.reverse_each { |session, boundary| session.pop_boundary(boundary) }
+      end
+
+      def with_owner_parent(&body)
+        return yield unless @owner
+
+        SpanContext.with_span(trace_id: @owner.trace_id,
+          span_id: @owner.current_parent_span_id, surface: :trace, &body)
+      end
+
+      def complete(result:, error:, ended_at:)
+        return if @completed
+
+        @completed = true
+        @records.each do |session, entry, _|
+          next if !entry || session.equal?(@owner)
+
+          entry.error = error
+          session.emit_configured(entry, result:)
+        end
+      end
+    end
+
     def prepare(root_source:, exclude:)
       path = root_source[0]
       root = first_party_root(path)
@@ -64,7 +175,7 @@ module Bitfab
     end
 
     def capture(client:, trace_function_key:, config:, max_depth:, max_spans:, include_wrappers:,
-      root_receiver:, root_method_name:)
+      root_receiver:, root_method_name:, max_captured_subtree_spans: DEFAULT_MAX_CAPTURED_SUBTREE_SPANS, mock_on_replay_default: false)
       return yield unless config.root
 
       span = SpanContext.current || ReplayContext.current_span
@@ -82,7 +193,9 @@ module Bitfab
           config:,
           max_depth:,
           max_spans:,
+          max_captured_subtree_spans:,
           include_wrappers:,
+          mock_on_replay_default:,
           root_receiver:,
           root_method_name:
         )
@@ -294,7 +407,7 @@ module Bitfab
     end
 
     class Session
-      attr_reader :client, :parent_session, :trace_function_key, :trace_id, :deferred_frames
+      attr_reader :client, :parent_session, :trace_function_key, :trace_id, :deferred_frames, :mock_on_replay_default
 
       def active? = @active
       def deferred? = @deferred
@@ -305,7 +418,8 @@ module Bitfab
 
       def initialize(client:, trace_function_key:, trace_id:, root_span_id:, test_run_id:,
         input_source_span_id:, config:, max_depth:, max_spans:, include_wrappers:,
-        root_receiver:, root_method_name:)
+        root_receiver:, root_method_name:, max_captured_subtree_spans: DEFAULT_MAX_CAPTURED_SUBTREE_SPANS, mock_on_replay_default: false)
+        @mock_on_replay_default = mock_on_replay_default
         @client = client
         @trace_function_key = trace_function_key
         @trace_id = trace_id
@@ -319,14 +433,17 @@ module Bitfab
         @skip_root_call = true
         @max_depth = max_depth
         @max_spans = max_spans
+        @max_captured_subtree_spans = [max_captured_subtree_spans, max_spans].min
         @exclude = config.exclude
         @include_wrappers = include_wrappers
         @stacks = {}
+        @configured_entries = {}
         @fiber_parents = {}
         @suppressed_depths = {}
         @owned_paths = {}
         @method_metadata = Hash.new { |owners, owner| owners[owner] = {} }
-        @spans_used = 0
+        @spans_used = 1
+        @captured_subtree_count = 0
         @truncated = false
         @exhausted = false
         @drained = false
@@ -409,14 +526,34 @@ module Bitfab
           @truncated = true
           return false
         end
-        if @spans_used >= @max_spans
-          @truncated = true
-          @exhausted = true
-          return false
-        end
+        return true if take_slot(discovered: false)
 
-        @spans_used += 1
-        true
+        @truncated = true
+        @exhausted = true
+        false
+      end
+
+      def configured_entry(name:, type:, args:, kwargs:, test_run_id:)
+        return unless reserve_node
+
+        entry = OpenSpan.new(
+          span_id: SecureRandom.uuid,
+          parent_span_id: current_parent_span_id,
+          name:, span_type: type, args:, kwargs:, test_run_id:, started_at: now,
+          span_data: {}
+        )
+        @configured_entries[entry.object_id] = entry
+        entry
+      end
+
+      def discard_configured(entry)
+        @configured_entries.delete(entry.object_id)
+      end
+
+      def emit_configured(entry, result:)
+        return unless discard_configured(entry)
+
+        emit(entry, result:)
       end
 
       def current_parent_span_id
@@ -465,6 +602,11 @@ module Bitfab
       private
 
       def finish_open_spans
+        @configured_entries.each_value do |entry|
+          entry.error ||= "Subtree capture ended before the call returned"
+          emit(entry, result: nil)
+        end
+        @configured_entries.clear
         @stacks.each_value do |stack|
           stack.reverse_each do |entry|
             next unless entry.is_a?(OpenSpan)
@@ -500,14 +642,10 @@ module Bitfab
           suppress(fiber_id)
           return
         end
-        if @spans_used >= @max_spans
+        content_off = content_off?(metadata.name)
+        if (limit = limit_reached(discovered: true, content_off:))
           @truncated = true
-          @exhausted = true
-          if captured_calls_drained?
-            @drained = true
-          else
-            suppress(fiber_id)
-          end
+          stop_discovery(fiber_id) unless limit == :max_captured_subtree_spans && content_off_names?
           return
         end
 
@@ -518,9 +656,10 @@ module Bitfab
           return
         end
 
-        @spans_used += 1
+        count_span(discovered: true, content_off:)
         span_id = SecureRandom.uuid
-        args, kwargs = inputs(event)
+        without_content = content_off || plan_unavailable?
+        args, kwargs = without_content ? [[], {}] : inputs(event)
         # A nested subtree root temporarily owns the ambient span context, but
         # an outer session must never parent its spans across trace boundaries.
         active_parent_span_id = if active_parent&.dig(:trace_id) == @trace_id
@@ -534,7 +673,8 @@ module Bitfab
           name: metadata.name,
           args:,
           kwargs:,
-          started_at: now
+          started_at: now,
+          span_data: without_content ? {"content_off_by_simulation_plan" => true} : nil
         )
       end
 
@@ -590,7 +730,9 @@ module Bitfab
           span_id: entry.span_id,
           parent_span_id: entry.parent_span_id,
           span_name: entry.name,
-          span_type: "function",
+          span_type: entry.span_type || "function",
+          instrumentation: "trace",
+          span_data: entry.span_data || {},
           function_name: entry.name,
           contexts: nil,
           prompt: nil,
@@ -600,7 +742,7 @@ module Bitfab
           error: entry.error,
           started_at: entry.started_at,
           ended_at: now,
-          test_run_id: @test_run_id,
+          test_run_id: entry.test_run_id || @test_run_id,
           input_source_span_id: @input_source_span_id
         )
       rescue Exception # rubocop:disable Lint/RescueException
@@ -739,14 +881,60 @@ module Bitfab
         @suppressed_depths[fiber_id] = 1
       end
 
+      def limit_reached(discovered:, content_off: false)
+        return :max_spans if @spans_used >= @max_spans
+
+        :max_captured_subtree_spans if discovered && !content_off && @captured_subtree_count >= @max_captured_subtree_spans
+      end
+
+      def take_slot(discovered:)
+        return false if limit_reached(discovered:)
+
+        count_span(discovered:, content_off: false)
+        true
+      end
+
+      def count_span(discovered:, content_off:)
+        @spans_used += 1
+        @captured_subtree_count += 1 if discovered && !content_off
+      end
+
+      def stop_discovery(fiber_id)
+        @exhausted = true
+        if captured_calls_drained?
+          @drained = true
+        else
+          suppress(fiber_id)
+        end
+      end
+
+      def content_off?(name)
+        @client.send(:content_off_by_simulation_plan?, trace_id: @trace_id, trace_function_key: @trace_function_key,
+          span_name: name, instrumentation: "trace")
+      rescue
+        false
+      end
+
+      def content_off_names?
+        @client.send(:simulation_plan_content_off_names?, trace_id: @trace_id, trace_function_key: @trace_function_key)
+      rescue
+        false
+      end
+
+      def plan_unavailable?
+        @client.send(:simulation_plan_unavailable?, instrumentation: "trace")
+      rescue
+        false
+      end
+
       def warn_if_truncated
         return unless @truncated
 
         Bitfab.warn_once(
           "subtree-truncated:#{@trace_function_key}",
           "'#{@trace_function_key}' hit a subtree capture limit (max_depth=#{@max_depth}, " \
-          "max_spans=#{@max_spans}) and its trace is incomplete. Raise the limits, or narrow " \
-          "what is traced with exclude: [...]."
+          "max_captured_subtree_spans=#{@max_captured_subtree_spans}, max_spans=#{@max_spans}) and its trace is incomplete. " \
+          "Raise the limits, add bitfab_node to the calls that matter, or narrow what is traced with exclude: [...]."
         )
       end
 

data/lib/bitfab/thread_propagation.rb

@@ -0,0 +1,111 @@
+# frozen_string_literal: true
+
+module Bitfab
+  module ThreadPropagation
+    ORIGIN_KEY = :__bitfab_submit_origin
+    SDK_PATH = File.expand_path(__dir__) + File::SEPARATOR
+    @installed = false
+    @mutex = Mutex.new
+
+    module_function
+
+    def install
+      @mutex.synchronize do
+        Thread.singleton_class.prepend(ThreadMethods) unless Thread.singleton_class.ancestors.include?(ThreadMethods)
+        @installed = true
+      end
+      begin
+        require "concurrent-ruby"
+        Concurrent::ThreadPoolExecutor.prepend(PoolMethods) unless Concurrent::ThreadPoolExecutor.ancestors.include?(PoolMethods)
+      rescue LoadError
+        nil
+      end
+    end
+
+    def installed?
+      @installed
+    end
+
+    def uninstall
+      @installed = false
+    end
+
+    def runtime
+      current = Thread.current
+      data = {"thread_id" => current.object_id, "thread_name" => current.name || ((current == Thread.main) ? "main" : "thread-#{current.object_id}")}
+      if (origin = current.thread_variable_get(ORIGIN_KEY))
+        data["parent_thread_id"] = origin["thread_id"]
+        data["submit_thread_id"] = origin["thread_id"]
+        data["submit_thread_name"] = origin["thread_name"]
+      end
+      data
+    end
+
+    def bind(block)
+      stack = SpanContext.stack.dup
+      stack << SpanContext.current if stack.empty? && SpanContext.current
+      replay = ReplayContext.current
+      seed = SeedContext.current
+      return block if stack.empty? && replay.nil? && seed.nil?
+
+      origin = runtime.slice("thread_id", "thread_name")
+      lambda do |*args, **kwargs|
+        current = Thread.current
+        previous_stack = current[SpanContext::STACK_KEY]
+        previous_bridge = current.thread_variable_get(SpanContext::FIBER_BRIDGE_STACK_KEY)
+        previous_replay = current.thread_variable_get(REPLAY_CONTEXT_KEY)
+        previous_seed = current.thread_variable_get(SeedContext::KEY)
+        previous_origin = current.thread_variable_get(ORIGIN_KEY)
+        local_replay = replay&.merge(span_stack: replay.fetch(:span_stack, []).dup)
+        current[SpanContext::STACK_KEY] = stack.dup
+        current.thread_variable_set(SpanContext::FIBER_BRIDGE_STACK_KEY, [])
+        current.thread_variable_set(REPLAY_CONTEXT_KEY, local_replay)
+        current.thread_variable_set(SeedContext::KEY, seed)
+        current.thread_variable_set(ORIGIN_KEY, origin)
+        begin
+          block.call(*args, **kwargs)
+        ensure
+          replay[:db_snapshot_accessed] = true if replay && local_replay[:db_snapshot_accessed]
+          current[SpanContext::STACK_KEY] = previous_stack
+          current.thread_variable_set(SpanContext::FIBER_BRIDGE_STACK_KEY, previous_bridge)
+          current.thread_variable_set(REPLAY_CONTEXT_KEY, previous_replay)
+          current.thread_variable_set(SeedContext::KEY, previous_seed)
+          current.thread_variable_set(ORIGIN_KEY, previous_origin)
+        end
+      end
+    end
+
+    def internal?(block)
+      path = block.source_location&.first
+      path && (path.start_with?(SDK_PATH) || path.include?("/opentelemetry") || path.include?("/concurrent-ruby"))
+    end
+
+    module ThreadMethods
+      %i[new start fork].each do |method_name|
+        define_method(method_name) do |*args, &block|
+          return super(*args, &block) if !ThreadPropagation.installed? || block.nil? || ThreadPropagation.internal?(block)
+
+          bound = begin
+            ThreadPropagation.bind(block)
+          rescue
+            block
+          end
+          super(*args, &bound)
+        end
+      end
+    end
+
+    module PoolMethods
+      def post(*args, &block)
+        return super unless ThreadPropagation.installed? && block
+
+        bound = begin
+          ThreadPropagation.bind(block)
+        rescue
+          block
+        end
+        super(*args, &bound)
+      end
+    end
+  end
+end