fiber_stream 0.2.0 → 0.3.0

data/lib/fiber_stream/pull/ractor_map_boundary.rb

@@ -10,7 +10,18 @@ module FiberStream
     # is below `workers`, preserving bounded backpressure and ordered output.
     class RactorMapBoundary
       TERMINAL_RESULT_CAPACITY = 1
-      READY_WAIT_INTERVAL = 0.001
+      Job = ::Data.define(:sequence, :value)
+      Shutdown = ::Data.define
+      Ready = ::Data.define(:worker_id)
+      WorkerValue = ::Data.define(:worker_id, :sequence, :value)
+      WorkerFailure = ::Data.define(:worker_id, :sequence, :kind, :cause_class_name, :cause_message)
+      Stopped = ::Data.define(:worker_id)
+      ResultValue = ::Data.define(:sequence, :value)
+      ResultDone = ::Data.define(:sequence)
+      ResultError = ::Data.define(:sequence, :error)
+
+      private_constant :Job, :Shutdown, :Ready, :WorkerValue, :WorkerFailure, :Stopped
+      private_constant :ResultValue, :ResultDone, :ResultError
 
       def initialize(upstream, workers, input_transfer, output_transfer, transform)
         @upstream = upstream
@@ -53,6 +64,8 @@ module FiberStream
         @done = true
         close_error = close_upstream
         close_admission(close_upstream: false)
+        close_ready_queue
+        close_result_queue
         request_worker_shutdown
         wait_for_workers
         close_error ||= @upstream_close_error
@@ -103,7 +116,7 @@ module FiberStream
           break unless worker
 
           message = pull_job_message
-          if message.fetch(0) == :job
+          if message.is_a?(Job)
             @in_flight += 1
             break unless deliver_job(worker, message)
           else
@@ -120,19 +133,23 @@ module FiberStream
 
         sequence = @next_sequence
         @next_sequence += 1
-        [:job, sequence, value]
+        Job.new(sequence, value)
       rescue StandardError => error
         close_upstream(record_error: false)
-        [:error, @next_sequence, error]
+        ResultError.new(sequence: @next_sequence, error:)
       end
 
       def terminal_done_message
         close_error = close_upstream
-        close_error ? [:error, @next_sequence, close_error] : [:done, @next_sequence]
+        if close_error
+          ResultError.new(sequence: @next_sequence, error: close_error)
+        else
+          ResultDone.new(sequence: @next_sequence)
+        end
       end
 
       def deliver_job(worker, message)
-        sequence = message.fetch(1)
+        sequence = message.sequence
         track_worker_job(worker, sequence)
 
         if @input_transfer == :move
@@ -143,8 +160,7 @@ module FiberStream
         true
       rescue StandardError => error
         clear_worker_job(worker)
-        sequence = message.fetch(1)
-        record_result([:error, sequence, build_ractor_map_error(sequence, :input_transfer, error)])
+        record_result(ResultError.new(sequence:, error: build_ractor_map_error(sequence, :input_transfer, error)))
         false
       end
 
@@ -162,28 +178,23 @@ module FiberStream
       end
 
       def emit(message)
-        case message.fetch(0)
-        when :value
-          emit_value(message)
-        when :done
+        case message
+        in ResultValue[sequence:, value:]
+          emit_value(sequence, value)
+        in ResultDone
           complete
-        when :error
-          fail_with_ordered_error(message)
+        in ResultError[sequence:, error:]
+          fail_with_ordered_error(sequence, error)
         end
       end
 
-      def emit_value(message)
-        sequence = message.fetch(1)
-        value = message.fetch(2)
+      def emit_value(sequence, value)
         @next_emit_sequence = sequence + 1
         @in_flight -= 1 if @in_flight.positive?
         value
       end
 
-      def fail_with_ordered_error(message)
-        sequence = message.fetch(1)
-        error = message.fetch(2)
-
+      def fail_with_ordered_error(sequence, error)
         if @failure_sequence && sequence > @failure_sequence
           @next_emit_sequence = sequence + 1
           @in_flight -= 1 if @in_flight.positive?
@@ -207,14 +218,14 @@ module FiberStream
       end
 
       def record_result(message)
