fiber_stream 0.3.0 → 0.5.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/reject.rb

@@ -0,0 +1,40 @@
+# frozen_string_literal: true
+
+module FiberStream
+  module Pull
+    # Complement filtering stage.
+    #
+    # A single downstream demand may pull multiple upstream elements until the
+    # predicate retains a value or upstream completes. Rejected elements are
+    # discarded immediately and are not buffered.
+    class Reject
+      def initialize(upstream, predicate)
+        @upstream = upstream
+        @predicate = predicate
+        @closed = false
+        @done = false
+      end
+
+      def next
+        return DONE if @closed || @done
+
+        loop do
+          value = @upstream.next
+          if Pull.done?(value)
+            @done = true
+            return DONE
+          end
+
+          return value unless @predicate.call(value)
+        end
+      end
+
+      def close
+        return if @closed
+
+        @closed = true
+        @upstream.close
+      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/tap.rb

@@ -0,0 +1,38 @@
+# frozen_string_literal: true
+
+module FiberStream
+  module Pull
+    # Stateless observing stage.
+    #
+    # It pulls one upstream element for each downstream demand, calls the
+    # observer for real elements, and emits the original element unchanged.
+    class Tap
+      def initialize(upstream, observer)
+        @upstream = upstream
+        @observer = observer
+        @closed = false
+        @done = false
+      end
+
+      def next
+        return DONE if @closed || @done
+
+        value = @upstream.next
+        if Pull.done?(value)
+          @done = true
+          return DONE
+        end
+
+        @observer.call(value)
+        value
+      end
+
+      def close
+        return if @closed
+
+        @closed = true
+        @upstream.close
+      end
+    end
+  end
+end

data/lib/fiber_stream/pull/throttle.rb

@@ -0,0 +1,43 @@
+# frozen_string_literal: true
+
+module FiberStream
+  module Pull
+    # Pull-driven rate-limiting stage.
+    #
+    # The stage pulls at most one upstream value, acquires one permit, and then
+    # emits that value unless the stage was closed while waiting.
+    class Throttle
+      def initialize(upstream, limiter)
+        @upstream = upstream
+        @limiter = limiter
+        @closed = false
+        @done = false
+      end
+
+      def next
+        return DONE if @closed || @done
+
+        value = @upstream.next
+        if Pull.done?(value)
+          @done = true
+          return DONE
+        end
+
+        @limiter.acquire(permits: 1)
+        if @closed
+          @done = true
+          return DONE
+        end
+
+        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
@@ -48,10 +85,30 @@ module FiberStream
       Map.new(upstream, transform)
     end
 
+    def self.filter_map(upstream, transform)
+      FilterMap.new(upstream, transform)
+    end
+
+    def self.compact(upstream)
+      Compact.new(upstream)
+    end
+
+    def self.map_concat(upstream, transform)
+      MapConcat.new(upstream, transform)
+    end
+
+    def self.tap(upstream, observer)
+      Tap.new(upstream, observer)
+    end
+
     def self.parallel_map(upstream, concurrency, transform)
       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
@@ -60,6 +117,10 @@ module FiberStream
       Select.new(upstream, predicate)
     end
 
+    def self.reject(upstream, predicate)
+      Reject.new(upstream, predicate)
+    end
+
     def self.take(upstream, count)
       Take.new(upstream, count)
     end
@@ -72,6 +133,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
@@ -88,6 +153,10 @@ module FiberStream
       BufferBoundary.new(upstream, count)
     end
 
+    def self.throttle(upstream, limiter)
+      Throttle.new(upstream, limiter)
+    end
+
     def self.lines(upstream, chomp, max_length)
       Lines.new(upstream, chomp, max_length)
     end
@@ -104,27 +173,37 @@ 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"
 require_relative "pull/map"
+require_relative "pull/filter_map"
+require_relative "pull/compact"
+require_relative "pull/map_concat"
+require_relative "pull/tap"
 require_relative "pull/select"
+require_relative "pull/reject"
 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"
 require_relative "pull/split"
 require_relative "pull/async_boundary"
 require_relative "pull/buffer_boundary"
+require_relative "pull/throttle"
 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, :FilterMap, :Compact, :MapConcat, :Tap, :Select, :Reject, :Take, :Drop,
+                     :Grouped, :Scan, :TakeWhile, :DropWhile, :Lines, :Split, :AsyncBoundary, :BufferBoundary,
+                     :Throttle, :ParallelMapBoundary, :ParallelUnorderedMapBoundary, :RactorMapBoundary
   end
 end