Merge pull request #3510 from samspills/sam/3506/dispatcher-await

Dispatcher: check for outstanding actions before release

Author: djspiewak

std/shared/src/main/scala/cats/effect/std/Dispatcher.scala

@@ -242,19 +242,14 @@ object Dispatcher {
         states
       })
       ec <- Resource.eval(F.executionContext)
-      alive <- Resource.make(F.delay(new AtomicBoolean(true)))(ref => F.delay(ref.set(false)))
 
       // supervisor for the main loop, which needs to always restart unless the Supervisor itself is canceled
       // critically, inner actions can be canceled without impacting the loop itself
       supervisor <- Supervisor[F](await, Some((_: Outcome[F, Throwable, _]) => true))
 
       _ <- {
-        def dispatcher(
-            doneR: AtomicBoolean,
-            latch: AtomicReference[() => Unit],
-            state: Array[AtomicReference[List[Registration]]]): F[Unit] = {
-
-          val step = for {
+        def step(state: Array[AtomicReference[List[Registration]]], await: F[Unit]): F[Unit] =
+          for {
             regs <- F delay {
               val buffer = mutable.ListBuffer.empty[Registration]
               var i = 0
@@ -271,12 +266,7 @@ object Dispatcher {
 
             _ <-
               if (regs.isEmpty) {
-                F.async_[Unit] { cb =>
-                  if (!latch.compareAndSet(Noop, () => cb(Completed))) {
-                    // state was changed between when we last set the latch and now; complete the callback immediately
-                    cb(Completed)
-                  }
-                }
+                await
               } else {
                 regs traverse_ {
                   case r @ Registration(action, prepareCancel) =>
@@ -289,10 +279,23 @@ object Dispatcher {
               }
           } yield ()
 
+        def dispatcher(
+            doneR: AtomicBoolean,
+            latch: AtomicReference[() => Unit],
+            state: Array[AtomicReference[List[Registration]]]): F[Unit] = {
+
+          val await =
+            F.async_[Unit] { cb =>
+              if (!latch.compareAndSet(Noop, () => cb(Completed))) {
+                // state was changed between when we last set the latch and now; complete the callback immediately
+                cb(Completed)
+              }
+            }
+
           F.delay(latch.set(Noop)) *> // reset latch
             // if we're marked as done, yield immediately to give other fibers a chance to shut us down
             // we might loop on this a few times since we're marked as done before the supervisor is canceled
-            F.delay(doneR.get()).ifM(F.cede, step)
+            F.delay(doneR.get()).ifM(F.cede, step(state, await))
         }
 
         0.until(workers).toList traverse_ { n =>
@@ -301,12 +304,15 @@ object Dispatcher {
             val worker = dispatcher(doneR, latch, states(n))
             val release = F.delay(latch.getAndSet(Open)())
             Resource.make(supervisor.supervise(worker)) { _ =>
-              // published by release
-              F.delay(doneR.lazySet(true)) *> release
+              F.delay(doneR.set(true)) *> step(states(n), F.unit) *> release
             }
           }
         }
       }
+
+      // Alive is the innermost resource so that when releasing
+      // the very first thing we do is set dispatcher to un-alive
+      alive <- Resource.make(F.delay(new AtomicBoolean(true)))(ref => F.delay(ref.set(false)))
     } yield {
       new Dispatcher[F] {
         override def unsafeRunAndForget[A](fa: F[A]): Unit = {

tests/shared/src/test/scala/cats/effect/std/DispatcherSpec.scala

@@ -321,6 +321,33 @@ class DispatcherSpec extends BaseSpec with DetectPlatform {
         } yield ok
       }
     }
+
+    "reject new tasks while shutting down" in real {
+      (IO.ref(false), IO.ref(false)).flatMapN { (resultR, rogueResultR) =>
+        dispatcher
+          .allocated
+          .flatMap {
+            case (runner, release) =>
+              IO(runner.unsafeRunAndForget(
+                IO.sleep(1.second).uncancelable.guarantee(resultR.set(true)))) *>
+                IO.sleep(100.millis) *>
+                release.both(
+                  IO.sleep(500.nanos) *>
+                    IO(runner.unsafeRunAndForget(rogueResultR.set(true))).attempt
+                )
+          }
+          .flatMap {
+            case (_, rogueSubmitResult) =>
+              for {
+                result <- resultR.get
+                rogueResult <- rogueResultR.get
+                _ <- IO(result must beTrue)
+                _ <- IO(rogueResult must beFalse)
+                _ <- IO(rogueSubmitResult must beLeft)
+              } yield ok
+          }
+      }
+    }
   }
 
   private def awaitTermination(dispatcher: Resource[IO, Dispatcher[IO]]) = {
@@ -353,6 +380,15 @@ class DispatcherSpec extends BaseSpec with DetectPlatform {
       }
     }
 
+    "issue #3506: await unsafeRunAndForget" in ticked { implicit ticker =>
+      val result = for {
+        resultR <- IO.ref(false)
+        _ <- dispatcher.use { runner => IO(runner.unsafeRunAndForget(resultR.set(true))) }
+        result <- resultR.get
+      } yield result
+      result must completeAs(true)
+    }
+
     "cancel active fibers when an error is produced" in real {
       case object TestException extends RuntimeException