fiber_stream 0.3.0 → 0.4.0

data/lib/fiber_stream/pull/parallel_unordered_map_boundary.rb

@@ -0,0 +1,311 @@
+# frozen_string_literal: true
+
+module FiberStream
+  module Pull
+    # Unordered scheduler-backed worker boundary for
+    # `Flow.parallel_unordered_map`.
+    #
+    # A single dispatcher pulls upstream and a bounded worker pool maps values.
+    # Downstream emits worker results in completion order. Admission is
+    # permit-based to keep queued, running, and completed pulled-but-unemitted
+    # work bounded by the configured concurrency.
+    class ParallelUnorderedMapBoundary
+      TERMINAL_RESULT_CAPACITY = 1
+      CancellationError = Class.new(StandardError)
+      JobMessage = Data.define(:sequence, :value)
+      ValueMessage = Data.define(:sequence, :value)
+      DoneMessage = Data.define
+      ErrorMessage = Data.define(:sequence, :error)
+      CloseErrorMessage = Data.define(:error)
+      private_constant :JobMessage, :ValueMessage, :DoneMessage, :ErrorMessage, :CloseErrorMessage
+
+      def initialize(upstream, concurrency, transform)
+        @upstream = upstream
+        @concurrency = concurrency
+        @transform = transform
+        @permits = Thread::SizedQueue.new(concurrency)
+        @jobs = Thread::SizedQueue.new(concurrency)
+        @results = Thread::SizedQueue.new(concurrency + TERMINAL_RESULT_CAPACITY)
+        @workers = []
+        @dispatcher = nil
+        @next_sequence = 0
+        @outstanding_jobs = 0
+        @terminal_message = nil
+        @started = false
+        @closed = false
+        @done = false
+        @admission_closed = false
+        @upstream_closing = false
+        @upstream_closed = false
+        @upstream_close_error = nil
+        @upstream_close_done = Thread::SizedQueue.new(1)
+
+        concurrency.times { @permits << true }
+      end
+
+      def next
+        return DONE if @closed || @done
+
+        start
+        next_message
+      end
+
+      def close
+        return if @closed
+
+        @closed = true
+        @done = true
+        close_error = close_upstream
+        close_internal_queues
+        close_error ||= @upstream_close_error
+        raise close_error if close_error
+      ensure
+        cancel_fibers
+      end
+
+      private
+
+      def start
+        return if @started
+
+        validate_scheduler!
+
+        @started = true
+        @concurrency.times do
+          @workers << Fiber.schedule { run_worker }
+        end
+        @dispatcher = Fiber.schedule { run_dispatcher }
+      end
+
+      def next_message
+        loop do
+          return emit_terminal(@terminal_message) if terminal_ready?
+
+          message = @results.pop
+          return complete if message.nil?
+
+          case message
+          in ValueMessage[sequence:, value:]
+            return emit_value(sequence, value)
+          in DoneMessage | CloseErrorMessage
+            @terminal_message = message
+          in ErrorMessage[sequence:, error:]
+            fail_with_error(sequence, error)
+          end
+        end
+      end
+
+      def emit_value(_sequence, value)
+        @outstanding_jobs -= 1
+        return_permit unless @admission_closed
+        value
+      end
+
+      def terminal_ready?
+        @terminal_message && @outstanding_jobs.zero?
+      end
+
+      def emit_terminal(message)
+        case message
+        in DoneMessage
+          complete
+        in CloseErrorMessage[error:]
+          fail_with_error(@next_sequence, error, close_admission: false)
+        end
+      end
+
+      def fail_with_error(_sequence, error, close_admission: true)
+        @done = true
+        close_admission() if close_admission
+        close_result_queue
+        cancel_fibers
+        raise error
+      end
+
+      def complete
+        @done = true
+        close_result_queue
+        DONE
+      end
+
+      def run_dispatcher
+        loop do
+          break if @closed || @admission_closed
+          break unless take_permit
+
+          message = pull_job_message
+          if message.is_a?(JobMessage)
+            break unless deliver_job(message)
+          else
+            close_admission(close_upstream: false)
+            deliver_result(message)
+            break
+          end
+        end
+      rescue CancellationError
+        nil
+      ensure
+        close_upstream unless @upstream_closed || @closed
+        close_job_queue
+      end
+
+      def pull_job_message
+        value = @upstream.next
+        return terminal_done_message if Pull.done?(value)
+
+        sequence = @next_sequence
+        @next_sequence += 1
+        @outstanding_jobs += 1
+        JobMessage.new(sequence:, value:)
+      rescue StandardError => error
+        close_upstream(record_error: false)
+        ErrorMessage.new(sequence: @next_sequence, error:)
+      end
+
+      def terminal_done_message
+        close_error = close_upstream
+        if close_error
+          CloseErrorMessage.new(error: close_error)
+        else
+          DoneMessage.new
+        end
+      end
+
+      def run_worker
+        loop do
+          break if @closed
+
+          message = @jobs.pop
+          break if message.nil?
+
+          deliver_result(map_job(message))
+        end
+      rescue CancellationError
+        nil
+      end
+
+      def map_job(message)
+        sequence = message.sequence
+        value = message.value
+        ValueMessage.new(sequence:, value: @transform.call(value))
+      rescue CancellationError
+        raise
+      rescue StandardError => error
+        ErrorMessage.new(sequence:, error:)
+      end
+
+      def close_admission(close_upstream: true)
+        return if @admission_closed
+
+        @admission_closed = true
+        close_upstream(record_error: false) if close_upstream
+        close_permit_queue
+        close_job_queue
+      end
+
+      def take_permit
+        @permits.pop
+      rescue ClosedQueueError
+        nil
+      end
+
+      def return_permit
+        @permits << true
+      rescue ClosedQueueError
+        nil
+      end
+
+      def deliver_job(message)
+        @jobs << message
+        true
+      rescue ClosedQueueError
+        false
+      end
+
+      def deliver_result(message)
+        @results << message
+        true
+      rescue ClosedQueueError
+        false
+      end
+
+      def close_internal_queues
+        close_permit_queue
+        close_job_queue
+        close_result_queue
+      end
+
+      def close_permit_queue
+        @permits.close
+      end
+
+      def close_job_queue
+        @jobs.close
+      end
+
+      def close_result_queue
+        @results.close
+      end
+
+      def close_upstream(record_error: true)
+        return wait_for_upstream_close(record_error:) if @upstream_closing
+        return nil if @upstream_closed
+
+        @upstream_closing = true
+        @upstream.close
+        nil
+      rescue StandardError => error
+        @upstream_close_error ||= error if record_error
+        error
+      ensure
+        if @upstream_closing
+          @upstream_closed = true
+          @upstream_closing = false
+          signal_upstream_close_done
+        end
+      end
+
+      def wait_for_upstream_close(record_error:)
+        @upstream_close_done.pop
+        return @upstream_close_error if record_error
+
+        nil
+      rescue ClosedQueueError
+        record_error ? @upstream_close_error : nil
+      end
+
+      def signal_upstream_close_done
+        @upstream_close_done << true
+      rescue ClosedQueueError
+        nil
+      ensure
+        @upstream_close_done.close
+      end
+
+      def cancel_fibers
+        scheduler = Fiber.scheduler
+        return unless scheduler.respond_to?(:fiber_interrupt)
+
+        (@workers + [@dispatcher]).compact.each do |fiber|
+          next unless fiber.alive?
+
+          scheduler.fiber_interrupt(fiber, CancellationError.new)
+        rescue StandardError
+          nil
+        end
+      end
+
+      def validate_scheduler!
+        return if Fiber.scheduler && !Fiber.current.blocking?
+
+        message =
+          if Fiber.scheduler
+            "Flow.parallel_unordered_map requires a non-blocking fiber"
+          else
+            "Flow.parallel_unordered_map requires Fiber.scheduler"
+          end
+        raise SchedulerRequiredError, message
+      end
+    end
+  end
+end

