From dfab06808d550004c8fc6821291c9f9ffbff88b7 Mon Sep 17 00:00:00 2001
From: Sarun Rattanasiri <midnight_w@gmx.tw>
Date: Thu, 13 Feb 2025 21:42:36 +0700
Subject: [PATCH] fix deadlock in dispatch_loop

---
 lib/bootsnap/cli/worker_pool.rb | 37 +++++++++++++++------------------
 1 file changed, 17 insertions(+), 20 deletions(-)

diff --git a/lib/bootsnap/cli/worker_pool.rb b/lib/bootsnap/cli/worker_pool.rb
index a4d40bad..8f2253b7 100644
--- a/lib/bootsnap/cli/worker_pool.rb
+++ b/lib/bootsnap/cli/worker_pool.rb
@@ -45,14 +45,9 @@ def initialize(jobs)
           @pid = nil
         end
 
-        def write(message, block: true)
+        def write(message)
           payload = Marshal.dump(message)
-          if block
-            to_io.write(payload)
-            true
-          else
-            to_io.write_nonblock(payload, exception: false) != :wait_writable
-          end
+          to_io.write_nonblock(payload)
         end
 
         def close
@@ -74,7 +69,7 @@ def spawn
           @pid = Process.fork do
             to_io.close
             work_loop
-            exit!(0)
+            exit!(true)
           end
           @pipe_out.close
           true
@@ -84,7 +79,7 @@ def spawn
       def initialize(size:, jobs: {})
         @size = size
         @jobs = jobs
-        @queue = Queue.new
+        @queue = Thread::Queue.new
         @pids = []
       end
 
@@ -98,25 +93,27 @@ def spawn
 
       def dispatch_loop
         loop do
-          case job = @queue.pop
-          when nil
+          job = @queue.pop
+          if job
+            IO.select(nil, @workers).tap do |(_nil, available)|
+              available.sample.write(job)
+            end
+          else
+            closed = []
             @workers.each do |worker|
               worker.write([:exit])
               worker.close
+              closed << worker
+            rescue IO::WaitWritable
+              next
             end
-            return true
-          else
-            unless @workers.sample.write(job, block: false)
-              free_worker.write(job)
-            end
+            @workers.delete_if(&closed.method(:include?))
+            return if @workers.empty?
+            IO.select(nil, @workers)
           end
         end
       end
 
-      def free_worker
-        IO.select(nil, @workers)[1].sample
-      end
-
       def push(*args)
         @queue.push(args)
         nil