fiber_stream 0.3.0 → 0.4.0

data/lib/fiber_stream/pull/ractor_producer_source.rb

@@ -0,0 +1,349 @@
+# frozen_string_literal: true
+
+module FiberStream
+  module Pull
+    # Setup adapter for high-level owned Ractor producer sources.
+    #
+    # Producer ractors and ports are created on first demand. Once every
+    # producer has returned its producer-owned ack port, this adapter delegates
+    # demand to the existing low-level Ractor port pull sources.
+    class RactorProducerSource
+      StartedProducer = Data.define(:side, :data_port, :setup_port, :ractor, :definition)
+      ReadyProducer = Data.define(:side, :data_port, :ack_port, :ractor)
+      PortPair = Data.define(:port, :ack_port, :producer_ractor)
+      SetupSuccess = Data.define(:producers)
+      SetupError = Data.define(:error)
+      SetupClosed = Data.define
+      private_constant :StartedProducer, :ReadyProducer, :PortPair, :SetupSuccess, :SetupError, :SetupClosed
+
+      def initialize(definitions, ack_transfer, merge)
+        @definitions = definitions
+        @ack_transfer = ack_transfer
+        @merge = merge
+        @setup_results = Thread::SizedQueue.new(1)
+        @state_mutex = Mutex.new
+        @started = false
+        @closed = false
+        @done = false
+        @shutdown_port = nil
+        @setup_thread = nil
+        @delegate = nil
+        @started_producers = []
+        @ready_producers = []
+      end
+
+      def next
+        return DONE if closed_or_done?
+
+        start
+        return DONE unless ensure_delegate
+
+        value = @delegate.next
+        mark_done if Pull.done?(value)
+        value
+      end
+
+      def close
+        already_closed = mark_closed
+        return if already_closed
+
+        wake_setup
+        wait_for_setup
+        close_setup_queue
+        close_error = close_delegate
+        cancel_ready_producers
+        wait_for_ractors(@ready_producers.map(&:ractor))
+        raise close_error if close_error
+      end
+
+      private
+
+      def closed_or_done?
+        @state_mutex.synchronize { @closed || @done }
+      end
+
+      def mark_done
+        @state_mutex.synchronize { @done = true }
+      end
+
+      def mark_closed
+        @state_mutex.synchronize do
+          already_closed = @closed
+          @closed = true
+          @done = true
+          already_closed
+        end
+      end
+
+      def start
+        @state_mutex.synchronize do
+          return if @started
+
+          @started = true
+          @shutdown_port = Ractor::Port.new
+          spawn_producers
+          @setup_thread = Thread.new { run_setup }
+        end
+      rescue Exception => error # rubocop:disable Lint/RescueException
+        setup_error = build_error(:producer_setup, error)
+        @setup_thread = Thread.new { cleanup_after_start_failure(setup_error) }
+      end
+
+      def cleanup_after_start_failure(setup_error)
+        setup_ports = @started_producers.map(&:setup_port)
+        ractors = @started_producers.map(&:ractor)
+        producer_by_setup_port = @started_producers.to_h { |producer| [producer.setup_port, producer] }
+
+        cleanup_remaining_setup(setup_ports, ractors, producer_by_setup_port)
+        deliver_setup(SetupError.new(error: setup_error))
+      end
+
+      def spawn_producers
+        @started_producers = []
+
+        @definitions.each_with_index do |definition, side|
+          data_port = Ractor::Port.new
+          setup_port = Ractor::Port.new
+          ractor = self.class.spawn_producer(data_port, setup_port, definition)
+          @started_producers << StartedProducer.new(side:, data_port:, setup_port:, ractor:, definition:)
+        end
+      end
+
+      def run_setup
+        remaining_setup_ports = @started_producers.map(&:setup_port)
+        remaining_ractors = @started_producers.map(&:ractor)
+        producer_by_setup_port = @started_producers.to_h { |producer| [producer.setup_port, producer] }
+
+        until remaining_setup_ports.empty?
+          selected, message = Ractor.select(@shutdown_port, *remaining_setup_ports, *remaining_ractors)
+          if selected == @shutdown_port
+            deliver_setup(SetupClosed.new)
+            cleanup_remaining_setup(remaining_setup_ports, remaining_ractors, producer_by_setup_port)
+            return
+          elsif producer_by_setup_port.key?(selected)
+            producer = producer_by_setup_port.fetch(selected)
+            validate_ack_port!(message)
+            @ready_producers << ReadyProducer.new(
+              side: producer.side,
+              data_port: producer.data_port,
+              ack_port: message,
+              ractor: producer.ractor
+            )
+            remaining_setup_ports.delete(selected)
+          else
+            raise "producer exited before setup completed"
+          end
+        end
+
+        deliver_setup(SetupSuccess.new(producers: @ready_producers.sort_by(&:side).freeze))
+      rescue Exception => error # rubocop:disable Lint/RescueException
+        setup_error = build_error(:producer_setup, error)
+        cancel_ready_producers
+        cleanup_remaining_setup(remaining_setup_ports, remaining_ractors, producer_by_setup_port)
+        deliver_setup(SetupError.new(error: setup_error))
+      end
+
+      def cleanup_remaining_setup(remaining_setup_ports, remaining_ractors, producer_by_setup_port)
+        producer_by_ractor = @started_producers.to_h { |producer| [producer.ractor, producer] }
+
+        until remaining_setup_ports.empty?
+          selected, message = Ractor.select(*remaining_setup_ports, *remaining_ractors)
+          if producer_by_setup_port.key?(selected)
+            producer = producer_by_setup_port.fetch(selected)
+            validate_ack_port!(message)
+            ready = ReadyProducer.new(
+              side: producer.side,
+              data_port: producer.data_port,
+              ack_port: message,
+              ractor: producer.ractor
+            )
+            @ready_producers << ready
+            send_cancel(ready)
+            remaining_setup_ports.delete(selected)
+          else
+            producer = producer_by_ractor.fetch(selected)
+            remaining_setup_ports.delete(producer.setup_port)
+            remaining_ractors.delete(selected)
+          end
+        end
+
+        wait_for_ractors(@started_producers.map(&:ractor))
+      rescue Exception # rubocop:disable Lint/RescueException
+        nil
+      end
+
+      def validate_ack_port!(ack_port)
+        return if ack_port.respond_to?(:send) && ack_port.method(:send).owner != Kernel
+
+        raise TypeError, "producer setup did not return a Ractor-style ack port"
+      end
+
+      def ensure_delegate
+        return true if delegate_installed?
+
+        case setup_result
+        in SetupSuccess[producers:]
+          delegate = build_delegate(producers)
+          should_close_delegate = false
+          @state_mutex.synchronize do
+            @delegate = delegate
+            should_close_delegate = @closed
+            @delegate = nil if should_close_delegate
+          end
+          if should_close_delegate
+            close_delegate_suppressing(delegate)
+            return false
+          end
+
+          true
+        in SetupError[error:]
+          mark_done
+          raise_error(error)
+        in SetupClosed
+          mark_done
+          false
+        end
+      end
+
+      def setup_result
+        result = @setup_results.pop
+        result || SetupClosed.new
+      rescue ClosedQueueError
+        SetupClosed.new
+      end
+
+      def build_delegate(producers)
+        if @merge
+          Pull.ractor_merge_ports(
+            producers.map do |producer|
+              PortPair.new(port: producer.data_port, ack_port: producer.ack_port, producer_ractor: producer.ractor)
+            end,
+            @ack_transfer,
+            true
+          )
+        else
+          producer = producers.fetch(0)
+          Pull.ractor_port(producer.data_port, producer.ack_port, @ack_transfer, true, producer.ractor)
+        end
+      end
+
+      def cancel_ready_producers
+        @ready_producers.each do |producer|
+          send_cancel(producer)
+        rescue Exception # rubocop:disable Lint/RescueException
+          nil
+        end
+      end
+
+      def send_cancel(producer)
+        send_control(producer.ack_port, RactorPort::Cancel.new(:closed))
+      end
+
+      def send_control(port, message)
+        if @ack_transfer == :move
+          port.send(message, move: true)
+        else
+          port.send(message)
+        end
+      end
+
+      def close_delegate
+        delegate = @state_mutex.synchronize { @delegate }
+        delegate&.close
+        nil
+      rescue StandardError => error
+        error
+      end
+
+      def close_delegate_suppressing(delegate)
+        delegate.close
+      rescue StandardError
+        nil
+      end
+
+      def delegate_installed?
+        @state_mutex.synchronize { !@delegate.nil? }
+      end
+
+      def wake_setup
+        @shutdown_port&.send(:shutdown)
+      rescue Exception # rubocop:disable Lint/RescueException
+        nil
+      end
+
+      def wait_for_setup
+        @setup_thread&.join
+      end
+
+      def wait_for_ractors(ractors)
+        return if ractors.empty?
+
+        Thread.new do
+          ractors.each do |ractor|
+            ractor.value
+          rescue Exception # rubocop:disable Lint/RescueException
+            nil
+          end
+        end.join
+      end
+
+      def close_setup_queue
+        @setup_results.close
+      end
+
+      def deliver_setup(result)
+        @setup_results.push(result)
+      rescue ClosedQueueError, ThreadError
+        nil
+      end
+
+      def build_error(kind, error)
+        RactorPortSourceError.new(
+          kind: kind,
+          cause_class_name: error.class.name,
+          cause_message: error.message,
+          cause: error
+        )
+      end
+
+      def raise_error(error)
+        if error.is_a?(RactorPortSourceError) && error.original_cause
+          raise error, cause: error.original_cause
+        end
+
+        raise error
+      end
+
+      class << self
+        def spawn_producer(data_port, setup_port, definition) # :nodoc:
+          Ractor.new(
+            data_port,
+            setup_port,
+            definition.block,
+            definition.transfer,
+            definition.args
+          ) do |outbox, setup, block, transfer, args|
+            ack_port = Ractor::Port.new
+            setup.send(ack_port)
+            producer = RactorProducer.new(outbox, ack_port, transfer)
+
+            begin
+              block.call(producer, *args)
+              producer.complete unless producer.terminal? || producer.cancelled?
+            rescue Exception => error # rubocop:disable Lint/RescueException
+              producer.fail(error) unless producer.terminal? || producer.cancelled?
+            end
+
+            if producer.send_failed?
+              ProducerSendFailed.new
+            elsif producer.cancelled?
+              ProducerCancelled.new
+            else
+              ProducerTerminal.new
+            end
+          end
+        end
+      end
+    end
+  end
+end

