fiber_stream 0.2.0 → 0.3.0

data/lib/fiber_stream/pull/grouped.rb

@@ -0,0 +1,46 @@
+# frozen_string_literal: true
+
+module FiberStream
+  module Pull
+    # Fixed-size grouping stage.
+    #
+    # It collects adjacent upstream elements into distinct arrays of up to
+    # `count` elements. A final partial group is emitted when upstream completes
+    # normally.
+    class Grouped
+      def initialize(upstream, count)
+        @upstream = upstream
+        @count = count
+        @closed = false
+        @done = false
+      end
+
+      def next
+        return DONE if @closed || @done
+
+        group = []
+
+        while group.length < @count
+          value = @upstream.next
+          if Pull.done?(value)
+            @done = true
+            return DONE if group.empty?
+
+            return group
+          end
+
+          group << value
+        end
+
+        group
+      end
+
+      def close
+        return if @closed
+
+        @closed = true
+        @upstream.close
+      end
+    end
+  end
+end

data/lib/fiber_stream/pull/merge.rb

@@ -0,0 +1,230 @@
+# frozen_string_literal: true
+
+module FiberStream
+  module Pull
+    # Scheduler-backed ready-order source merge.
+    #
+    # Each input source is materialized by a scheduled producer fiber on first
+    # downstream demand. Producers publish values, completion, and failures into
+    # a bounded mailbox; downstream emits values in mailbox arrival order while
+    # preserving each input's own order.
+    class Merge
+      SIDE_ORDER = [:left, :right].freeze
+      CancellationError = Class.new(StandardError)
+      ValueMessage = Data.define(:side, :value)
+      DoneMessage = Data.define(:side)
+      ErrorMessage = Data.define(:side, :error)
+      private_constant :ValueMessage, :DoneMessage, :ErrorMessage
+
+      def initialize(left_materializer, right_materializer)
+        @materializers = { left: left_materializer, right: right_materializer }
+        @streams = { left: nil, right: nil }
+        @stream_closed = { left: false, right: false }
+        @side_done = { left: false, right: false }
+        @producers = {}
+        @mailbox = nil
+        @started = false
+        @closed = false
+        @done = false
+      end
+
+      def next
+        return DONE if @closed || @done
+
+        start
+        next_message
+      end
+
+      def close
+        return if @closed
+
+        @closed = true
+        @done = true
+        close_error = close_materialized_streams
+        close_mailbox
+        raise close_error if close_error
+      ensure
+        cancel_producers
+      end
+
+      private
+
+      def start
+        return if @started
+
+        validate_scheduler!
+
+        @mailbox = MergeMailbox.new(1)
+        @started = true
+        SIDE_ORDER.each do |side|
+          @producers[side] = Fiber.schedule { run_producer(side) }
+        end
+      end
+
+      def next_message
+        loop do
+          message = @mailbox.pop
+          return complete if message.nil?
+
+          case message
+          in ValueMessage[value:]
+            return value
+          in DoneMessage[side:]
+            mark_side_done(side)
+            return complete if all_done?
+          in ErrorMessage[error:]
+            return fail_with(error)
+          end
+        end
+      rescue MergeMailbox::Closed
+        complete
+      end
+
+      def run_producer(side)
+        stream = materialize_side(side)
+
+        loop do
+          break if @closed
+
+          message = pull_message(side, stream)
+          break unless deliver(message)
+          break unless message.is_a?(ValueMessage)
+        end
+      rescue MergeMailbox::Closed, CancellationError
+        nil
+      rescue StandardError => error
+        close_side(side, record_error: false)
+        deliver(ErrorMessage.new(side:, error:)) unless @closed
+      end
+
+      def materialize_side(side)
+        stream = @materializers.fetch(side).call
+        @streams[side] = stream
+        close_side(side) if @closed
+        stream
+      end
+
+      def pull_message(side, stream)
+        value = stream.next
+        return terminal_done_message(side) if Pull.done?(value)
+
+        ValueMessage.new(side:, value:)
+      rescue StandardError => error
+        close_side(side, record_error: false)
+        ErrorMessage.new(side:, error:)
+      end
+
+      def terminal_done_message(side)
+        close_error = close_side(side)
+        close_error ? ErrorMessage.new(side:, error: close_error) : DoneMessage.new(side:)
+      end
+
+      def deliver(message)
+        @mailbox.push(message)
+        true
+      rescue MergeMailbox::Closed
+        false
+      end
+
+      def mark_side_done(side)
+        @side_done[side] = true
+      end
+
+      def all_done?
+        SIDE_ORDER.all? { |side| @side_done.fetch(side) }
+      end
+
+      def complete
+        @done = true
+        close_mailbox
+        DONE
+      end
+
+      def fail_with(error)
+        @done = true
+        close_mailbox
+        close_materialized_streams
+        cancel_producers
+        raise error
+      end
+
+      def close_materialized_streams
+        first_error = nil
+
+        SIDE_ORDER.each do |side|
+          close_error = close_side(side)
+          first_error ||= close_error
+        end
+
+        first_error
+      end
+
+      def close_side(side, record_error: true)
+        return nil if @stream_closed.fetch(side)
+
+        stream = @streams[side]
+        return nil unless stream
+
+        @stream_closed[side] = true
+        @streams[side] = nil
+        stream.close
+        nil
+      rescue StandardError => error
+        error if record_error
+      end
+
+      def close_mailbox
+        @mailbox&.close
+      end
+
+      def cancel_producers
+        scheduler = Fiber.scheduler
+        return unless scheduler.respond_to?(:fiber_interrupt)
+
+        @producers.each_value 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
+            "Source.merge requires a non-blocking fiber"
+          else
+            "Source.merge requires Fiber.scheduler"
+          end
+        raise SchedulerRequiredError, message
+      end
+
+      class MergeMailbox
+        Closed = Class.new(StandardError)
+
+        def initialize(capacity)
+          @queue = Thread::SizedQueue.new(capacity)
+        end
+
+        def push(message)
+          @queue << message
+        rescue ClosedQueueError
+          raise Closed
+        end
+
+        def pop
+          @queue.pop
+        rescue ClosedQueueError
+          raise Closed
+        end
+
+        def close
+          @queue.close
+        end
+      end
+    end
+  end
+end

