zio · somdoron · Jan 21, 2025 · Jan 21, 2025 · Jan 22, 2025 · Jan 22, 2025
diff --git a/core-tests/shared/src/test/scala/zio/ZIOSpec.scala b/core-tests/shared/src/test/scala/zio/ZIOSpec.scala
@@ -307,6 +307,29 @@ object ZIOSpec extends ZIOBaseSpec {
         } yield assert(result)((equalTo(t)))
       }
     ) @@ zioTag(errors),
+    suite("catchAllFailure")(
+      test("recovers from all failures") {
+        val s   = "division by zero"
+        val zio = ZIO.fail(new IllegalArgumentException(s))
+        for {
+          result <- zio.catchFailureCause(e => ZIO.succeed(e.failureOption.get.getMessage))
+        } yield assert(result)(equalTo(s))
+      },
+      test("leaves defects") {
+        val t   = new IllegalArgumentException("division by zero")
+        val zio = ZIO.attempt(true) *> ZIO.die(t)
+        for {
+          exit <- zio.catchFailureCause(e => ZIO.succeed(e.failureOption.get.getMessage)).exit
+        } yield assert(exit)(dies(equalTo(t)))
+      },
+      test("leaves values") {
+        val t   = new IllegalArgumentException("division by zero")
+        val zio = ZIO.attempt(t)
+        for {
+          result <- zio.catchFailureCause(e => ZIO.succeed(e.failureOption.get.getMessage))
+        } yield assert(result)((equalTo(t)))
+      }
+    ) @@ zioTag(errors),
     suite("catchSomeCause")(
       test("catches matching cause") {
         ZIO.interrupt.catchSomeCause {
@@ -363,6 +386,38 @@ object ZIOSpec extends ZIOBaseSpec {
         } yield assert(result)((equalTo(t)))
       }
     ) @@ zioTag(errors),
+    suite("catchSomeFailure")(
+      test("catches matching cause") {
+        ZIO
+          .fail("fail")
+          .catchSomeFailureCause {
+            case c if c.failureOption.get == "fail" => ZIO.succeed(true)
+          }
+          .sandbox
+          .map(
+            assert(_)(isTrue)
+          )
+      },
+      test("fails if cause doesn't match") {
+        ZIO
+          .fail("no-match")
+          .catchSomeFailureCause {
+            case c if c.failureOption.get == "fail" => ZIO.succeed(true)
+          }
+          .sandbox
+          .either
+          .map(
+            assert(_)(isLeft(equalTo(Cause.fail("no-match"))))
+          )
+      },
+      test("doesn't catch defects") {
+        for {
+          result <- (ZIO.attempt(42) *> ZIO.dieMessage("die")).catchSomeFailureCause { case _ =>
+                      ZIO.succeed(true)
+                    }.exit
+        } yield assert(result)(dies(hasMessage(equalTo("die"))))
+      }
+    ) @@ zioTag(errors),
     suite("collect")(
       test("returns failure ignoring value") {
         for {
@@ -3866,6 +3921,24 @@ object ZIOSpec extends ZIOBaseSpec {
         } yield assert(effect)(equalTo(42))
       }
     ),
+    suite("tapFailure")(
+      test("doesn't peek at the defect of this effect") {
+        for {
+          ref    <- Ref.make(false)
+          result <- (ZIO.attempt(42) *> ZIO.dieMessage("die")).tapFailureCause(_ => ref.set(true)).exit
+          effect <- ref.get
+        } yield assert(result)(dies(hasMessage(equalTo("die")))) &&
+          assert(effect)(isFalse)
+      },
+      test("effectually peeks at the failure of this effect") {
+        for {
+          ref    <- Ref.make(false)
+          result <- ZIO.fail("fail").tapFailureCause(_ => ref.set(true)).exit
+          effect <- ref.get
+        } yield assert(result)(fails(equalTo("fail"))) &&
+          assert(effect)(isTrue)
+      }
+    ),
     suite("tapSome")(
       test("is identity if the function doesn't match") {
         for {

diff --git a/core/shared/src/main/scala/zio/Cause.scala b/core/shared/src/main/scala/zio/Cause.scala
@@ -364,6 +364,9 @@ sealed abstract class Cause[+E] extends Product with Serializable { self =>
   final def isFailure: Boolean =
     failureOption.isDefined
 
+  final def isFailureOnly: Boolean =
+    foldContext(())(Folder.IsFailureOnly)
+
   /**
    * Determines if the `Cause` contains an interruption.
    */
@@ -412,6 +415,28 @@ sealed abstract class Cause[+E] extends Product with Serializable { self =>
       (causeOption, stackless) => causeOption.map(Stackless(_, stackless))
     )
 
+  def keepFailures: Option[Cause[E]] =
+    foldLog[Option[Cause[E]]](
+      None,
+      (e, trace, spans, annotations) => Some(Fail(e, trace, spans, annotations)),
+      (_, _, _, _) => None,
+      (_, _, _, _) => None
+    )(
+      {
+        case (Some(l), Some(r)) => Some(Then(l, r))
+        case (Some(l), None)    => Some(l)
+        case (None, Some(r))    => Some(r)
+        case (None, None)       => None
+      },
+      {
+        case (Some(l), Some(r)) => Some(Both(l, r))
+        case (Some(l), None)    => Some(l)
+        case (None, Some(r))    => Some(r)
+        case (None, None)       => None
+      },
+      (causeOption, stackless) => causeOption.map(Stackless(_, stackless))
+    )
+
   /**
    * Linearizes this cause to a set of parallel causes where each parallel cause
    * contains a linear sequence of failures.
@@ -822,6 +847,17 @@ object Cause extends Serializable {
       def stacklessCase(context: Any, value: Boolean, stackless: Boolean): Boolean = value
     }
 
+    case object IsFailureOnly extends Folder[Any, Any, Boolean] {
+      def empty(context: Any): Boolean                                                   = true
+      def failCase(context: Any, error: Any, stackTrace: StackTrace): Boolean            = true
+      def dieCase(context: Any, t: Throwable, stackTrace: StackTrace): Boolean           = false
+      def interruptCase(context: Any, fiberId: FiberId, stackTrace: StackTrace): Boolean = false
+
+      def bothCase(context: Any, left: Boolean, right: Boolean): Boolean           = left && right
+      def thenCase(context: Any, left: Boolean, right: Boolean): Boolean           = left && right
+      def stacklessCase(context: Any, value: Boolean, stackless: Boolean): Boolean = value
+    }
+
     final case class Filter[E](p: Cause[E] => Boolean) extends Folder[Any, E, Cause[E]] {
       def empty(context: Any): Cause[E] = Cause.empty
 

diff --git a/core/shared/src/main/scala/zio/ZIO.scala b/core/shared/src/main/scala/zio/ZIO.scala
@@ -28,6 +28,13 @@ import scala.concurrent.ExecutionContext
 import scala.reflect.ClassTag
 import scala.util.control.NoStackTrace
 import izumi.reflect.macrortti.LightTypeTag
+import zio.Cause.Empty
+import zio.Cause.Then
+import zio.Cause.Interrupt
+import zio.Cause.Die
+import zio.Cause.Stackless
+import zio.Cause.Fail
+import zio.Cause.Both
 
 /**
  * A `ZIO[R, E, A]` value is an immutable value (called an "effect") that
@@ -297,12 +304,17 @@ sealed trait ZIO[-R, +E, +A]
     self.foldTraceZIO[R1, E2, A1](h, ZIO.successFn)
 
   /**
-   * Recovers from all errors with provided Cause.
+   * Recovers from all errors and defects with provided Cause.
    *
    * {{{
    * openFile("config.json").catchAllCause(_ => ZIO.succeed(defaultConfig))
    * }}}
    *
+   * '''WARNING''': There is no sensible way to recover from defects. This
+   * method should be used only at the boundary between ZIO and an external
+   * system, to transmit information on a defect for diagnostic or explanatory
+   * purposes. Consider using `catchAll` or `catchFailureCause` instead.
+   *
    * @see
    *   [[absorb]], [[sandbox]], [[mapErrorCause]] - other functions that can
    *   recover from defects
@@ -329,6 +341,28 @@ sealed trait ZIO[-R, +E, +A]
   ): ZIO[R1, E1, A1] =
     catchSomeDefect { case t => h(t) }
 
+  /**
+   * Recovers from all failures with provided function. Equivalent to `catchAll`
+   * but with the provided Cause. If you only care about the error, use
+   * `catchAll` instead.
+   *
+   * This method doesn't recover from defects. If you need to recover from
+   * defects, use `catchAllCause` instead.
+   *
+   * {{{
+   * effect.catchFailureCause(c => ZIO.logErrorCause("failure", c))
+   * }}}
+   */
+  final def catchFailureCause[R1 <: R, E2, A1 >: A](
+    h: Cause.Fail[E] => ZIO[R1, E2, A1]
+  )(implicit ev: CanFail[E], trace: Trace): ZIO[R1, E2, A1] =
+    self.foldTraceZIO[R1, E2, A1](
+      { case (e, stackTrace) =>
+        h(Cause.Fail(e, stackTrace))
+      },
+      ZIO.successFn
+    )
+
   /**
    * Recovers from all NonFatal Throwables.
    *
@@ -376,13 +410,18 @@ sealed trait ZIO[-R, +E, +A]
   }
 
   /**
-   * Recovers from some or all of the error cases with provided cause.
+   * Recovers from some or all of the error or defect cases with provided cause.
    *
    * {{{
    * openFile("data.json").catchSomeCause {
    *   case c if (c.interrupted) => openFile("backup.json")
    * }
    * }}}
+   *
+   * '''WARNING''': There is no sensible way to recover from defects. This
+   * method should be used only at the boundary between ZIO and an external
+   * system, to transmit information on a defect for diagnostic or explanatory
+   * purposes. Consider using `catchSomeFailureCause` instead.
    */
   final def catchSomeCause[R1 <: R, E1 >: E, A1 >: A](
     pf: PartialFunction[Cause[E], ZIO[R1, E1, A1]]
@@ -412,6 +451,33 @@ sealed trait ZIO[-R, +E, +A]
   )(implicit trace: Trace): ZIO[R1, E1, A1] =
     unrefineWith(pf)(ZIO.failFn).catchAll(identity)
 
+  /**
+   * Recovers from some or all of the failure cases with provided cause.
+   *
+   * This method only recovers from failures. If you need to recover from
+   * defects as well, use `catchSomeCause` or `catchSomeDefect` instead.
+   *
+   * {{{
+   * effect.catchSomeFailureCause {
+   *   case _: FileNotFoundException => createFile()
+   * }
+   * }}}
+   */
+  final def catchSomeFailureCause[R1 <: R, E1 >: E, A1 >: A](
+    pf: PartialFunction[Cause.Fail[E], ZIO[R1, E1, A1]]
+  )(implicit ev: CanFail[E], trace: Trace): ZIO[R1, E1, A1] = {
+    def tryRescue(c: Cause[E]): ZIO[R1, E1, A1] =
+      if (c.isFailureOnly) {
+        c.find { case f: Cause.Fail[E] => f } match {
+          case Some(f) => pf.applyOrElse(f, (_: Cause.Fail[E]) => Exit.failCause(c))
+          case None    => Exit.failCause(c)
+        }
+      } else {
+        Exit.failCause(c)
+      }
+    self.foldCauseZIO(tryRescue, ZIO.successFn)
+  }
+
   /**
    * Returns an effect that succeeds with the cause of failure of this effect,
    * or `Cause.empty` if the effect did succeed.
@@ -2101,6 +2167,10 @@ sealed trait ZIO[-R, +E, +A]
   /**
    * Returns an effect that effectually "peeks" at the cause of the failure of
    * this effect.
+   *
+   * This method "peeks" at both the failure and defect of this effect. If you
+   * only need to "peek" at the failure, use `tapFailureCause` instead.
+   *
    * {{{
    * readFile("data.json").tapErrorCause(logCause(_))
    * }}}
@@ -2121,6 +2191,26 @@ sealed trait ZIO[-R, +E, +A]
       ZIO.successFn
     )
 
+  /**
+   * Returns an effect that effectfully "peeks" at the failure of this effect.
+   *
+   * This method only "peeks" at the failure of this effect. If you need to
+   * "peek" at defects as well, use `tapErrorCause` or `tapDefect` instead.
+   *
+   * {{{
+   * readFile("data.json").tapError(logError(_))
+   * }}}
+   */
+      final def tapFailureCause[R1 <: R, E1 >: E](
+    f: Cause.Fail[E] => ZIO[R1, E1, Any]
+  )(implicit ev: CanFail[E], trace: Trace): ZIO[R1, E1, A] =
+    self.foldCauseZIO(
+      c =>
+        c.find { case failure: Cause.Fail[E] => failure }
+          .fold[ZIO[R1, E1, Nothing]](Exit.failCause(c))(f(_) *> Exit.failCause(c)),
+      ZIO.successFn
+    )
+
   /**
    * Returns an effect that effectfully "peeks" at the success of this effect.
    * If the partial function isn't defined at the input, the result is