bitfab 0.33.5 → 0.34.0

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
 
@@ -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
@@ -876,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.
@@ -904,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:,
@@ -919,7 +950,6 @@ module Bitfab
         mock_strategy:,
         mock_overrides:,
         fetch_span_output:,
-        pending_persistence:,
         db_branch_lease:,
         source_bitfab_trace_id:
       ) do
@@ -945,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(direct_sender:)
+        Otel.create_transport(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.5"
+  VERSION = "0.34.0"
 end

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: bitfab
 version: !ruby/object:Gem::Version
-  version: 0.33.5
+  version: 0.34.0
 platform: ruby
 authors:
 - Harvest Team
@@ -23,6 +23,26 @@ 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: rake
   requirement: !ruby/object:Gem::Requirement
@@ -122,11 +142,13 @@ files:
 - lib/bitfab/db_snapshot.rb
 - lib/bitfab/http_client.rb
 - lib/bitfab/mock_override.rb
+- lib/bitfab/otel.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