zio · iravid · Nov 17, 2019 · Oct 22, 2019 · Oct 19, 2019 · Nov 14, 2019
diff --git a/core-tests/shared/src/test/scala/zio/RefMSpec.scala b/core-tests/shared/src/test/scala/zio/RefMSpec.scala
@@ -3,6 +3,8 @@ package zio
 import zio.test._
 import zio.test.Assertion._
 import zio.RefMSpecUtils._
+import zio.clock.Clock
+import zio.duration.durationInt
 
 object RefMSpec
     extends ZIOBaseSpec(
@@ -101,6 +103,23 @@ object RefMSpec
             refM  <- RefM.make[State](Active)
             value <- refM.modifySome("State doesn't change") { case Active => IO.fail(failure) }.run
           } yield assert(value, fails(equalTo(failure)))
+        },
+        testM("modifySome with fatal error") {
+          for {
+            refM  <- RefM.make[State](Active)
+            value <- refM.modifySome("State doesn't change") { case Active => IO.dieMessage(fatalError) }.run
+          } yield assert(value, dies(hasMessage(fatalError)))
+        },
+        testM("interrupt parent fiber and update") {
+          for {
+            promise     <- Promise.make[Nothing, RefM[State]]
+            latch       <- Promise.make[Nothing, Unit]
+            makeAndWait = promise.complete(RefM.make[State](Active)) *> latch.await
+            fiber       <- makeAndWait.fork
+            refM        <- promise.await
+            _           <- fiber.interrupt
+            value       <- refM.update(_ => ZIO.succeed(Closed)).timeout(1.second).provide(Clock.Live)
+          } yield assert(value, equalTo(Some(Closed)))
         }
       )
     )
@@ -109,6 +128,7 @@ object RefMSpecUtils {
 
   val (current, update) = ("value", "new value")
   val failure           = "failure"
+  val fatalError        = ":-0"
 
   sealed trait State
   case object Active  extends State

diff --git a/core-tests/shared/src/test/scala/zio/ZManagedSpec.scala b/core-tests/shared/src/test/scala/zio/ZManagedSpec.scala
@@ -3,10 +3,10 @@ package zio
 import zio.Cause.Interrupt
 import zio.duration._
 import zio.Exit.Failure
-import zio.test.{ Gen, testM, _ }
 import zio.test.Assertion._
 import zio.test.environment._
 import ZManagedSpecUtil._
+import zio.test._
 
 object ZManagedSpec
     extends ZIOBaseSpec(
@@ -880,6 +880,107 @@ object ZManagedSpec
             } yield assert(result, equalTo(List("Second", "First")))
           }
         ),