data/lib/fiber_stream/pull/scan.rb

@@ -0,0 +1,38 @@
+# frozen_string_literal: true
+
+module FiberStream
+  module Pull
+    # Running-accumulator stage.
+    #
+    # It pulls one immediate upstream value for each downstream demand, updates
+    # the accumulator with the reducer, and emits the updated accumulator.
+    class Scan
+      def initialize(upstream, initial, reducer)
+        @upstream = upstream
+        @accumulator = initial
+        @reducer = reducer
+        @closed = false
+        @done = false
+      end
+
+      def next
+        return DONE if @closed || @done
+
+        value = @upstream.next
+        if Pull.done?(value)
+          @done = true
+          return DONE
+        end
+
+        @accumulator = @reducer.call(@accumulator, value)
+      end
+
+      def close
+        return if @closed
+
+        @closed = true
+        @upstream.close
+      end
+    end
+  end
+end

data/lib/fiber_stream/pull.rb

@@ -11,11 +11,40 @@ module FiberStream
   # The sentinel must never escape through public APIs.
   module Pull
     DONE = Object.new.freeze
+    ProducerTerminal = Data.define
+    ProducerCancelled = Data.define
+    ProducerSendFailed = Data.define
+    private_constant :ProducerTerminal, :ProducerCancelled, :ProducerSendFailed
 
     def self.done?(value)
       value.equal?(DONE)
     end
 
