From d0f7d7448717e08aaa10edd11fc2d8bd92ccd5f7 Mon Sep 17 00:00:00 2001
From: Adam Fraser <adam.fraser@gmail.com>
Date: Fri, 2 Dec 2022 09:19:32 +0000
Subject: [PATCH] preserve stack track in zio#retry

---
 core/shared/src/main/scala/zio/ZIO.scala | 26 ++++++++++++++++++++++--
 1 file changed, 24 insertions(+), 2 deletions(-)

diff --git a/core/shared/src/main/scala/zio/ZIO.scala b/core/shared/src/main/scala/zio/ZIO.scala
index 0631fdb5e2b7..b864f96f7b77 100644
--- a/core/shared/src/main/scala/zio/ZIO.scala
+++ b/core/shared/src/main/scala/zio/ZIO.scala
@@ -1642,7 +1642,24 @@ sealed trait ZIO[-R, +E, +A]
   final def retry[R1 <: R, S](
     policy: => Schedule[R1, E, S]
   )(implicit ev: CanFail[E], trace: Trace): ZIO[R1, E, A] =
-    retryOrElse(policy, (e: E, _: S) => ZIO.fail(e))
+    ZIO.suspendSucceed {
+
+      def loop(driver: Schedule.Driver[Any, R1, E, S]): ZIO[R1, E, A] =
+        self.catchAllCause { cause =>
+          cause.failureOrCause.fold(
+            e =>
+              driver
+                .next(e)
+                .foldZIO(
+                  _ => driver.last.orDie.flatMap(_ => ZIO.refailCause(cause)),
+                  _ => loop(driver)
+                ),
+            cause => ZIO.refailCause(cause)
+          )
+        }
+
+      policy.driver.flatMap(loop(_))
+    }
 
   /**
    * Retries this effect the specified number of times.
@@ -1651,7 +1668,12 @@ sealed trait ZIO[-R, +E, +A]
     ZIO.suspendSucceed {
 
       def loop(n: Int): ZIO[R, E, A] =
-        self.catchAll(e => if (n <= 0) ZIO.fail(e) else ZIO.yieldNow *> loop(n - 1))
+        self.catchAllCause { cause =>
+          cause.failureOrCause.fold(
+            _ => if (n <= 0) ZIO.refailCause(cause) else ZIO.yieldNow *> loop(n - 1),
+            cause => ZIO.refailCause(cause)
+          )
+        }
 
       loop(n)
     }