fix IOs.ensure error handling

getkyo · Oct 5, 2023 · de1b7df · de1b7df
1 parent 493122a
commit de1b7df
Show file tree

Hide file tree

Showing 2 changed files with 34 additions and 3 deletions.
diff --git a/kyo-core/shared/src/main/scala/kyo/ios.scala b/kyo-core/shared/src/main/scala/kyo/ios.scala
@@ -104,7 +104,7 @@ object ios {
       runLazyLoop(v)
     }
 
-    private[kyo] def ensure[T, S](f: => Unit > IOs)(v: => T > S): T > (IOs with S) = {
+    private[kyo] def ensure[T, S](f: => Unit > IOs)(v: T > S): T > (IOs with S) = {
       lazy val run: Unit =
         try IOs.run(f)
         catch {
@@ -128,7 +128,14 @@ object ios {
                     case _ =>
                       p
                   }
-                ensureLoop(kyo(v, s, l), np)
+                val v2 =
+                  try kyo(v, s, l)
+                  catch {
+                    case ex if (NonFatal(ex)) =>
+                      run
+                      throw ex
+                  }
+                ensureLoop(v2, np)
               }
             }
           case _ =>

diff --git a/kyo-core/shared/src/test/scala/kyoTest/iosTest.scala b/kyo-core/shared/src/test/scala/kyoTest/iosTest.scala
@@ -1,5 +1,6 @@
-package kyoTest
+package kyo
 
+import kyoTest.KyoTest
 import kyo.concurrent.atomics.AtomicInt
 import kyo._
 import kyo.ios._
@@ -12,6 +13,7 @@ import scala.concurrent.duration._
 import scala.util.Success
 import scala.util.Try
 import org.scalatest.compatible.Assertion
+import scala.util.Failure
 
 class iosTest extends KyoTest {
 
@@ -142,4 +144,26 @@ class iosTest extends KyoTest {
       )
     }
   }
+
+  "ensure" - {
+    "success" in {
+      var called = false
+      checkEquals[Try[Int], Any](
+          Tries.run(IOs.run(IOs.ensure[Int, Any] { called = true }(1))),
+          Try(1)
+      )
+      assert(called)
+    }
+    "failure" in {
+      val ex     = new Exception
+      var called = false
+      checkEquals[Try[Int], Any](
+          Tries.run(IOs.run(IOs.ensure { called = true } {
+            IOs[Int, Any](throw ex)
+          })),
+          Failure(ex)
+      )
+      assert(called)
+    }
+  }
 }