+    def self.each_value(stream)
+      loop do
+        value = stream.next
+        break if done?(value)
+
+        yield value
+      end
+    end
+
+    def self.ractor_producer_termination_error(result)
+      cause_message =
+        case result
+        in ProducerSendFailed
+          "producer exited after failing to send the ack-permitted message"
+        else
+          "producer exited before sending the ack-permitted message"
+        end
+
+      RactorPortSourceError.new(
+        kind: :producer_failure,
+        cause_class_name: "RuntimeError",
+        cause_message: cause_message
+      )
+    end
+
     def self.each(enumerable)
       Each.new(enumerable)
     end
@@ -24,14 +53,22 @@ module FiberStream
       IOSource.new(io, chunk_size, close_io)
     end
 
-    def self.ractor_port(port, ack_port, ack_transfer, cancel)
-      RactorPortSource.new(port, ack_port, ack_transfer, cancel)
+    def self.ractor_port(port, ack_port, ack_transfer, cancel, producer_ractor = nil)
+      RactorPortSource.new(port, ack_port, ack_transfer, cancel, producer_ractor)
     end
 
     def self.ractor_merge_ports(port_pairs, ack_transfer, cancel)
       RactorMergePortsSource.new(port_pairs, ack_transfer, cancel)
     end
 
+    def self.ractor_producer(definitions, ack_transfer)
+      RactorProducerSource.new(definitions, ack_transfer, false)
+    end
+
+    def self.ractor_merge_producers(definitions, ack_transfer)
+      RactorProducerSource.new(definitions, ack_transfer, true)
+    end
+
     def self.concat(left_materializer, right_materializer)
       Concat.new(left_materializer, right_materializer)
     end