+        suite("memoize")(
+          testM("acquires and releases exactly once") {
+            for {
+              effects  <- Ref.make[List[Int]](Nil)
+              res      = (x: Int) => ZManaged.make(effects.update(x :: _))(_ => effects.update(x :: _))
+              program  = res(1) *> res(2) *> res(3)
+              memoized = program.memoize
+              _ <- memoized.use { managed =>
+                    val use = managed.use_(ZIO.unit)
+                    use *> use *> use
+                  }
+              res <- effects.get
+            } yield assert(res, equalTo(List(1, 2, 3, 3, 2, 1)))
+          },
+          testM("acquires and releases nothing if the inner managed is never used") {
+            for {
+              acquired <- Ref.make(false)
+              released <- Ref.make(false)
+              managed  = Managed.make(acquired.set(true))(_ => released.set(true))
+              _        <- managed.memoize.use_(ZIO.unit)
+              res      <- assertM(acquired.get zip released.get, equalTo((false, false)))
+            } yield res
+          },
+          testM("behaves like an ordinary ZManaged if flattened") {
+            for {
+              resource <- Ref.make(0)
+              acquire  = resource.update(_ + 1)
+              release  = resource.update(_ - 1)
+              managed  = ZManaged.make(acquire)(_ => release).memoize.flatten
+              res1     <- managed.use_(assertM(resource.get, equalTo(1)))
+              res2     <- assertM(resource.get, equalTo(0))
+            } yield res1 && res2
+          },
+          testM("properly raises an error if acquiring fails") {
+            for {
+              released <- Ref.make(false)
+              error    = ":-o"
+              managed  = Managed.make(ZIO.fail(error))(_ => released.set(true))
+              res1 <- managed.memoize.use { memoized =>
+                       for {
+                         v1 <- memoized.use_(ZIO.unit).either
+                         v2 <- memoized.use_(ZIO.unit).either
+                       } yield assert(v1, equalTo(v2)) && assert(v1, isLeft(equalTo(error)))
+                     }
+              res2 <- assertM(released.get, isFalse)
+            } yield res1 && res2
+          },
+          testM("behaves properly if acquiring dies") {
+            for {
+              released <- Ref.make(false)
+              ohNoes   = ";-0"
+              managed  = Managed.make(ZIO.dieMessage(ohNoes))(_ => released.set(true))
+              res1 <- managed.memoize.use { memoized =>
+                       assertM(memoized.use_(ZIO.unit).run, dies(hasMessage(ohNoes)))
+                     }
+              res2 <- assertM(released.get, isFalse)
+            } yield res1 && res2
+          },
+          testM("behaves properly if releasing dies") {
+            val myBad   = "#@*!"
+            val managed = Managed.make(ZIO.unit)(_ => ZIO.dieMessage(myBad))
+
+            val program = managed.memoize.use { memoized =>
+              memoized.use_(ZIO.unit)
+            }
+
+            assertM(program.run, dies(hasMessage(myBad)))
+          },
+          testM("behaves properly if use dies") {
+            val darn = "darn"
+            for {
+              latch    <- Promise.make[Nothing, Unit]
+              released <- Ref.make(false)
+              managed  = Managed.make(ZIO.unit)(_ => released.set(true) *> latch.succeed(()))
+              v1 <- managed.memoize.use { memoized =>
+                     memoized.use_(ZIO.dieMessage(darn))
+                   }.run
+              v2 <- released.get
+            } yield assert(v1, dies(hasMessage(darn))) && assert(v2, isTrue)
+          },
+          testM("behaves properly if use is interrupted") {
+            for {
+              latch1   <- Promise.make[Nothing, Unit]
+              latch2   <- Promise.make[Nothing, Unit]
+              latch3   <- Promise.make[Nothing, Unit]
+              resource <- Ref.make(0)
+              acquire  = resource.update(_ + 1)
+              release  = resource.update(_ - 1) *> latch3.succeed(())
+              managed  = ZManaged.make(acquire)(_ => release)
+              fiber <- managed.memoize.use { memoized =>
+                        memoized.use_(latch1.succeed(()) *> latch2.await)
+                      }.fork
+              _    <- latch1.await
+              res1 <- assertM(resource.get, equalTo(1))
+              _    <- fiber.interrupt
+              _    <- latch3.await
+              res2 <- assertM(resource.get, equalTo(0))
+              res3 <- assertM(latch2.isDone, isFalse)
+            } yield res1 && res2 && res3
+          }
+        ),
         suite("catch")(
           testM("catchAllCause") {
             val zm: ZManaged[Any, String, String] =

diff --git a/core/shared/src/main/scala/zio/RefM.scala b/core/shared/src/main/scala/zio/RefM.scala
@@ -146,7 +146,7 @@ object RefM extends Serializable {
       interrupted.get.flatMap {
         case Some(cause) => onDefect(cause)
         case None =>
-          update(a).foldM(e => onDefect(Cause.fail(e)) <* promise.fail(e), {
+          update(a).foldCauseM(c => onDefect(c).ensuring(promise.halt(c)), {
             case (b, a) => ref.set(a) <* promise.succeed(b)
           })
       }
@@ -169,7 +169,7 @@ object RefM extends Serializable {
     for {
       ref   <- Ref.make(a)
       queue <- Queue.bounded[Bundle[_, A, _]](n)
-      _     <- queue.take.flatMap(b => ref.get.flatMap(a => b.run(a, ref, onDefect))).forever.fork
+      _     <- queue.take.flatMap(b => ref.get.flatMap(a => b.run(a, ref, onDefect))).forever.fork.daemon
     } yield new RefM[A](ref, queue)
 
 }
diff --git a/core/shared/src/main/scala/zio/ZManaged.scala b/core/shared/src/main/scala/zio/ZManaged.scala
@@ -872,6 +872,31 @@ final class ZManaged[-R, +E, +A] private (reservation: ZIO[R, E, Reservation[R,
       }
     }
 
+  def memoize: ZManaged[R, E, ZManaged[R, E, A]] =
+    ZManaged {
+      RefM.make[Option[(Reservation[R, E, A], Exit[E, A])]](None).map { ref =>
+        val acquire1: ZIO[R, E, A] =
+          ref.modify {
+            case v @ Some((_, e)) => ZIO.succeed(e -> v)
+            case None =>
+              ZIO.uninterruptibleMask { restore =>
+                self.reserve.flatMap { res =>
+                  restore(res.acquire).run.map(e => e -> Some(res -> e))
+                }
+              }
+          }.flatMap(ZIO.done)
+
+        val acquire2: ZIO[R, E, ZManaged[R, E, A]] =
+          ZIO.succeed(acquire1.toManaged_)
+
+        val release2 = (_: Exit[_, _]) =>
+          ref.updateSome {
+            case Some((res, e)) => res.release(e).as(None)
+          }
+
+        Reservation(acquire2, release2)
+      }
+    }
 }
 
 object ZManaged {

diff --git a/test/shared/src/main/scala/zio/test/Assertion.scala b/test/shared/src/main/scala/zio/test/Assertion.scala
@@ -257,6 +257,12 @@ object Assertion {
       case _                   => None
     }
 
+  /**
+   * Makes a new assertion that requires an exception to have a certain message.
+   */
+  final def hasMessage(message: String): Assertion[Throwable] =
+    assertion[Throwable]("hasMessage")(param(message))(th => th.getMessage == message)
+
   /**
    * Makes a new assertion that requires a given string to end with the specified suffix.
    */