bitfab 0.33.4 → 0.33.6

data/lib/bitfab/otel_collector.rb

@@ -0,0 +1,28 @@
+# frozen_string_literal: true
+
+require "opentelemetry-exporter-otlp"
+
+module Bitfab
+  module Otel
+    # OTLP/HTTP protobuf delivery to a customer's Collector.
+    class CollectorExporter < OpenTelemetry::Exporter::OTLP::Exporter
+      def self.traces_endpoint(endpoint)
+        base = endpoint.chomp("/")
+        base.end_with?("/v1/traces") ? base : "#{base}/v1/traces"
+      end
+
+      def initialize(api_key:, **options)
+        super(**options)
+        @api_key = api_key
+      end
+
+      # The OTLP exporter reads headers once at construction, while a Bitfab key
+      # is resolved at request time (an env var may load after the client is built).
+      def export(span_data, timeout: nil)
+        key = @api_key.respond_to?(:call) ? @api_key.call : @api_key
+        @headers = @headers.merge("Authorization" => "Bearer #{key}")
+        super
+      end
+    end
+  end
+end

data/lib/bitfab/replay.rb

@@ -9,6 +9,7 @@ require_relative "constants"
 require_relative "mock_override"
 require_relative "serialize"
 require_relative "traceable"
+require_relative "transport"
 
 module Bitfab
   # Replay mock strategies. Mirrors the Python and TypeScript SDKs.
@@ -29,12 +30,8 @@ module Bitfab
 
     # Execute a block with replay context set on the current thread.
     # The context is automatically cleared when the block completes.
-    #
-    # pending_persistence, when given, collects the root span's persistence
-    # threads (span uploads + trace completion) so the replay runner can join
-    # them before complete_replay builds the trace-ID mapping.
     def with_context(test_run_id:, input_source_span_id: nil, input_source_trace_id: nil, trace_id: nil,
-      mock_tree: nil, mock_strategy: nil, mock_overrides: nil, fetch_span_output: nil, pending_persistence: nil,
+      mock_tree: nil, mock_strategy: nil, mock_overrides: nil, fetch_span_output: nil,
       db_branch_lease: nil, source_bitfab_trace_id: nil)
       previous = Thread.current[REPLAY_CONTEXT_KEY]
       ctx = {
@@ -43,7 +40,6 @@ module Bitfab
         input_source_trace_id:,
         trace_id:
       }
-      ctx[:pending_persistence] = pending_persistence if pending_persistence
       if mock_tree
         ctx[:mock_tree] = mock_tree
         ctx[:mock_strategy] = mock_strategy || "marked"
@@ -82,6 +78,8 @@ module Bitfab
   # Replay historical traces through a traced method and create a test run.
   module Replay
     CODE_CHANGE_UNSET = Object.new.freeze
+    PERSISTENCE_TIMEOUT_SECONDS = 30.0
+    PERSISTENCE_POLL_SECONDS = 0.1
 
     module_function
 
@@ -127,10 +125,10 @@ module Bitfab
     #   source_trace_id/source_span_id aliases), and returns [new_args, new_kwargs].
     #   Runs per item inside the same rescue as the method, so a raising adapter
     #   sets that item's :error rather than crashing the run.
-    # @param db_branch [Hash, nil] passing it at all turns database branching
-    #   on, so +{}+ enables it with the mirror project's own sizing. Keys
-    #   +:min_cu+/+:max_cu+ are the branch compute's autoscaling floor and
-    #   ceiling in Neon Compute Units, and +:warmup_sql+ warms the branch's
+    # @param db_branch [Boolean, Hash, nil] +true+ branches with the mirror
+    #   project's own sizing; a Hash tunes it and +false+/+nil+ leave branching
+    #   off. Keys +:min_cu+/+:max_cu+ are the branch compute's autoscaling floor
+    #   and ceiling in Neon Compute Units, and +:warmup_sql+ warms the branch's
     #   cache before the replayed method sees it. Read the resolved branch
     #   inside the method with +Bitfab.current_replay_branch+.
     # @param on_progress [#call, nil] optional callback invoked once per item as
@@ -192,7 +190,7 @@ module Bitfab
       # the count), so it's omitted from the request entirely.
       effective_limit = trace_ids ? nil : (limit || 5)
 
-      include_db_branch_lease = !db_branch.nil?
+      include_db_branch_lease = db_branch_enabled?(db_branch)
 
       # code_change_files controls capture: omitted auto-captures, an array
       # wins, and explicit nil opts out. Preserve a caller-supplied description
@@ -233,11 +231,15 @@ module Bitfab
         []
       end
 
