Retain reference to blocking fibers.

1 files changed, 8 insertions, 6 deletions

diff --git a/test/fiber/scheduler.rb b/test/fiber/scheduler.rb
index 4138015e4b..96b22856d1 100644
--- a/test/fiber/scheduler.rb
+++ b/test/fiber/scheduler.rb
@@ -30,7 +30,7 @@ class Scheduler
     @closed = false
 
     @lock = Thread::Mutex.new
-    @blocking = 0
+    @blocking = Hash.new.compare_by_identity
     @ready = []
 
     @urgent = IO.pipe
@@ -57,7 +57,7 @@ class Scheduler
   def run
     # $stderr.puts [__method__, Fiber.current].inspect
 
-    while @readable.any? or @writable.any? or @waiting.any? or @blocking.positive?
+    while @readable.any? or @writable.any? or @waiting.any? or @blocking.any?
       # Can only handle file descriptors up to 1024...
       readable, writable = IO.select(@readable.keys + [@urgent.first], @writable.keys, [], next_timeout)
 
@@ -211,20 +211,22 @@ class Scheduler
   def block(blocker, timeout = nil)
     # $stderr.puts [__method__, blocker, timeout].inspect
 
+    fiber = Fiber.current
+
     if timeout
-      @waiting[Fiber.current] = current_time + timeout
+      @waiting[fiber] = current_time + timeout
       begin
         Fiber.yield
       ensure
         # Remove from @waiting in the case #unblock was called before the timeout expired:
-        @waiting.delete(Fiber.current)
+        @waiting.delete(fiber)
       end
     else
-      @blocking += 1
+      @blocking[fiber] = true
       begin
         Fiber.yield
       ensure
-        @blocking -= 1
+        @blocking.delete(fiber)
       end
     end
   end