fiber_stream 0.2.0 → 0.4.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
@@ -67,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,
@@ -103,7 +115,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 +132,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 +159,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 +177,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 +217,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 +292,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 +349,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 +385,7 @@ module FiberStream
             cause_message: cause_message
           )
 
-        [:error, sequence, error]
+        ResultError.new(sequence:, error:)
       end
 
       def worker_for_id(worker_id)
@@ -411,7 +417,7 @@ module FiberStream
 
         @worker_shutdown_sent = true
         @workers.each do |worker|
-          worker.send([:shutdown])
+          worker.send(Shutdown.new)
         rescue StandardError
           nil
         end
@@ -420,7 +426,6 @@ module FiberStream
       def wait_for_workers
         return unless @coordinator
 
-        sleep READY_WAIT_INTERVAL while @coordinator.alive?
         @coordinator.join
       end
 
@@ -453,48 +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
-
-          begin
-            port.send([:ready, id])
+      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
 
-            loop do
-              message = Ractor.receive
-              break if message.fetch(0) == :shutdown
+            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
 
-              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
-                begin
-                  if transfer == :move
-                    port.send([:value, id, current_sequence, mapped_value], move: true)
+                  begin
+                    mapped_value = mapper.call(value)
+                  rescue Exception => error # rubocop:disable Lint/RescueException
+                    break unless send_failure.call(current_sequence, :worker, error)
                   else
-                    port.send([:value, id, current_sequence, mapped_value])
+                    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
-                rescue Exception => error # rubocop:disable Lint/RescueException
-                  port.send([:error, id, current_sequence, :output_transfer, error.class.name, error.message])
+
+                  current_sequence = nil
+                  break unless send_control.call(Ready.new(id))
                 end
               end
-
-              current_sequence = nil
-              port.send([:ready, id])
+            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
-            port.send([:error, id, sequence, :worker_termination, error.class.name, error.message])
-          ensure
-            port.send([:stopped, id])
           end
         end
       end
-
-      private_class_method :spawn_worker
     end
   end
 end