-      # Every item joined its own trace-persistence threads (span uploads +
-      # completion) in execute_item, so all replay traces are on the server
-      # by now: no flush needed. complete_replay finalizes the run and returns
-      # the token/diagnostic mapping; its failures propagate loudly because a
-      # run that never completed can't be finalized.
+      # Spans and completions ride a batched transport, so the run is only safe
+      # to finalize once the server confirms every replay trace it queued: the
+      # trace-ID mapping complete_replay builds would otherwise race the
+      # in-flight batches and hand back nil for every item.
+      wait_for_replay_persistence(http_client, test_run_id, result_items.map { |item| item[:_sdk_trace_id] })
+
+      # complete_replay finalizes the run and returns the token/diagnostic
+      # mapping; its failures propagate loudly because a run that never
+      # completed can't be finalized.
       complete_response = http_client.complete_replay(test_run_id)
       trace_id_map = complete_response&.dig("traceIds")
       # Per-replay-trace token usage keyed by server trace id: the REPLAYED
@@ -277,7 +279,7 @@ module Bitfab
           raise "Replay completed but the server has no persisted trace for " \
             "any of the #{completed_count} completed item(s) " \
             "(test_run_id #{test_run_id}).#{server_count} Trace uploads were " \
-            "joined, so either the uploads failed or the replayed method is " \
+            "flushed, so either the uploads failed or the replayed method is " \
             "not traced (no root span was emitted)."
         end
         # SOME completed items missing: per-item upload failure. Warn, but
@@ -335,16 +337,26 @@ module Bitfab
     CC_MAX_FILE_BYTES = 500_000
     CC_MAX_TOTAL_BYTES = 2_000_000
 
+    # Whether the caller asked for database branching. +true+ and a Hash both
+    # turn it on; +false+ and +nil+ leave it off. Kept separate from
+    # +db_branch_settings+ because +db_branch: true+ enables branching while
+    # contributing no settings, so the presence of settings can't stand in for
+    # the switch.
+    def db_branch_enabled?(db_branch)
+      !db_branch.nil? && db_branch != false
+    end
+
     # Convert the caller's +db_branch+ options to the wire shape, or nil when
     # nothing was set. Dropping the key entirely keeps the request identical to
-    # what older SDKs send.
+    # what older SDKs send. Booleans carry no settings: they only move the
+    # switch.
     #
     # Keys: :min_cu and :max_cu are the branch compute's autoscaling floor and
     # ceiling in Neon Compute Units (equal values pin the size, keeping items
     # comparable); :warmup_sql is appended to the branch's readiness check so it
     # warms the cache before the replayed function sees the lease.
     def db_branch_settings(db_branch)
-      return nil if db_branch.nil? || db_branch.empty?
+      return nil if db_branch.nil? || db_branch == true || db_branch == false || db_branch.empty?
 
       # Symbol or string keys, matching normalize_code_change_files: a hash
       # assembled from JSON or YAML would otherwise be read as empty, leaving
@@ -866,6 +878,40 @@ module Bitfab
       }
     end
 
+    # Block until the server has persisted every replay trace this run queued.
+    #
+    # Flushing only proves the batches left the process. The server writes a
+    # trace's spans and completion independently, so the barrier also polls
+    # replay status with the span count each trace owes: a trace is ready only
+    # once its completion and all of its spans have landed.
+    def wait_for_replay_persistence(http_client, test_run_id, sdk_trace_ids)
+      expected_span_counts = Transport.take_replay_span_counts(sdk_trace_ids.compact)
+      return if expected_span_counts.empty?
+
+      unless Bitfab.flush_traces(timeout: PERSISTENCE_TIMEOUT_SECONDS)
+        raise "Replay traces could not be flushed before the delivery deadline " \
+          "(test_run_id #{test_run_id})."
+      end
+
+      deadline = Otel.monotonic_now + PERSISTENCE_TIMEOUT_SECONDS
+      missing = expected_span_counts.keys
+
+      while missing.any?
+        status = http_client.get_replay_status(test_run_id, expected_span_counts)
+        ready = status["traceIds"]
+        ready = {} unless ready.is_a?(Hash)
+        missing = expected_span_counts.keys - ready.keys
+        return if missing.empty?
+        break if Otel.monotonic_now >= deadline
+
+        sleep((deadline - Otel.monotonic_now).clamp(0, PERSISTENCE_POLL_SECONDS))
+      end
+
+      raise "Replay traces were not fully persisted before the delivery deadline " \
+        "(test_run_id #{test_run_id}, missing #{missing.length} of " \
+        "#{expected_span_counts.length} trace(s))."
+    end
+
     # Normalize a complete-replay tokens hash (string-keyed JSON) into the
     # symbol-keyed shape the replay item exposes. Nil when the server reported
     # no token data for this trace.
@@ -894,11 +940,6 @@ module Bitfab
       # complete-replay loop). Carried on the item under :_sdk_trace_id, never
       # surfaced as the public :trace_id.
       sdk_trace_id = SecureRandom.uuid
-      # Collects the root span's persistence threads (span uploads + trace
-      # completion). Joined below so this item's trace is on the server
-      # before run() calls complete_replay: otherwise the server's trace-ID
-      # mapping races the uploads and the item's trace_id comes back nil.
-      pending_persistence = []
 
       ReplayContext.with_context(
         test_run_id:,
@@ -909,7 +950,6 @@ module Bitfab
         mock_strategy:,
         mock_overrides:,
         fetch_span_output:,
-        pending_persistence:,
         db_branch_lease:,
         source_bitfab_trace_id:
       ) do