-        if message.fetch(0) == :error
-          sequence = message.fetch(1)
+        if message.is_a?(ResultError)
+          sequence = message.sequence
           @failure_sequence = sequence if @failure_sequence.nil? || sequence < @failure_sequence
           close_admission
           request_worker_shutdown
         end
 
-        @pending[message.fetch(1)] = message
+        @pending[message.sequence] = message
       end
 
       def drain_available_results
@@ -282,39 +293,39 @@ module FiberStream
       end
 
       def handle_worker_message(message, live_workers)
-        case message.fetch(0)
-        when :ready
-          deliver_ready_worker(message.fetch(1))
+        case message
+        in Ready[worker_id]
+          deliver_ready_worker(worker_id)
           0
-        when :value
+        in WorkerValue
           handle_worker_value_message(message)
           0
-        when :error
+        in WorkerFailure
           handle_worker_error_message(message)
           0
-        when :stopped
+        in Stopped
           handle_worker_stopped_message(message, live_workers)
         end
       end
 
       def handle_worker_value_message(message)
-        worker = worker_for_id(message.fetch(1))
-        sequence = message.fetch(2)
-        value = message.fetch(3)
+        worker = worker_for_id(message.worker_id)
+        sequence = message.sequence
+        value = message.value
 
         clear_worker_job(worker)
-        deliver_result([:value, sequence, value])
+        deliver_result(ResultValue.new(sequence:, value:))
       end
 
       def handle_worker_error_message(message)
-        worker = worker_for_id(message.fetch(1))
+        worker = worker_for_id(message.worker_id)
 
         clear_worker_job(worker)
         deliver_result(normalize_worker_error_message(message))
       end
 
       def handle_worker_stopped_message(message, live_workers)
-        worker = worker_for_id(message.fetch(1))
+        worker = worker_for_id(message.worker_id)
         live_workers.delete(worker)
         sequence = clear_worker_job(worker)
         deliver_worker_termination_error(worker, sequence) if sequence && !@closed && !@worker_shutdown_sent
@@ -339,38 +350,34 @@ module FiberStream
             cause: cause
           )
 
-        deliver_result([:error, sequence, error])
+        deliver_result(ResultError.new(sequence:, error:))
       end
 
       def deliver_ready_worker(worker_id)
         return if @closed
 
-        push_until_delivered_or_closed(@ready_workers, worker_for_id(worker_id), suppress_data: false)
+        push_until_delivered_or_closed(@ready_workers, worker_for_id(worker_id))
       end
 
       def deliver_result(message)
         return if @closed
 
-        push_until_delivered_or_closed(@results, message, suppress_data: true)
+        push_until_delivered_or_closed(@results, message)
       end
 
-      def push_until_delivered_or_closed(queue, message, suppress_data:)
-        loop do
-          return if @closed && suppress_data
-          return if @closed && !suppress_data
+      def push_until_delivered_or_closed(queue, message)
+        return if @closed
 
-          queue.push(message, true)
-          return
-        rescue ThreadError, ClosedQueueError
-          sleep READY_WAIT_INTERVAL
-        end
+        queue.push(message)
+      rescue ThreadError, ClosedQueueError
+        nil
       end
 
       def normalize_worker_error_message(message)
-        sequence = message.fetch(2)
-        kind = message.fetch(3)
-        cause_class_name = message.fetch(4)
-        cause_message = message.fetch(5)
+        sequence = message.sequence
+        kind = message.kind
+        cause_class_name = message.cause_class_name
+        cause_message = message.cause_message
         error =
           RactorMapError.new(
             sequence: sequence,
@@ -379,7 +386,7 @@ module FiberStream
             cause_message: cause_message
           )
 
-        [:error, sequence, error]
+        ResultError.new(sequence:, error:)
       end
 
       def worker_for_id(worker_id)
@@ -411,7 +418,7 @@ module FiberStream
 
         @worker_shutdown_sent = true
         @workers.each do |worker|
-          worker.send([:shutdown])
+          worker.send(Shutdown.new)
         rescue StandardError
           nil
         end
@@ -420,7 +427,6 @@ module FiberStream
       def wait_for_workers
         return unless @coordinator
 
-        sleep READY_WAIT_INTERVAL while @coordinator.alive?
         @coordinator.join
       end
 
