Merge pull request #3432 from djspiewak/bug/timer-firing

Fixed issues in the timer handling state machine integration

Author: djspiewak

core/jvm/src/main/scala/cats/effect/unsafe/IORuntimeBuilderPlatform.scala

@@ -18,14 +18,20 @@ package cats.effect.unsafe
 
 private[unsafe] abstract class IORuntimeBuilderPlatform { self: IORuntimeBuilder =>
 
+  // TODO unify this with the defaults in IORuntime.global and IOApp
   protected def platformSpecificBuild: IORuntime = {
     val (compute, computeShutdown) =
       customCompute.getOrElse(
         IORuntime.createWorkStealingComputeThreadPool(reportFailure = failureReporter))
+    val xformedCompute = computeTransform(compute)
+
+    val (scheduler, schedulerShutdown) = xformedCompute match {
+      case sched: Scheduler => customScheduler.getOrElse((sched, () => ()))
+      case _ => customScheduler.getOrElse(IORuntime.createDefaultScheduler())
+    }
+
     val (blocking, blockingShutdown) =
       customBlocking.getOrElse(IORuntime.createDefaultBlockingExecutionContext())
-    val (scheduler, schedulerShutdown) =
-      customScheduler.getOrElse(IORuntime.createDefaultScheduler())
     val shutdown = () => {
       computeShutdown()
       blockingShutdown()

core/jvm/src/main/scala/cats/effect/unsafe/WorkerThread.scala

@@ -327,21 +327,41 @@ private final class WorkerThread(
       }
     }
 
-    def parkUntilNextSleeper(): Unit = {
-      if (!isInterrupted()) {
+    // returns true if timed out, false if unparked
+    def parkUntilNextSleeper(): Boolean = {
+      while (!done.get()) {
         val now = System.nanoTime()
         val head = sleepersQueue.head()
         val nanos = head.triggerTime - now
+
+        // System.err.println(s"parking for $nanos nanos")
         LockSupport.parkNanos(pool, nanos)
 
-        if (parked.getAndSet(false)) {
-          pool.doneSleeping()
+        if (isInterrupted()) {
+          pool.shutdown()
+        } else {
+          if (parked.get()) {
+            // we were either awakened spuriously, or we timed out
+
+            if (head.triggerTime - System.nanoTime() <= 0) {
+              // we timed out
+              if (parked.getAndSet(false)) {
+                pool.doneSleeping()
+              }
+
+              return true
+            }
+          } else {
+            // we were awakened
+            return false
+          }
         }
       }
+
+      false
     }
 
     while (!done.get()) {
-
       if (blocking) {
         // The worker thread was blocked before. It is no longer part of the
         // core pool and needs to be cached.
@@ -394,25 +414,6 @@ private final class WorkerThread(
 
       val sleepers = sleepersQueue
 
-      if (sleepers.nonEmpty) {
-        val now = System.nanoTime()
-
-        var cont = true
-        while (cont) {
-          val head = sleepers.head()
-
-          if (head.triggerTime - now <= 0) {
-            if (head.get()) {
-              head.callback(RightUnit)
-            }
-            sleepers.popHead()
-            cont = sleepers.nonEmpty
-          } else {
-            cont = false
-          }
-        }
-      }
-
       ((state & ExternalQueueTicksMask): @switch) match {
         case 0 =>
           if (pool.blockedThreadDetectionEnabled) {
@@ -538,13 +539,20 @@ private final class WorkerThread(
               // Park the thread.
               if (sleepers.isEmpty) {
                 parkLoop()
+
+                // After the worker thread has been unparked, look for work in the
+                // external queue.
+                state = 3
               } else {
-                parkUntilNextSleeper()
+                if (parkUntilNextSleeper()) {
+                  // we made it to the end of our sleeping, so go straight to local queue stuff
+                  pool.transitionWorkerFromSearching(rnd)
+                  state = 4
+                } else {
+                  // we were interrupted, look for more work in the external queue
+                  state = 3
+                }
               }
-
-              // After the worker thread has been unparked, look for work in the
-              // external queue.
-              state = 3
             }
           }
 
@@ -584,13 +592,20 @@ private final class WorkerThread(
             // Park the thread.
             if (sleepers.isEmpty) {
               parkLoop()
+
+              // After the worker thread has been unparked, look for work in the
+              // external queue.
+              state = 3
             } else {
-              parkUntilNextSleeper()
+              if (parkUntilNextSleeper()) {
+                // we made it to the end of our sleeping, so go straight to local queue stuff
+                pool.transitionWorkerFromSearching(rnd)
+                state = 4
+              } else {
+                // we were interrupted, look for more work in the external queue
+                state = 3
+              }
             }
-
-            // After the worker thread has been unparked, look for work in the
-            // external queue.
-            state = 3
           }
 
         case 3 =>
@@ -647,6 +662,26 @@ private final class WorkerThread(
           }
 
         case _ =>
+          if (sleepers.nonEmpty) {
+            val now = System.nanoTime()
+
+            var cont = true
+            while (cont) {
+              val head = sleepers.head()
+
+              if (head.triggerTime - now <= 0) {
+                if (head.get()) {
+                  // System.err.println(s"dequeued with state = $state; sleepers = $sleepers")
+                  head.callback(RightUnit)
+                }
+                sleepers.popHead()
+                cont = sleepers.nonEmpty
+              } else {
+                cont = false
+              }
+            }
+          }
+
           // Check the queue bypass reference before dequeueing from the local
           // queue.
           val fiber = if (cedeBypass eq null) {

tests/jvm/src/test/scala/cats/effect/IOPlatformSpecification.scala

@@ -340,6 +340,36 @@ trait IOPlatformSpecification { self: BaseSpec with ScalaCheck =>
           runtime.shutdown()
         }
       }
+
+      // this test ensures that the parkUntilNextSleeper bit works
+      "run a timer when parking thread" in {
+        val (pool, shutdown) = IORuntime.createWorkStealingComputeThreadPool(threads = 1)
+
+        implicit val runtime: IORuntime = IORuntime.builder().setCompute(pool, shutdown).build()
+
+        try {
+          // longer sleep all-but guarantees this timer is fired *after* the worker is parked
+          val test = IO.sleep(500.millis) *> IO.pure(true)
+          test.unsafeRunTimed(5.seconds) must beSome(true)
+        } finally {
+          runtime.shutdown()
+        }
+      }
+
+      // this test ensures that we always see the timer, even when it fires just as we're about to park
+      "run a timer when detecting just prior to park" in {
+        val (pool, shutdown) = IORuntime.createWorkStealingComputeThreadPool(threads = 1)
+
+        implicit val runtime: IORuntime = IORuntime.builder().setCompute(pool, shutdown).build()
+
+        try {
+          // shorter sleep makes it more likely this timer fires *before* the worker is parked
+          val test = IO.sleep(1.milli) *> IO.pure(true)
+          test.unsafeRunTimed(1.second) must beSome(true)
+        } finally {
+          runtime.shutdown()
+        }
+      }
     }
   }
 }