@@ -935,12 +975,6 @@ module Bitfab
         fn_error = e.message
       end
 
-      # Wait for this item's trace (spans + completion) to be fully persisted
-      # before the item resolves. Runs on the error path too: a raising
-      # method still emits a root span whose trace must land before
-      # complete_replay. Joins are bounded by the HTTP layer's own timeouts.
-      pending_persistence.each(&:join)
-
       {
         input: args,
         result: fn_result,

data/lib/bitfab/serialize.rb

@@ -236,5 +236,45 @@ module Bitfab
         [[raw_input], {}]
       end
     end
+
+    # JSON-encode a request or carrier body without ever raising on a stray
+    # value.
+    #
+    # Upstream serialization (serialize_value) should already have flattened
+    # user data. This is the boundary backstop: if anything non-serializable
+    # still slips through, it is run through serialize_value (which never
+    # raises and stubs strays) instead of letting JSON.generate raise and drop
+    # the whole span/trace silently. A degraded payload warns loudly so the
+    # trace isn't quietly left incomplete or not replayable.
+    def safe_generate(payload)
+      JSON.generate(payload)
+    rescue => e
+      Bitfab.warn_once(
+        "request-body-stubbed",
+        "a request body held a non-serializable value (#{e.message}); it was " \
+        "stubbed so the span still sends, but the trace may be incomplete or " \
+        "not replayable. Capture a JSON-safe projection of this input to make " \
+        "it replayable."
+      )
+
+      begin
+        JSON.generate(sanitize_payload(payload))
+      rescue
+        # Truly pathological. Still never drop silently: send a marker body.
+        JSON.generate({"error" => "payload_serialize_failed"})
+      end
+    end
+
+    # Serialize each top-level value so a bad or oversize value is stubbed in
+    # place while the payload keeps its object shape. Running serialize_value on
+    # the whole payload could collapse the entire body to a single stub string
+    # (oversize/cyclic), sending a JSON string instead of a span object.
+    def sanitize_payload(payload)
+      return {"error" => "payload_serialize_failed"} unless payload.is_a?(Hash)
+
+      payload.each_with_object({}) do |(k, v), acc|
+        acc[k.to_s] = serialize_value(v)
+      end
+    end
   end
 end

data/lib/bitfab/transport.rb

@@ -0,0 +1,27 @@
+# frozen_string_literal: true
+
+require_relative "otel"
+
+module Bitfab
+  # Framework-neutral seam between span/trace sending and the delivery
+  # mechanism, so callers never depend on OpenTelemetry directly.
+  module Transport
+    class << self
+      def create_trace_transport(api_key:, direct_sender:)
+        Otel.create_transport(api_key:, direct_sender:)
+      end
+
+      def flush_trace_transports(timeout = 30.0)
+        Otel.flush_transports(timeout)
+      end
+
+      def shutdown_trace_transports(timeout = 30.0)
+        Otel.shutdown_transports(timeout)
+      end
+
+      def take_replay_span_counts(trace_ids)
+        Otel.take_replay_span_counts(trace_ids)
+      end
+    end
+  end
+end

data/lib/bitfab/version.rb

@@ -1,5 +1,5 @@
 # frozen_string_literal: true
 
 module Bitfab
-  VERSION = "0.33.4"
+  VERSION = "0.33.6"
 end

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: bitfab
 version: !ruby/object:Gem::Version
-  version: 0.33.4
+  version: 0.33.6
 platform: ruby
 authors:
 - Harvest Team
@@ -23,6 +23,40 @@ dependencies:
     - - ">="
       - !ruby/object:Gem::Version
         version: '0'
+- !ruby/object:Gem::Dependency
+  name: opentelemetry-sdk
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: '1.2'
+    - - "<"
+      - !ruby/object:Gem::Version
+        version: '2'
+  type: :runtime
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: '1.2'
+    - - "<"
+      - !ruby/object:Gem::Version
+        version: '2'
+- !ruby/object:Gem::Dependency
+  name: opentelemetry-exporter-otlp
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: '0.30'
+  type: :development
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: '0.30'
 - !ruby/object:Gem::Dependency
   name: rake
   requirement: !ruby/object:Gem::Requirement
@@ -122,11 +156,14 @@ files:
 - lib/bitfab/db_snapshot.rb
 - lib/bitfab/http_client.rb
 - lib/bitfab/mock_override.rb
+- lib/bitfab/otel.rb
+- lib/bitfab/otel_collector.rb
 - lib/bitfab/replay.rb
 - lib/bitfab/replay_branch.rb
 - lib/bitfab/serialize.rb
 - lib/bitfab/span_context.rb
 - lib/bitfab/traceable.rb
+- lib/bitfab/transport.rb
 - lib/bitfab/version.rb
 - lib/bitfab/warn_once.rb
 homepage: https://bitfab.ai