@@ -456,40 +462,58 @@ module FiberStream
       def self.spawn_worker(worker_id, result_port, transform, output_transfer)
         Ractor.new(worker_id, result_port, transform, output_transfer) do |id, port, mapper, transfer|
           current_sequence = nil
+          send_control =
+            lambda do |message|
+              port.send(message)
+              true
+            rescue Exception # rubocop:disable Lint/RescueException
+              false
+            end
+          send_failure =
+            lambda do |sequence, kind, error|
+              send_control.call(WorkerFailure.new(id, sequence, kind, error.class.name, error.message))
+            rescue Exception # rubocop:disable Lint/RescueException
+              false
+            end
 
           begin
-            port.send([:ready, id])
-
-            loop do
-              message = Ractor.receive
-              break if message.fetch(0) == :shutdown
-
-              current_sequence = message.fetch(1)
-              value = message.fetch(2)
-              begin
-                mapped_value = mapper.call(value)
-              rescue Exception => error # rubocop:disable Lint/RescueException
-                port.send([:error, id, current_sequence, :worker, error.class.name, error.message])
-              else
+            if send_control.call(Ready.new(id))
+              loop do
+                message = Ractor.receive
+                case message
+                in Shutdown
+                  break
+                in Job[sequence, value]
+                  current_sequence = sequence
+                else
+                  raise TypeError, "invalid ractor_map worker message: #{message.class}"
+                end
+
                 begin
-                  if transfer == :move
-                    port.send([:value, id, current_sequence, mapped_value], move: true)
-                  else
-                    port.send([:value, id, current_sequence, mapped_value])
-                  end
+                  mapped_value = mapper.call(value)
                 rescue Exception => error # rubocop:disable Lint/RescueException
-                  port.send([:error, id, current_sequence, :output_transfer, error.class.name, error.message])
+                  break unless send_failure.call(current_sequence, :worker, error)
+                else
+                  begin
+                    if transfer == :move
+                      port.send(WorkerValue.new(id, current_sequence, mapped_value), move: true)
+                    else
+                      port.send(WorkerValue.new(id, current_sequence, mapped_value))
+                    end
+                  rescue Exception => error # rubocop:disable Lint/RescueException
+                    break unless send_failure.call(current_sequence, :output_transfer, error)
+                  end
                 end
-              end
 
-              current_sequence = nil
-              port.send([:ready, id])
+                current_sequence = nil
+                break unless send_control.call(Ready.new(id))
+              end
             end
           rescue Exception => error # rubocop:disable Lint/RescueException
             sequence = current_sequence || -1
-            port.send([:error, id, sequence, :worker_termination, error.class.name, error.message])
+            send_failure.call(sequence, :worker_termination, error)
           ensure
-            port.send([:stopped, id])
+            send_control.call(Stopped.new(id))
           end
         end
       end

data/lib/fiber_stream/pull/ractor_merge_ports_source.rb

