AsyncDataloader: avoid ClosedError race condition with straggler fibers

lib/graphql/dataloader/async_dataloader.rb

@@ -36,10 +36,14 @@ def yield(source = Fiber[:__graphql_current_dataloader_source])
         run = task.graphql_async_dataloader_run
         trace = run.trace
         trace&.dataloader_fiber_yield(source)
-        run.tasks_channel.push([:paused_task, task])
+        if !run.push_task_message(:paused_task, task)
+          task.stop
+        end
         condition = task.graphql_async_dataloader_condition
         condition.wait
-        run.tasks_channel.push([:resumed_task, task])
+        if !run.push_task_message(:resumed_task, task)
+          task.stop
+        end
         trace&.dataloader_fiber_resume(source)
         nil
       end
@@ -69,7 +73,7 @@ def initialize(dataloader, total_fiber_limit, jobs_fiber_limit)
 
         attr_accessor :trace, :root_task
 
-        attr_reader :jobs, :lazies_at_depth, :jobs_fiber_limit,  :snoozed_jobs_condition, :snoozed_sources_condition, :tasks_channel
+        attr_reader :jobs, :lazies_at_depth, :jobs_fiber_limit,  :snoozed_jobs_condition, :snoozed_sources_condition
 
         def jobs_bandwidth?
           running_count < @jobs_fiber_limit
@@ -84,6 +88,20 @@ def close_queues
           @tasks_channel_task.cancel
         end
 
+        # Push to the tasks_channel, tolerating a closed channel: on the error path, `run_queue`
+        # closes the channel while sibling tasks can still run one more slice before
+        # `root_task.cancel` reaches them. Record `:task_error` payloads so they aren't lost, and
+        # return false so the caller can stop the task instead of raising `ClosedError` into user code.
+        def push_task_message(msg, data)
+          @tasks_channel.push([msg, data])
+          true
+        rescue Async::Queue::ClosedError
+          if msg == :task_error
+            @task_error ||= data
+          end
+          false
+        end
+
         def wait_for_activity
           @activity.wait
         end
@@ -336,14 +354,14 @@ def spawn_tasks(run, mode, condition, pending_work, num_tasks)
             end
             nil
           rescue StandardError => err
-            run.tasks_channel.push([:task_error, err])
+            run.push_task_message(:task_error, err)
           else
-            run.tasks_channel.push([:finished_task, task])
+            run.push_task_message(:finished_task, task)
           ensure
             cleanup_fiber
             trace&.dataloader_fiber_exit
           end
-          run.tasks_channel.push([:started_task, new_task])
+          run.push_task_message(:started_task, new_task)
         end
       end
     end

spec/graphql/dataloader/async_dataloader_spec.rb

@@ -2,6 +2,7 @@
 require "spec_helper"
 if RUBY_VERSION >= "3.2.0"
   require "async"
+  require "timeout"
   describe GraphQL::Dataloader::AsyncDataloader do
     class AsyncSchema < GraphQL::Schema
       class SleepSource < GraphQL::Dataloader::Source
@@ -549,5 +550,48 @@ def fetch(keys)
         assert watchdog.join
       end
     end
+
+    describe "when a job errors while sibling tasks are parked in non-dataloader IO" do
+      # When a job errors, `run_queue` closes the tasks_channel while sibling tasks are still
+      # parked in non-Dataloader IO; a straggler's next push used to leak
+      # `Async::Queue::ClosedError` into resolver code and lose its own error report.
+      it "doesn't leak Async::Queue::ClosedError into straggler fibers or lose the original error" do
+        rng = Random.new(20260814)
+        leaked_error = nil
+
+        250.times do |i|
+          break if leaked_error
+
+          dataloader = GraphQL::Dataloader::AsyncDataloader.new
+
+          3.times do |j|
+            io_delay = rng.rand(0.003)
+            fetch_delay = rng.rand(0.002)
+            dataloader.append_job do
+              sleep(io_delay)
+              begin
+                dataloader.with(SlowSource, fetch_delay).load([i, j])
+              rescue Async::Queue::ClosedError => err
+                leaked_error = err
+                raise
+              end
+            end
+          end
+
+          err_delay = rng.rand(0.003)
+          dataloader.append_job do
+            sleep(err_delay)
+            raise "boom-#{i}"
+          end
+
+          err = assert_raises(RuntimeError) do
+            Timeout.timeout(15) { dataloader.run }
+          end
+          assert_equal "boom-#{i}", err.message
+        end
+
+        assert_nil leaked_error, "Async::Queue::ClosedError leaked into straggler fibers"
+      end
+    end
   end
 end