data/lib/fiber_stream/pull/parallel_map_boundary.rb

@@ -11,6 +11,11 @@ module FiberStream
     class ParallelMapBoundary
       TERMINAL_RESULT_CAPACITY = 1
       CancellationError = Class.new(StandardError)
+      JobMessage = Data.define(:sequence, :value)
+      ValueMessage = Data.define(:sequence, :value)
+      DoneMessage = Data.define(:sequence)
+      ErrorMessage = Data.define(:sequence, :error)
+      private_constant :JobMessage, :ValueMessage, :DoneMessage, :ErrorMessage
 
       def initialize(upstream, concurrency, transform)
         @upstream = upstream
@@ -85,28 +90,23 @@ module FiberStream
       end
 
       def emit(message)
-        case message.fetch(0)
-        when :value
-          emit_value(message)
-        when :done
+        case message
+        in ValueMessage[sequence:, value:]
+          emit_value(sequence, value)
+        in DoneMessage
           complete
-        when :error
-          fail_with_ordered_error(message)
+        in ErrorMessage[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
         return_permit unless @admission_closed
         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
           return next_message
@@ -130,7 +130,7 @@ module FiberStream
           break unless take_permit
 
           message = pull_job_message
-          if message.fetch(0) == :job
+          if message.is_a?(JobMessage)
             break unless deliver_job(message)
           else
             close_admission(close_upstream: false)
@@ -151,15 +151,19 @@ module FiberStream
 
         sequence = @next_sequence
         @next_sequence += 1
-        [:job, sequence, value]
+        JobMessage.new(sequence:, value:)
       rescue StandardError => error
         close_upstream(record_error: false)
-        [:error, @next_sequence, error]
+        ErrorMessage.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
+          ErrorMessage.new(sequence: @next_sequence, error: close_error)
+        else
+          DoneMessage.new(sequence: @next_sequence)
+        end
       end
 
       def run_worker
@@ -176,23 +180,23 @@ module FiberStream
       end
 
       def map_job(message)
-        sequence = message.fetch(1)
-        value = message.fetch(2)
-        [:value, sequence, @transform.call(value)]
+        sequence = message.sequence
+        value = message.value
+        ValueMessage.new(sequence:, value: @transform.call(value))
       rescue CancellationError
         raise
       rescue StandardError => error
-        [:error, sequence, error]
+        ErrorMessage.new(sequence:, error:)
       end
 
       def record_result(message)
-        if message.fetch(0) == :error
-          sequence = message.fetch(1)
+        if message.is_a?(ErrorMessage)
+          sequence = message.sequence
           @failure_sequence = sequence if @failure_sequence.nil? || sequence < @failure_sequence
           close_admission
         end
 
-        @pending[message.fetch(1)] = message
+        @pending[message.sequence] = message
       end
 
       def drain_available_results