@@ -52,6 +89,10 @@ module FiberStream
       ParallelMapBoundary.new(upstream, concurrency, transform)
     end
 
+    def self.parallel_unordered_map(upstream, concurrency, transform)
+      ParallelUnorderedMapBoundary.new(upstream, concurrency, transform)
+    end
+
     def self.ractor_map(upstream, workers, input_transfer, output_transfer, transform)
       RactorMapBoundary.new(upstream, workers, input_transfer, output_transfer, transform)
     end
@@ -72,6 +113,10 @@ module FiberStream
       Grouped.new(upstream, count)
     end
 
+    def self.scan(upstream, initial, reducer)
+      Scan.new(upstream, initial, reducer)
+    end
+
     def self.take_while(upstream, predicate)
       TakeWhile.new(upstream, predicate)
     end
@@ -104,6 +149,7 @@ require_relative "pull/each"
 require_relative "pull/io_source"
 require_relative "pull/ractor_port_source"
 require_relative "pull/ractor_merge_ports_source"
+require_relative "pull/ractor_producer_source"
 require_relative "pull/concat"
 require_relative "pull/zip"
 require_relative "pull/merge"
@@ -112,6 +158,7 @@ require_relative "pull/select"
 require_relative "pull/take"
 require_relative "pull/drop"
 require_relative "pull/grouped"
+require_relative "pull/scan"
 require_relative "pull/take_while"
 require_relative "pull/drop_while"
 require_relative "pull/lines"
@@ -119,12 +166,14 @@ require_relative "pull/split"
 require_relative "pull/async_boundary"
 require_relative "pull/buffer_boundary"
 require_relative "pull/parallel_map_boundary"
+require_relative "pull/parallel_unordered_map_boundary"
 require_relative "pull/ractor_map_boundary"
 
 module FiberStream
   module Pull
-    private_constant :Each, :IOSource, :RactorPortSource, :RactorMergePortsSource, :Concat, :Zip, :Merge, :Map,
-                     :Select, :Take, :Drop, :Grouped, :TakeWhile, :DropWhile, :Lines, :Split, :AsyncBoundary,
-                     :BufferBoundary, :ParallelMapBoundary, :RactorMapBoundary
+    private_constant :Each, :IOSource, :RactorPortSource, :RactorMergePortsSource, :RactorProducerSource, :Concat,
+                     :Zip, :Merge, :Map, :Select, :Take, :Drop, :Grouped, :Scan, :TakeWhile, :DropWhile, :Lines, :Split,
+                     :AsyncBoundary, :BufferBoundary, :ParallelMapBoundary, :ParallelUnorderedMapBoundary,
+                     :RactorMapBoundary
   end
 end

data/lib/fiber_stream/ractor_producer.rb