data/lib/fiber_stream/pull/ractor_map_boundary.rb

@@ -80,8 +80,7 @@ module FiberStream
         @started = true
         @result_port = Ractor::Port.new
         @workers_count.times do |worker_id|
-          @workers << self.class.__send__(
-            :spawn_worker,
+          @workers << self.class.spawn_worker(
             worker_id,
             @result_port,
             @transform,
@@ -459,66 +458,66 @@ module FiberStream
         )
       end
 
-      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
+      class << self
+        def spawn_worker(worker_id, result_port, transform, output_transfer) # :nodoc:
+          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
-            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 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
-                  mapped_value = mapper.call(value)
-                rescue Exception => error # rubocop:disable Lint/RescueException
-                  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
+                    mapped_value = mapper.call(value)
                   rescue Exception => error # rubocop:disable Lint/RescueException
-                    break unless send_failure.call(current_sequence, :output_transfer, error)
+                    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
-                break unless send_control.call(Ready.new(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
+              send_failure.call(sequence, :worker_termination, error)
+            ensure
+              send_control.call(Stopped.new(id))
             end
-          rescue Exception => error # rubocop:disable Lint/RescueException
-            sequence = current_sequence || -1
-            send_failure.call(sequence, :worker_termination, error)
-          ensure
-            send_control.call(Stopped.new(id))
           end
         end
       end
-
-      private_class_method :spawn_worker
     end
   end
 end

data/lib/fiber_stream/pull/ractor_merge_ports_source.rb

@@ -9,7 +9,7 @@ module FiberStream
     # one outstanding ack, and downstream demand replenishes only the producer
     # that emitted the previous value.
     class RactorMergePortsSource
-      PortPair = Data.define(:side, :port, :ack_port)
+      PortPair = Data.define(:side, :port, :ack_port, :producer_ractor)
       StartCommand = Data.define
       RequestAckCommand = Data.define(:side)
       ShutdownCommand = Data.define
@@ -21,7 +21,8 @@ module FiberStream
 
       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)
+          producer_ractor = pair.respond_to?(:producer_ractor) ? pair.producer_ractor : nil
+          PortPair.new(side:, port: pair.port, ack_port: pair.ack_port, producer_ractor:)
         end.freeze
         @ack_transfer = ack_transfer
         @cancel_enabled = cancel
@@ -151,14 +152,28 @@ module FiberStream
       def run_coordinator
         outstanding_ack = @pairs.to_h { |pair| [pair.side, false] }
         active_ports = @pairs.map(&:port)
+        active_ractors = @pairs.filter_map(&:producer_ractor)
         pair_by_port = @pairs.to_h { |pair| [pair.port, pair] }
+        pair_by_ractor =
+          @pairs.each_with_object({}) do |pair, pairs|
+            pairs[pair.producer_ractor] = pair if pair.producer_ractor
+          end
 
         loop do
           break if active_ports.empty?
 
-          selected, message = Ractor.select(@control_port, *active_ports)
+          selected, message = Ractor.select(@control_port, *active_ports, *active_ractors)
           if selected == @control_port
             break if handle_control_message(message, outstanding_ack)
+          elsif pair_by_ractor.key?(selected)
+            pair = pair_by_ractor.fetch(selected)
+            active_ractors.delete(selected)
+            case message
+            in ProducerTerminal | ProducerCancelled
+              next
+            else
+              deliver_result(ErrorResult.new(side: pair.side, error: Pull.ractor_producer_termination_error(message)))
+            end
           else
             pair = pair_by_port.fetch(selected)
             outstanding_ack[pair.side] = false

data/lib/fiber_stream/pull/ractor_port_source.rb

@@ -12,13 +12,18 @@ module FiberStream
       ProtocolMessage = Data.define(:message)
       ErrorMessage = Data.define(:error)
       ClosedMessage = Data.define
+      SelectedProtocol = Data.define(:message)
+      SelectedShutdown = Data.define
+      SelectedProducerTerminated = Data.define(:result)
       private_constant :ProtocolMessage, :ErrorMessage, :ClosedMessage
+      private_constant :SelectedProtocol, :SelectedShutdown, :SelectedProducerTerminated
 
-      def initialize(port, ack_port, ack_transfer, cancel)
+      def initialize(port, ack_port, ack_transfer, cancel, producer_ractor = nil)
         @port = port
         @ack_port = ack_port
         @ack_transfer = ack_transfer
         @cancel_enabled = cancel
+        @producer_ractor = producer_ractor
         @demands = Thread::SizedQueue.new(1)
         @results = Thread::SizedQueue.new(1)
         @shutdown_port = nil
@@ -28,6 +33,7 @@ module FiberStream
         @closed = false
         @done = false
         @producer_terminal = false
+        @producer_ractor_terminal = false
         @cancel_sent = false
       end
 
@@ -156,10 +162,15 @@ module FiberStream
             break
           end
 
-          selected, message = select_message
-          break if selected == @shutdown_port || closed?
-
-          deliver_result(ProtocolMessage.new(message:))
+          case select_message
+          in SelectedShutdown
+            break
+          in SelectedProducerTerminated[result:]
+            deliver_result(ErrorMessage.new(error: Pull.ractor_producer_termination_error(result)))
+            break
+          in SelectedProtocol[message:]
+            deliver_result(ProtocolMessage.new(message:))
+          end
         end
       rescue StandardError => error
         deliver_result(ErrorMessage.new(error: build_error(:receive, error)))
@@ -168,7 +179,29 @@ module FiberStream
       end
 
       def select_message
-        Ractor.select(@port, @shutdown_port)
+        loop do
+          selected, message =
+            if @producer_ractor && !@producer_ractor_terminal
+              Ractor.select(@port, @shutdown_port, @producer_ractor)
+            else
+              Ractor.select(@port, @shutdown_port)
+            end
+
+          if selected == @producer_ractor
+            @producer_ractor_terminal = true
+            case message
+            in ProducerTerminal | ProducerCancelled
+              next
+            else
+              return SelectedProducerTerminated.new(result: message)
+            end
+
+          end
+
+          return SelectedShutdown.new if selected == @shutdown_port || closed?
+
+          return SelectedProtocol.new(message:)
+        end
       end
 
       def send_ack