Replace manual Fibers with Async gem primitives for component streaming

react_on_rails_pro/app/helpers/react_on_rails_pro_helper.rb

@@ -128,7 +128,11 @@ def stream_react_component(component_name, options = {})
     # Because setting prerender to false is equivalent to calling react_component with prerender: false
     options[:prerender] = true
     options = options.merge(immediate_hydration: true) unless options.key?(:immediate_hydration)
-    run_stream_inside_fiber do
+
+    # Extract streaming-specific callback
+    on_complete = options.delete(:on_complete)
+
+    consumer_stream_async(on_complete: on_complete) do
       internal_stream_react_component(component_name, options)
     end
   end
@@ -185,7 +189,11 @@ def rsc_payload_react_component(component_name, options = {})
     # rsc_payload_react_component doesn't have the prerender option
     # Because setting prerender to false will not do anything
     options[:prerender] = true
-    run_stream_inside_fiber do
+
+    # Extract streaming-specific callback
+    on_complete = options.delete(:on_complete)
+
+    consumer_stream_async(on_complete: on_complete) do
       internal_rsc_payload_react_component(component_name, options)
     end
   end
@@ -246,30 +254,29 @@ def handle_stream_cache_hit(component_name, raw_options, auto_load_bundle, cache
     load_pack_for_generated_component(component_name, render_options)
 
     initial_result, *rest_chunks = cached_chunks
-    hit_fiber = Fiber.new do
-      rest_chunks.each { |chunk| Fiber.yield(chunk) }
-      nil
+
+    # Enqueue remaining chunks asynchronously
+    @async_barrier.async do
+      rest_chunks.each { |chunk| @main_output_queue.enqueue(chunk) }
     end
-    @rorp_rendering_fibers << hit_fiber
+
+    # Return first chunk directly
     initial_result
   end
 
   def handle_stream_cache_miss(component_name, raw_options, auto_load_bundle, view_cache_key, &block)
-    # Kick off the normal streaming helper to get the initial result and the original fiber
-    initial_result = render_stream_component_with_props(component_name, raw_options, auto_load_bundle, &block)
-    original_fiber = @rorp_rendering_fibers.pop
-
-    buffered_chunks = [initial_result]
-    wrapper_fiber = Fiber.new do
-      while (chunk = original_fiber.resume)
-        buffered_chunks << chunk
-        Fiber.yield(chunk)
-      end
-      Rails.cache.write(view_cache_key, buffered_chunks, raw_options[:cache_options] || {})
-      nil
-    end
-    @rorp_rendering_fibers << wrapper_fiber
-    initial_result
+    cache_aware_options = raw_options.merge(
+      on_complete: lambda { |chunks|
+        Rails.cache.write(view_cache_key, chunks, raw_options[:cache_options] || {})
+      }
+    )
+
+    render_stream_component_with_props(
+      component_name,
+      cache_aware_options,
+      auto_load_bundle,
+      &block
+    )
   end
 
   def render_stream_component_with_props(component_name, raw_options, auto_load_bundle)
@@ -291,25 +298,46 @@ def check_caching_options!(raw_options, block)
     raise ReactOnRailsPro::Error, "Option 'cache_key' is required for React on Rails caching"
   end
 
-  def run_stream_inside_fiber
-    if @rorp_rendering_fibers.nil?
+  def consumer_stream_async(on_complete:)
+    require "async/variable"
+
+    if @async_barrier.nil?
       raise ReactOnRails::Error,
             "You must call stream_view_containing_react_components to render the view containing the react component"
     end
 
-    rendering_fiber = Fiber.new do
+    # Create a variable to hold the first chunk for synchronous return
+    first_chunk_var = Async::Variable.new
+    all_chunks = [] if on_complete # Only collect if callback provided
+
+    # Start an async task on the barrier to stream all chunks
+    @async_barrier.async do
       stream = yield
+      is_first = true
+
       stream.each_chunk do |chunk|
-        Fiber.yield chunk
+        all_chunks << chunk if on_complete # Collect for callback
+
+        if is_first
+          # Store first chunk in variable for synchronous access
+          first_chunk_var.value = chunk
+          is_first = false
+        else
+          # Enqueue remaining chunks to main output queue
+          @main_output_queue.enqueue(chunk)
+        end
       end
-    end
 
-    @rorp_rendering_fibers << rendering_fiber
+      # Handle case where stream has no chunks
+      first_chunk_var.value = nil if is_first
+
+      # Call callback with all chunks when streaming completes
+      on_complete&.call(all_chunks)
+    end
 
-    # return the first chunk of the fiber
-    # It contains the initial html of the component
-    # all updates will be appended to the stream sent to browser
-    rendering_fiber.resume
+    # Wait for and return the first chunk (blocking)
+    first_chunk_var.wait
+    first_chunk_var.value
   end
 
   def internal_stream_react_component(component_name, options = {})

react_on_rails_pro/lib/react_on_rails_pro/concerns/stream.rb

@@ -31,73 +31,54 @@ module Stream
     #
     # @see ReactOnRails::Helper#stream_react_component
     def stream_view_containing_react_components(template:, close_stream_at_end: true, **render_options)
-      @rorp_rendering_fibers = []
-      template_string = render_to_string(template: template, **render_options)
-      # View may contain extra newlines, chunk already contains a newline
-      # Having multiple newlines between chunks causes hydration errors
-      # So we strip extra newlines from the template string and add a single newline
-      response.stream.write(template_string)
-
-      begin
-        drain_streams_concurrently
-      ensure
-        response.stream.close if close_stream_at_end
-      end
-    end
-
-    private
-
-    def drain_streams_concurrently
       require "async"
+      require "async/barrier"
       require "async/limited_queue"
 
-      return if @rorp_rendering_fibers.empty?
-
-      Sync do |parent|
-        # To avoid memory bloat, we use a limited queue to buffer chunks in memory.
+      Sync do |parent_task|
+        # Initialize async primitives for concurrent component streaming
+        @async_barrier = Async::Barrier.new
         buffer_size = ReactOnRailsPro.configuration.concurrent_component_streaming_buffer_size
-        queue = Async::LimitedQueue.new(buffer_size)
+        @main_output_queue = Async::LimitedQueue.new(buffer_size)
 
-        writer = build_writer_task(parent: parent, queue: queue)
-        tasks = build_producer_tasks(parent: parent, queue: queue)
+        # Render template - components will start streaming immediately
+        template_string = render_to_string(template: template, **render_options)
+        # View may contain extra newlines, chunk already contains a newline
+        # Having multiple newlines between chunks causes hydration errors
+        # So we strip extra newlines from the template string and add a single newline
+        response.stream.write(template_string)
 
-        # This structure ensures that even if a producer task fails, we always
-        # signal the writer to stop and then wait for it to finish draining
-        # any remaining items from the queue before propagating the error.
         begin
-          tasks.each(&:wait)
-        ensure
-          # `close` signals end-of-stream; when writer tries to dequeue, it will get nil, so it will exit.
-          queue.close
-          writer.wait
+          drain_streams_concurrently(parent_task)
+          # Do not close the response stream in an ensure block.
+          # If an error occurs we may need the stream open to send diagnostic/error details
+          # (for example, ApplicationController#rescue_from in the dummy app).
+          response.stream.close if close_stream_at_end
         end
       end
     end
 
-    def build_producer_tasks(parent:, queue:)
-      @rorp_rendering_fibers.each_with_index.map do |fiber, idx|
-        parent.async do
-          loop do
-            chunk = fiber.resume
-            break unless chunk
+    private
 
-            # Will be blocked if the queue is full until a chunk is dequeued
-            queue.enqueue([idx, chunk])
-          end
+    def drain_streams_concurrently(parent_task)
+      writing_task = parent_task.async do
+        # Drain all remaining chunks from the queue to the response stream
+        while (chunk = @main_output_queue.dequeue)
+          response.stream.write(chunk)
         end
       end
-    end
 
-    def build_writer_task(parent:, queue:)
-      parent.async do
-        loop do
-          pair = queue.dequeue
-          break if pair.nil?
-
-          _idx_from_queue, item = pair
-          response.stream.write(item)
-        end
+      # Wait for all component streaming tasks to complete
+      begin
+        @async_barrier.wait
+      rescue StandardError => e
+        @async_barrier.stop
+        raise e
       end
+    ensure
+      # Close the queue to signal end of streaming
+      @main_output_queue.close
+      writing_task.wait
     end
   end
 end