@@ -0,0 +1,358 @@
+# frozen_string_literal: true
+
+module FiberStream
+  module Pull
+    # Pull stream for `Source.ractor_merge_ports`.
+    #
+    # A coordinator thread owns blocking Ractor waits and forwards producer
+    # messages through a bounded result mailbox. Each producer receives at most
+    # one outstanding ack, and downstream demand replenishes only the producer
+    # that emitted the previous value.
+    class RactorMergePortsSource
+      PortPair = Data.define(:side, :port, :ack_port)
+      StartCommand = Data.define
+      RequestAckCommand = Data.define(:side)
+      ShutdownCommand = Data.define
+      ValueResult = Data.define(:side, :value)
+      DoneResult = Data.define(:side)
+      ErrorResult = Data.define(:side, :error)
+      private_constant :PortPair, :StartCommand, :RequestAckCommand, :ShutdownCommand
+      private_constant :ValueResult, :DoneResult, :ErrorResult
+
+      def initialize(port_pairs, ack_transfer, cancel)
+        @pairs = port_pairs.each_with_index.map do |pair, side|
+          PortPair.new(side:, port: pair.port, ack_port: pair.ack_port)
+        end.freeze
+        @ack_transfer = ack_transfer
+        @cancel_enabled = cancel
+        @result_mailbox = RactorMergeResultMailbox.new(@pairs.size)
+        @control_port = nil
+        @coordinator = nil
+        @state_mutex = Mutex.new
+        @producer_terminal = @pairs.to_h { |pair| [pair.side, false] }
+        @side_done = @pairs.to_h { |pair| [pair.side, false] }
+        @cancel_sent = @pairs.to_h { |pair| [pair.side, false] }
+        @pending_ack_sides = {}
+        @started = false
+        @closed = false
+        @done = false
+      end
+
+      def next
+        return DONE if closed_or_done?
+
+        start
+        request_pending_acks
+        next_result
+      end
+
+      def close
+        started = mark_closed
+        return if started.nil?
+        return unless started
+
+        close_result_mailbox
+        wake_coordinator
+        wait_for_coordinator
+        cancel_error = cancel_non_terminal_producers
+        raise cancel_error if cancel_error
+      end
+
+      private
+
+      def closed_or_done?
+        @state_mutex.synchronize { @closed || @done }
+      end
+
+      def mark_closed
+        already_closed = false
+        started = false
+
+        @state_mutex.synchronize do
+          already_closed = @closed
+          started = @started
+          @closed = true
+          @done = true
+        end
+
+        already_closed ? nil : started
+      end
+
+      def start
+        started_now = false
+
+        @state_mutex.synchronize do
+          return if @started || @closed
+
+          @control_port = Ractor::Port.new
+          @started = true
+          @coordinator = Thread.new { run_coordinator }
+          started_now = true
+        end
+
+        send_control_command(StartCommand.new) if started_now && !closed_or_done?
+      end
+
+      def request_pending_acks
+        sides =
+          @state_mutex.synchronize do
+            pending = @pending_ack_sides.keys
+            @pending_ack_sides.clear
+            pending
+          end
+
+        sides.each { |side| send_control_command(RequestAckCommand.new(side:)) }
+      end
+
+      def next_result
+        loop do
+          result = @result_mailbox.pop
+          return complete if result.nil?
+
+          case result
+          in ValueResult[side:, value:]
+            record_pending_ack(side)
+            return value
+          in DoneResult[side:]
+            mark_side_done(side)
+            return complete if all_done?
+          in ErrorResult[error:]
+            mark_done
+            raise_error(error)
+          end
+        end
+      rescue RactorMergeResultMailbox::Closed
+        complete
+      end
+
+      def record_pending_ack(side)
+        @state_mutex.synchronize do
+          @pending_ack_sides[side] = true unless @closed || @producer_terminal.fetch(side)
+        end
+      end
+
+      def mark_side_done(side)
+        @state_mutex.synchronize { @side_done[side] = true }
+      end
+
+      def all_done?
+        @state_mutex.synchronize { @side_done.values.all? }
+      end
+
+      def mark_done
+        @state_mutex.synchronize { @done = true }
+      end
+
+      def complete
+        mark_done
+        DONE
+      end
+
+      def run_coordinator
+        outstanding_ack = @pairs.to_h { |pair| [pair.side, false] }
+        active_ports = @pairs.map(&:port)
+        pair_by_port = @pairs.to_h { |pair| [pair.port, pair] }
+
+        loop do
+          break if active_ports.empty?
+
+          selected, message = Ractor.select(@control_port, *active_ports)
+          if selected == @control_port
+            break if handle_control_message(message, outstanding_ack)
+          else
+            pair = pair_by_port.fetch(selected)
+            outstanding_ack[pair.side] = false
+            result = build_result(pair, message)
+            active_ports.delete(pair.port) if producer_terminal?(pair.side)
+            deliver_result(result)
+          end
+        end
+      rescue StandardError => error
+        deliver_result(ErrorResult.new(side: nil, error: build_error(:receive, error)))
+      end
+
+      def handle_control_message(message, outstanding_ack)
+        case message
+        in StartCommand
+          @pairs.each { |pair| ack_pair(pair, outstanding_ack) }
+          false
+        in RequestAckCommand[side:]
+          pair = @pairs.fetch(side)
+          ack_pair(pair, outstanding_ack)
+          false
+        in ShutdownCommand
+          true
+        else
+          false
+        end
+      end
+
+      def ack_pair(pair, outstanding_ack)
+        return if outstanding_ack.fetch(pair.side)
+        return if producer_terminal?(pair.side)
+
+        ack_error = send_ack(pair)
+        if ack_error
+          deliver_result(ErrorResult.new(side: pair.side, error: ack_error))
+        else
+          outstanding_ack[pair.side] = true
+        end
+      end
+
+      def send_ack(pair)
+        send_control(pair.ack_port, RactorPort::Ack.new)
+        nil
+      rescue StandardError => error
+        build_error(:ack_transfer, error)
+      end
+
+      def build_result(pair, message)
+        side = pair.side
+
+        case message
+        in RactorPort::Element[value]
+          ValueResult.new(side:, value:)
+        in RactorPort::Complete
+          mark_producer_terminal(side)
+          DoneResult.new(side:)
+        in RactorPort::Failure[String => cause_class_name, String => cause_message]
+          mark_producer_terminal(side)
+          ErrorResult.new(
+            side:,
+            error: RactorPortSourceError.new(
+              kind: :producer_failure,
+              cause_class_name: cause_class_name,
+              cause_message: cause_message
+            )
+          )
+        in RactorPort::Failure
+          ErrorResult.new(side:, error: invalid_message_error(message, "Failure payloads must be Strings"))
+        else
+          ErrorResult.new(side:, error: invalid_message_error(message, "invalid RactorPort message"))
+        end
+      end
+
+      def invalid_message_error(message, cause_message)
+        RactorPortSourceError.new(
+          kind: :invalid_message,
+          cause_class_name: message.class.name,
+          cause_message: cause_message
+        )
+      end
+
+      def build_error(kind, error)
+        RactorPortSourceError.new(
+          kind: kind,
+          cause_class_name: error.class.name,
+          cause_message: error.message,
+          cause: error
+        )
+      end
+
+      def mark_producer_terminal(side)
+        @state_mutex.synchronize { @producer_terminal[side] = true }
+      end
+
+      def producer_terminal?(side)
+        @state_mutex.synchronize { @producer_terminal.fetch(side) }
+      end
+
+      def cancel_non_terminal_producers
+        first_error = nil
+
+        @pairs.each do |pair|
+          next unless should_cancel_pair?(pair)
+
+          cancel_error = send_cancel(pair)
+          first_error ||= cancel_error
+        end
+
+        first_error
+      end
+
+      def should_cancel_pair?(pair)
+        @state_mutex.synchronize do
+          return false unless @cancel_enabled
+          return false if @producer_terminal.fetch(pair.side)
+          return false if @cancel_sent.fetch(pair.side)
+
+          @cancel_sent[pair.side] = true
+          true
+        end
+      end
+
+      def send_cancel(pair)
+        send_control(pair.ack_port, RactorPort::Cancel.new(:closed))
+        nil
+      rescue StandardError => error
+        build_error(:cancel_transfer, error)
+      end
+
+      def send_control(port, message)
+        if @ack_transfer == :move
+          port.send(message, move: true)
+        else
+          port.send(message)
+        end
+      end
+
+      def raise_error(error)
+        if error.is_a?(RactorPortSourceError) && error.original_cause
+          raise error, cause: error.original_cause
+        end
+
+        raise error
+      end
+
+      def deliver_result(result)
+        @result_mailbox.push(result)
+      rescue RactorMergeResultMailbox::Closed
+        nil
+      end
+
+      def send_control_command(command)
+        @control_port&.send(command)
+      rescue StandardError
+        nil
+      end
+
+      def wake_coordinator
+        send_control_command(ShutdownCommand.new)
+      end
+
+      def wait_for_coordinator
+        return unless @coordinator
+
+        @coordinator.join
+      end
+
+      def close_result_mailbox
+        @result_mailbox.close
+      end
+
+      class RactorMergeResultMailbox
+        Closed = Class.new(StandardError)
+
+        def initialize(capacity)
+          @queue = Thread::SizedQueue.new(capacity)
+        end
+
+        def push(result)
+          @queue << result
+        rescue ClosedQueueError
+          raise Closed
+        end
+
+        def pop
+          @queue.pop
+        rescue ClosedQueueError
+          raise Closed
+        end
+
+        def close
+          @queue.close
+        end
+      end
+    end
+  end
+end