@@ -0,0 +1,167 @@
+# frozen_string_literal: true
+
+module FiberStream
+  # Producer-side context for `Source.ractor_producer`.
+  #
+  # Producer blocks call `emit`, `complete`, or `fail` to send one protocol
+  # message after receiving one downstream acknowledgment. A `false` return
+  # means cooperative cancellation was observed before the requested message
+  # could be sent.
+  class RactorProducer
+    def initialize(data_port, ack_port, transfer)
+      @data_port = data_port
+      @ack_port = ack_port
+      @transfer = transfer
+      @terminal = false
+      @cancelled = false
+      @send_failed = false
+    end
+
+    def emit(value, transfer: nil)
+      return false if terminal? || cancelled?
+
+      message_transfer = validate_transfer_override!(transfer)
+      return false unless wait_for_ack
+
+      send_emitted_message(RactorPort::Element.new(value), message_transfer)
+    end
+
+    def complete
+      return false if terminal? || cancelled?
+      return false unless wait_for_ack
+
+      send_terminal_message(RactorPort::Complete.new)
+    end
+
+    def fail(error = nil, cause_class_name: nil, cause_message: nil)
+      return false if terminal? || cancelled?
+
+      failure = failure_message(error, cause_class_name, cause_message)
+      return false unless wait_for_ack
+
+      send_terminal_message(failure)
+    end
+
+    def cancelled?
+      @cancelled
+    end
+
+    def terminal? # :nodoc:
+      @terminal
+    end
+
+    def send_failed? # :nodoc:
+      @send_failed
+    end
+
+    private
+
+    def send_emitted_message(message, transfer)
+      send_data_message(message, transfer)
+      true
+    rescue Exception => error # rubocop:disable Lint/RescueException
+      report_same_ack_failure(error)
+      false
+    end
+
+    def send_terminal_message(message)
+      send_data_message(message, @transfer)
+      @terminal = true
+      true
+    rescue Exception => send_error # rubocop:disable Lint/RescueException
+      report_same_ack_failure(send_error)
+      false
+    end
+
+    def validate_transfer_override!(transfer)
+      return @transfer if transfer.nil?
+      return transfer if [:copy, :move].include?(transfer)
+
+      raise ArgumentError, "transfer must be :copy or :move"
+    end
+
+    def wait_for_ack
+      case @ack_port.receive
+      in RactorPort::Ack
+        true
+      in RactorPort::Cancel
+        @cancelled = true
+        false
+      else
+        raise TypeError, "invalid ractor producer control message"
+      end
+    end
+
+    def send_data_message(message, transfer)
+      if transfer == :move
+        @data_port.send(message, move: true)
+      else
+        @data_port.send(message)
+      end
+    end
+
+    def failure_message(error, cause_class_name, cause_message)
+      if error
+        return RactorPort::Failure.new(safe_class_name(error), safe_message(error))
+      end
+
+      unless cause_class_name.is_a?(String) && cause_message.is_a?(String)
+        raise ArgumentError, "fail requires an error or String failure metadata"
+      end
+
+      RactorPort::Failure.new(cause_class_name, cause_message)
+    end
+
+    def safe_class_name(error)
+      name = error.class.name
+      name.is_a?(String) && !name.empty? ? name : "Exception"
+    rescue Exception # rubocop:disable Lint/RescueException
+      "Exception"
+    end
+
+    def safe_message(error)
+      message = error.message
+      message.is_a?(String) ? message : ""
+    rescue Exception # rubocop:disable Lint/RescueException
+      ""
+    end
+
+    def report_same_ack_failure(error)
+      send_data_message(RactorPort::Failure.new(safe_class_name(error), safe_message(error)), :copy)
+      @terminal = true
+    rescue Exception # rubocop:disable Lint/RescueException
+      @terminal = true
+      @send_failed = true
+    end
+  end
+
+  # Builder passed to `Source.ractor_merge_producers`.
+  #
+  # Each `producer` call records one lazily started owned producer definition.
+  # Registration validates producer block isolation and transfer policy but
+  # does not create Ractor ports or start producer code.
+  class RactorProducerGroup
+    Definition = Data.define(:args, :transfer, :block)
+    private_constant :Definition
+
+    def initialize(default_transfer)
+      @default_transfer = default_transfer
+      @definitions = []
+    end
+
+    def producer(*args, transfer: nil, &block)
+      raise ArgumentError, "missing block" unless block
+      unless transfer.nil? || [:copy, :move].include?(transfer)
+        raise ArgumentError, "transfer must be :copy or :move"
+      end
+      raise TypeError, "block must be shareable" unless Ractor.shareable?(block)
+
+      @definitions << Definition.new(args:, transfer: transfer || @default_transfer, block:)
+      self
+    end
+
+    def definitions
+      @definitions.dup.freeze
+    end
+  end
+end

data/lib/fiber_stream/running_pipeline.rb

@@ -7,6 +7,10 @@ module FiberStream
     CancelledMessage = Data.define(:error)
     private_constant :ValueMessage, :ErrorMessage, :CancelledMessage
 
+    def self.start(scheduler, &run) # :nodoc:
+      new(scheduler, &run)
+    end
+
     def initialize(scheduler, &run)
       @scheduler = scheduler
       @completion = nil