bitfab 0.51.2 → 0.51.3

data/lib/bitfab/subtree.rb

@@ -0,0 +1,774 @@
+# frozen_string_literal: true
+
+require "rbconfig"
+require "securerandom"
+require "time"
+
+module Bitfab
+  # Captures Ruby method calls beneath one traced root via TracePoint.
+  module Subtree
+    DEFAULT_MAX_DEPTH = 30
+    DEFAULT_MAX_SPANS = 500
+    SESSION_KEY = :bitfab_subtree_session
+    DISPATCHER_KEY = :bitfab_subtree_dispatcher
+    FIBER_CONTEXT_KEY = :bitfab_subtree_fiber_context
+    REALPATH_CACHE = {}
+    REALPATH_MUTEX = Mutex.new
+    SDK_PREFIX = "#{File.realpath(File.expand_path("..", __dir__))}#{File::SEPARATOR}".freeze
+
+    CaptureConfig = Struct.new(:root, :root_path, :exclude, keyword_init: true)
+    FiberContext = Struct.new(:session, :fiber, :thread, :span_stack, :inherited_span, :parent, keyword_init: true)
+
+    Boundary = Struct.new(
+      :receiver,
+      :method_name,
+      :span_id,
+      :counts_toward_depth,
+      :skip_next_call,
+      keyword_init: true
+    )
+
+    MethodMetadata = Struct.new(
+      :identity,
+      :lineno,
+      :name,
+      :excluded,
+      :passthrough,
+      keyword_init: true
+    )
+
+    OpenSpan = Struct.new(
+      :identity,
+      :span_id,
+      :parent_span_id,
+      :name,
+      :args,
+      :kwargs,
+      :started_at,
+      :error,
+      keyword_init: true
+    )
+
+    module_function
+
+    def prepare(root_source:, exclude:)
+      path = root_source[0]
+      root = first_party_root(path)
+      CaptureConfig.new(
+        root: root && with_separator(realpath(root)),
+        root_path: path,
+        exclude: exclude.to_h { |name| [name.to_s, true] }.freeze
+      ).freeze
+    rescue
+      CaptureConfig.new(root: nil, root_path: nil, exclude: {}.freeze).freeze
+    end
+
+    def capture(client:, trace_function_key:, config:, max_depth:, max_spans:, include_wrappers:,
+      root_receiver:, root_method_name:)
+      return yield unless config.root
+
+      span = SpanContext.current || ReplayContext.current_span
+      return yield unless span
+
+      session = begin
+        replay = ReplayContext.current
+        Session.new(
+          client:,
+          trace_function_key:,
+          trace_id: span[:trace_id],
+          root_span_id: span[:span_id],
+          test_run_id: replay&.dig(:test_run_id),
+          input_source_span_id: replay&.dig(:input_source_span_id),
+          config:,
+          max_depth:,
+          max_spans:,
+          include_wrappers:,
+          root_receiver:,
+          root_method_name:
+        )
+      rescue Exception # rubocop:disable Lint/RescueException
+        return yield
+      end
+      result = session.run(defer_enumerator: true) { yield }
+      result.is_a?(Enumerator) ? session.wrap_enumerator(result) : result
+    end
+
+    def current_session
+      Thread.current[SESSION_KEY] || fiber_context&.session
+    end
+
+    def fiber_context
+      context = Fiber[FIBER_CONTEXT_KEY]
+      context = context.parent while context && !context.session&.active?
+      context if context && context.thread.equal?(Thread.current)
+    end
+
+    def with_borrowed_span(trace_id:, span_id:, receiver:, method_name:)
+      sessions = active_sessions.select { |session| session.trace_id == trace_id }
+      return yield if sessions.empty?
+
+      with_boundaries(
+        sessions,
+        receiver:,
+        method_name:,
+        span_id:,
+        counts_toward_depth: true
+      ) { yield }
+    end
+
+    def with_suppressed_node(receiver:, method_name:, sessions: active_sessions)
+      return yield if sessions.empty?
+
+      with_boundaries(
+        sessions,
+        receiver:,
+        method_name:,
+        span_id: nil,
+        counts_toward_depth: false
+      ) { yield }
+    end
+
+    def active_sessions
+      sessions = []
+      session = current_session
+      while session
+        sessions << session
+        session = session.parent_session
+      end
+      sessions
+    end
+
+    def current_parent_span_id(trace_id)
+      session = active_sessions.find { |candidate| candidate.trace_id == trace_id }
+      session&.current_parent_span_id
+    end
+
+    def with_boundaries(sessions, receiver:, method_name:, span_id:, counts_toward_depth:)
+      boundaries = sessions.map do |session|
+        [session, session.push_boundary(
+          receiver:,
+          method_name:,
+          span_id:,
+          counts_toward_depth:
+        )]
+      end
+      yield
+    ensure
+      boundaries&.reverse_each { |session, boundary| session.pop_boundary(boundary) }
+    end
+    private_class_method :with_boundaries
+
+    def active_trace?(trace_function_key)
+      session = current_session
+      while session
+        return true if session.trace_function_key == trace_function_key
+
+        session = session.parent_session
+      end
+      false
+    end
+
+    def dispatcher
+      Thread.current.thread_variable_get(DISPATCHER_KEY) || begin
+        value = Dispatcher.new
+        Thread.current.thread_variable_set(DISPATCHER_KEY, value)
+        value
+      end
+    end
+
+    def first_party_root(path)
+      return nil unless path && !path.start_with?("<")
+
+      resolved = realpath(path)
+      return nil unless resolved
+
+      if defined?(Gem)
+        gem_root = Gem.loaded_specs.values.map(&:full_gem_path).filter_map { |entry| realpath(entry) }
+          .select { |entry| resolved.start_with?(with_separator(entry)) }.max_by(&:length)
+        return gem_root if gem_root
+
+        Gem.path.each do |home|
+          prefix = with_separator(realpath(File.join(home, "gems")))
+          next unless prefix && resolved.start_with?(prefix)
+
+          return File.join(prefix, resolved.delete_prefix(prefix).split(File::SEPARATOR).first)
+        end
+      end
+
+      directory = File.dirname(resolved)
+      candidate = directory
+      loop do
+        if File.file?(File.join(candidate, "Gemfile")) || !Dir.glob(File.join(candidate, "*.gemspec")).empty?
+          return candidate
+        end
+
+        parent = File.dirname(candidate)
+        break if parent == candidate
+
+        candidate = parent
+      end
+      directory
+    rescue SystemCallError
+      nil
+    end
+
+    def realpath(path)
+      REALPATH_CACHE.fetch(path) do
+        REALPATH_MUTEX.synchronize do
+          REALPATH_CACHE.fetch(path) { REALPATH_CACHE[path] = File.realpath(path) }
+        end
+      end
+    rescue SystemCallError, TypeError
+      nil
+    end
+
+    def with_separator(path)
+      return nil unless path
+
+      path.end_with?(File::SEPARATOR) ? path : "#{path}#{File::SEPARATOR}"
+    end
+
+    class Dispatcher
+      def initialize
+        @sessions = []
+        @frames = []
+        @current_sessions = []
+        @trace = TracePoint.new(:call, :return, :raise, :rescue, :fiber_switch) do |event|
+          if event.event == :fiber_switch
+            switch_fiber
+          else
+            @current_sessions.each { |session| session.handle(event) }
+          end
+        end
+      end
+
+      def run(session)
+        previous_frame = Fiber[FIBER_CONTEXT_KEY]
+        # Fiber storage inherits the frame by reference; stacks remain local.
+        # Expired frames retain only their enclosing context, so children can
+        # keep participating in an outer capture without reviving an inner one.
+        frame = FiberContext.new(session:, fiber: Fiber.current, thread: Thread.current,
+          span_stack: SpanContext.stack, inherited_span: SpanContext.current, parent: Subtree.fiber_context)
+        Fiber[FIBER_CONTEXT_KEY] = frame
+        @frames.concat(session.deferred_frames)
+        session.deferred_frames.clear
+        @frames << frame
+        @sessions << session
+        @current_sessions = Subtree.active_sessions.reverse
+        @trace.enable(target_thread: Thread.current) unless @trace.enabled?
+        yield
+      ensure
+        @sessions.delete(session)
+        @sessions.each { |entry| entry.remove_parent(session) }
+        @frames.delete_if do |entry|
+          next false unless entry.session.equal?(session)
+
+          if session.deferred?
+            session.deferred_frames << entry
+          else
+            entry.session = entry.span_stack = entry.inherited_span = nil
+          end
+          true
+        end
+        Fiber[FIBER_CONTEXT_KEY] = previous_frame
+        @current_sessions = @current_sessions.reject { |entry| entry.equal?(session) }
+        @trace.disable if @sessions.empty?
+      end
+
+      private
+
+      def switch_fiber
+        context = Subtree.fiber_context
+        if context && !context.fiber.equal?(Fiber.current)
+          parent_fiber_id = context.fiber.object_id
+          inherited_span = context.span_stack.reverse_each.find { |span| span[:trace_id] == context.session.trace_id } ||
+            context.inherited_span
+          context = FiberContext.new(session: context.session, fiber: Fiber.current,
+            thread: Thread.current, span_stack: SpanContext.stack, inherited_span:, parent: context.parent)
+          Fiber[FIBER_CONTEXT_KEY] = context
+          @frames << context
+          Subtree.active_sessions.each { |session| session.inherit_parent(parent_fiber_id) }
+        end
+        @current_sessions = Subtree.active_sessions.reverse
+      end
+    end
+
+    class Session
+      attr_reader :client, :parent_session, :trace_function_key, :trace_id, :deferred_frames
+
+      def active? = @active
+      def deferred? = @deferred
+
+      def remove_parent(session)
+        @parent_session = session.parent_session if @parent_session.equal?(session)
+      end
+
+      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:)
+        @client = client
+        @trace_function_key = trace_function_key
+        @trace_id = trace_id
+        @root_span_id = root_span_id
+        @test_run_id = test_run_id
+        @input_source_span_id = input_source_span_id
+        @root = config.root
+        @root_path = config.root_path
+        @root_receiver = root_receiver
+        @root_method_name = root_method_name.to_sym
+        @skip_root_call = true
+        @max_depth = max_depth
+        @max_spans = max_spans
+        @exclude = config.exclude
+        @include_wrappers = include_wrappers
+        @stacks = {}
+        @fiber_parents = {}
+        @suppressed_depths = {}
+        @owned_paths = {}
+        @method_metadata = Hash.new { |owners, owner| owners[owner] = {} }
+        @spans_used = 0
+        @truncated = false
+        @exhausted = false
+        @drained = false
+        @deferred_frames = []
+      end
+
+      def run(defer_enumerator: false)
+        previous = Thread.current[SESSION_KEY]
+        @parent_session = Subtree.current_session
+        @active = true
+        @deferred = false
+        Thread.current[SESSION_KEY] = self
+        entered = false
+        completed = false
+        result = Subtree.dispatcher.run(self) do
+          entered = true
+          value = yield
+          @deferred = defer_enumerator && value.is_a?(Enumerator)
+          completed = true
+          value
+        end
+        result
+      rescue Exception # rubocop:disable Lint/RescueException
+        return result if completed
+        raise if entered
+
+        yield
+      ensure
+        Thread.current[SESSION_KEY] = previous
+        @active = false
+        @parent_session = nil
+        unless @deferred
+          @fiber_parents.clear
+          finish_open_spans
+          warn_if_truncated
+        end
+      end
+
+      def wrap_enumerator(source)
+        Enumerator.new do |yielder|
+          run do
+            source.each { |value| yielder << value }
+          end
+        end
+      end
+
+      def handle(event)
+        return if @drained
+
+        if event.event == :raise
+          record_error(event)
+          return
+        elsif event.event == :rescue
+          clear_rescued_error(event)
+          return
+        end
+
+        return if event.event == :call && skip_boundary_call?(event)
+
+        resolved_path = owned_path(event.path)
+        return unless resolved_path
+
+        case event.event
+        when :call
+          open_span(event, resolved_path)
+        when :return
+          close(event, resolved_path)
+        end
+      rescue Exception # rubocop:disable Lint/RescueException
+        nil
+      end
+
+      def excluded?(method_name)
+        @exclude[method_name.to_s] == true
+      end
+
+      def reserve_node
+        stack = current_stack(Fiber.current.object_id)
+        if captured_depth(stack) >= @max_depth
+          @truncated = true
+          return false
+        end
+        if @spans_used >= @max_spans
+          @truncated = true
+          @exhausted = true
+          return false
+        end
+
+        @spans_used += 1
+        true
+      end
+
+      def current_parent_span_id
+        stack = @stacks[Fiber.current.object_id] || []
+        stack.reverse_each.find { |entry| entry.span_id }&.span_id ||
+          @fiber_parents.dig(Fiber.current.object_id, :span_id) || @root_span_id
+      end
+
+      def inherit_parent(parent_fiber_id)
+        fiber_id = Fiber.current.object_id
+        return if @fiber_parents.key?(fiber_id)
+
+        stack = @stacks[parent_fiber_id] || []
+        parent = @fiber_parents[parent_fiber_id]
+        @fiber_parents[fiber_id] = {
+          span_id: stack.reverse_each.find { |entry| entry.span_id }&.span_id || parent&.dig(:span_id) || @root_span_id,
+          depth: local_depth(stack) + parent&.dig(:depth).to_i
+        }
+      end
+
+      def push_boundary(receiver:, method_name:, span_id:, counts_toward_depth:)
+        boundary = Boundary.new(
+          receiver:,
+          method_name: method_name.to_sym,
+          span_id:,
+          counts_toward_depth:,
+          skip_next_call: true
+        )
+        current_stack(Fiber.current.object_id) << boundary
+        boundary
+      end
+
+      def pop_boundary(boundary)
+        fiber_id = Fiber.current.object_id
+        stack = @stacks[fiber_id]
+        return unless stack
+
+        if stack.last.equal?(boundary)
+          stack.pop
+        else
+          stack.delete(boundary)
+        end
+        @stacks.delete(fiber_id) if stack.empty?
+      end
+
+      private
+
+      def finish_open_spans
+        @stacks.each_value do |stack|
+          stack.reverse_each do |entry|
+            next unless entry.is_a?(OpenSpan)
+
+            entry.error ||= "Subtree capture ended before the call returned"
+            emit(entry, result: nil)
+          end
+        end
+      ensure
+        @stacks.clear
+        @suppressed_depths.clear
+      end
+
+      def open_span(event, resolved_path)
+        fiber_id = Fiber.current.object_id
+        if @suppressed_depths[fiber_id].to_i.positive?
+          @suppressed_depths[fiber_id] += 1
+          return
+        end
+
+        if @exhausted
+          suppress(fiber_id)
+          return
+        end
+
+        metadata = metadata_for(event, resolved_path)
+        return if metadata.excluded
+        return if !@include_wrappers && metadata.passthrough
+
+        stack = current_stack(fiber_id)
+        if captured_depth(stack) >= @max_depth
+          @truncated = true
+          suppress(fiber_id)
+          return
+        end
+        if @spans_used >= @max_spans
+          @truncated = true
+          @exhausted = true
+          if captured_calls_drained?
+            @drained = true
+          else
+            suppress(fiber_id)
+          end
+          return
+        end
+
+        active_parent = SpanContext.current || ReplayContext.current_span
+        explicit_receiver = active_parent && active_parent[:explicit_span_receiver]
+        if explicit_receiver && active_parent[:trace_id] == @trace_id && explicit_receiver.equal?(event.self) &&
+            active_parent[:explicit_span_method_name] == called_method_id(event)
+          return
+        end
+
+        @spans_used += 1
+        span_id = SecureRandom.uuid
+        args, kwargs = 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
+          active_parent[:span_id]
+        end
+        subtree_parent_span_id = stack.reverse_each.find { |entry| entry.span_id }&.span_id
+        stack << OpenSpan.new(
+          identity: metadata.identity,
+          span_id:,
+          parent_span_id: subtree_parent_span_id || @fiber_parents.dig(fiber_id, :span_id) || active_parent_span_id || @root_span_id,
+          name: metadata.name,
+          args:,
+          kwargs:,
+          started_at: now
+        )
+      end
+
+      def close(event, resolved_path)
+        fiber_id = Fiber.current.object_id
+        suppressed_depth = @suppressed_depths[fiber_id].to_i
+        if suppressed_depth.positive?
+          if suppressed_depth == 1
+            @suppressed_depths.delete(fiber_id)
+            @drained = true if @exhausted && captured_calls_drained?
+          else
+            @suppressed_depths[fiber_id] = suppressed_depth - 1
+          end
+          return
+        end
+
+        stack = @stacks[fiber_id]
+        return unless stack
+
+        entry = stack.last
+        return unless entry.is_a?(OpenSpan) && matching_identity?(entry.identity, event, resolved_path)
+
+        stack.pop
+        emit(entry, result: event.return_value)
+        parent = stack.last
+        parent.error ||= entry.error if entry.error && parent.is_a?(OpenSpan)
+        @stacks.delete(fiber_id) if stack.empty?
+        @drained = true if @exhausted && captured_calls_drained?
+      end
+
+      def record_error(event)
+        entry = @stacks[Fiber.current.object_id]&.last
+        return unless entry.is_a?(OpenSpan)
+
+        exception = event.raised_exception
+        entry.error = "#{exception.class}: #{exception.message}"
+      end
+
+      def clear_rescued_error(event)
+        entry = @stacks[Fiber.current.object_id]&.last
+        return unless entry.is_a?(OpenSpan)
+
+        entry.error = nil
+      end
+
+      def emit(entry, result:)
+        # Recording a lazy return value must not execute the caller's iteration.
+        result = Serialize.unserializable_stub(result, "lazy_enumerator") if result.is_a?(Enumerator)
+        @client.send(
+          :send_span,
+          trace_function_key: @trace_function_key,
+          trace_id: @trace_id,
+          span_id: entry.span_id,
+          parent_span_id: entry.parent_span_id,
+          span_name: entry.name,
+          span_type: "function",
+          function_name: entry.name,
+          contexts: nil,
+          prompt: nil,
+          args: entry.args,
+          kwargs: entry.kwargs,
+          result:,
+          error: entry.error,
+          started_at: entry.started_at,
+          ended_at: now,
+          test_run_id: @test_run_id,
+          input_source_span_id: @input_source_span_id
+        )
+      rescue Exception # rubocop:disable Lint/RescueException
+        nil
+      end
+
+      def inputs(event)
+        args = []
+        kwargs = {}
+        event.parameters.each do |kind, name|
+          next unless name
+
+          value = event.binding.local_variable_get(name)
+          case kind
+          when :req, :opt
+            args << value
+          when :rest
+            args.concat(value)
+          when :keyreq, :key
+            kwargs[name] = value
+          when :keyrest
+            kwargs.merge!(value)
+          end
+        end
+        [args, kwargs]
+      rescue NameError, TypeError
+        [[], {}]
+      end
+
+      def span_name(event, method_id)
+        if @root_receiver.equal?(event.self) && @root_method_name == method_id
+          return method_id.to_s
+        end
+
+        owner = event.defined_class
+        owner_name = if owner.respond_to?(:singleton_class?) && owner.singleton_class?
+          event.self.respond_to?(:name) ? event.self.name : nil
+        elsif owner.respond_to?(:name)
+          owner.name
+        end
+        return method_id.to_s if !owner_name || owner_name.empty? || owner_name == "Object"
+
+        "#{owner_name}.#{method_id}"
+      end
+
+      def skip_boundary_call?(event)
+        boundary = @stacks[Fiber.current.object_id]&.last
+        if boundary.is_a?(Boundary) && boundary.skip_next_call &&
+            boundary.receiver.equal?(event.self) && boundary.method_name == called_method_id(event)
+          boundary.skip_next_call = false
+          return true
+        end
+
+        return false unless @skip_root_call
+        return false unless @root_receiver.equal?(event.self) && @root_method_name == called_method_id(event)
+
+        @skip_root_call = false
+        if event.path != @root_path
+          root = Subtree.first_party_root(event.path)
+          resolved_root = root && Subtree.with_separator(Subtree.realpath(root))
+          if resolved_root
+            @root = resolved_root
+            @root_path = event.path
+            @owned_paths.clear
+          end
+        end
+        true
+      end
+
+      def passthrough_wrapper?(parameters)
+        kinds = parameters.map(&:first)
+        kinds.any? { |kind| kind == :rest || kind == :keyrest } &&
+          kinds.all? { |kind| %i[rest keyrest block].include?(kind) }
+      end
+
+      def metadata_for(event, resolved_path)
+        method_id = called_method_id(event)
+        definitions = @method_metadata[event.defined_class]
+        cached = definitions[method_id]
+        return cached if cached && cached.identity[0] == resolved_path && cached.lineno == event.lineno
+
+        definitions[method_id] = MethodMetadata.new(
+          identity: [resolved_path, method_id, event.defined_class].freeze,
+          lineno: event.lineno,
+          name: span_name(event, method_id),
+          excluded: @exclude[method_id.to_s],
+          passthrough: passthrough_wrapper?(event.parameters)
+        ).freeze
+      end
+
+      def matching_identity?(identity, event, resolved_path)
+        identity[0] == resolved_path && identity[1] == called_method_id(event) &&
+          identity[2] == event.defined_class
+      end
+
+      def called_method_id(event)
+        event.callee_id || event.method_id
+      end
+
+      def owned_path(path)
+        @owned_paths.fetch(path) do
+          resolved = Subtree.realpath(path)
+          # The selected application may itself be installed under Gem.path.
+          # Its own boundary wins over ancestor exclusions, never the SDK's.
+          @owned_paths[path] = if resolved&.start_with?(@root) &&
+              !resolved.start_with?(SDK_PREFIX) &&
+              !Subtree.library_prefixes.any? { |prefix| resolved.start_with?(prefix) && !@root.start_with?(prefix) }
+            resolved
+          else
+            false
+          end
+        end
+      end
+
+      def current_stack(fiber_id)
+        @stacks[fiber_id] ||= []
+      end
+
+      def captured_depth(stack)
+        local_depth(stack) + @fiber_parents.dig(Fiber.current.object_id, :depth).to_i
+      end
+
+      def local_depth(stack)
+        stack.count do |entry|
+          !entry.is_a?(Boundary) || entry.counts_toward_depth
+        end
+      end
+
+      def captured_calls_drained?
+        @suppressed_depths.empty? && @stacks.values.all? do |stack|
+          stack.none? { |entry| entry.is_a?(OpenSpan) }
+        end
+      end
+
+      def suppress(fiber_id)
+        @suppressed_depths[fiber_id] = 1
+      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: [...]."
+        )
+      end
+
+      def now
+        Bitfab.now_iso_timestamp
+      end
+    end
+
+    def library_prefixes
+      @library_prefixes ||= begin
+        paths = [
+          RbConfig::CONFIG["rubylibdir"],
+          RbConfig::CONFIG["archdir"],
+          RbConfig::CONFIG["vendorlibdir"],
+          RbConfig::CONFIG["sitelibdir"],
+          File.expand_path("..", __dir__)
+        ]
+        paths.concat(Gem.path) if defined?(Gem)
+        paths.compact.filter_map do |path|
+          with_separator(realpath(path))
+        end.freeze
+      end
+    end
+  end
+end