Merge pull request #1455 from wemrysi/feature/throw-on-main-canceled

Raise exception when IOApp main fiber is canceled
typelevel · Nov 27, 2020 · 0edddb5 · 0edddb5
2 parents 2eb654f + 9815784
commit 0edddb5
Show file tree

Hide file tree

Showing 3 changed files with 37 additions and 12 deletions.
diff --git a/core/js/src/main/scala/cats/effect/IOApp.scala b/core/js/src/main/scala/cats/effect/IOApp.scala
@@ -42,11 +42,19 @@ trait IOApp {
       else
         args.toList
 
-    IO.race(run(argList), keepAlive)
+    Spawn[IO]
+      .raceOutcome[ExitCode, Nothing](run(argList), keepAlive)
+      .flatMap {
+        case Left(Outcome.Canceled()) =>
+          IO.raiseError(new RuntimeException("IOApp main fiber canceled"))
+        case Left(Outcome.Errored(t)) => IO.raiseError(t)
+        case Left(Outcome.Succeeded(code)) => code
+        case Right(Outcome.Errored(t)) => IO.raiseError(t)
+        case Right(_) => sys.error("impossible")
+      }
       .unsafeRunAsync({
         case Left(t) => throw t
-        case Right(Left(code)) => reportExitCode(code)
-        case Right(Right(_)) => sys.error("impossible")
+        case Right(code) => reportExitCode(code)
       })(unsafe.IORuntime.global)
   }
 

diff --git a/core/jvm/src/main/scala/cats/effect/IOApp.scala b/core/jvm/src/main/scala/cats/effect/IOApp.scala
@@ -34,15 +34,21 @@ trait IOApp {
 
     val ioa = run(args.toList)
 
-    val fiber = ioa.unsafeRunFiber(
-      { t =>
-        error = t
-        latch.countDown()
-      },
-      { a =>
-        result = a
-        latch.countDown()
-      })(runtime)
+    val fiber =
+      ioa
+        .onCancel(IO {
+          error = new RuntimeException("IOApp main fiber canceled")
+          latch.countDown()
+        })
+        .unsafeRunFiber(
+          { t =>
+            error = t
+            latch.countDown()
+          },
+          { a =>
+            result = a
+            latch.countDown()
+          })(runtime)
 
     def handleShutdown(): Unit = {
       if (latch.getCount() > 0) {

diff --git a/core/jvm/src/test/scala/cats/effect/IOAppSpec.scala b/core/jvm/src/test/scala/cats/effect/IOAppSpec.scala
@@ -108,6 +108,12 @@ class IOAppSpec extends Specification {
         h.awaitStatus() mustEqual 1
         h.stderr() must contain("Boom!")
       }
+
+      "exit on canceled" in {
+        val h = java(Canceled, List.empty)
+        h.awaitStatus() mustEqual 1
+        h.stderr() must contain("canceled")
+      }
     }
   }
 
@@ -179,4 +185,9 @@ package examples {
         _ <- IO.never[Unit]
       } yield ExitCode.Success
   }
+
+  object Canceled extends IOApp {
+    def run(args: List[String]): IO[ExitCode] =
+      IO.canceled.as(ExitCode.Success)
+  }
 }