From 2b73e6ba712d35e6ec767bf722edf542cc2e47c6 Mon Sep 17 00:00:00 2001
From: Benoit Daloze <eregontp@gmail.com>
Date: Mon, 21 Sep 2020 16:04:20 +0200
Subject: [PATCH] Simplify the implementation of Scheduler#block

* This shows block() with a timeout is similar to #kernel_sleep and also
  does not need to change `@blocking`.
---
 test/fiber/scheduler.rb | 28 +++++++++++++++-------------
 1 file changed, 15 insertions(+), 13 deletions(-)

diff --git a/test/fiber/scheduler.rb b/test/fiber/scheduler.rb
index d93d0f106d..c685e6ffc0 100644
--- a/test/fiber/scheduler.rb
+++ b/test/fiber/scheduler.rb
@@ -102,8 +102,8 @@ class Scheduler
     self.run
   ensure
     @closed = true
-    
-    # We freeze to detect any inadvertant modifications after the scheduler is closed:
+
+    # We freeze to detect any unintended modifications after the scheduler is closed:
     self.freeze
   end
 
@@ -144,19 +144,21 @@ class Scheduler
   # Used when blocking on synchronization (Mutex#lock, Queue#pop, SizedQueue#push, ...)
   def block(blocker, timeout = nil)
     # p [__method__, blocker, timeout]
-    @blocking += 1
-
     if timeout
       @waiting[Fiber.current] = current_time + timeout
-    end
-
-    Fiber.yield
-  ensure
-    @blocking -= 1
-
-    # Remove from @waiting in the case #unblock was called before the timeout expired:
-    if timeout
-      @waiting.delete(Fiber.current)
+      begin
+        Fiber.yield
+      ensure
+        # Remove from @waiting in the case #unblock was called before the timeout expired:
+        @waiting.delete(Fiber.current)
+      end
+    else
+      @blocking += 1
+      begin
+        Fiber.yield
+      ensure
+        @blocking -= 1
+      end
     end
   end