Asynchronous Cache

zio-query/shared/src/main/scala/zio/query/Cache.scala

@@ -20,38 +20,36 @@ import zio._
 import zio.stacktracer.TracingImplicits.disableAutoTrace
 
 /**
- * A `Cache` maintains an internal state with a mapping from requests to `Ref`s
- * that will contain the result of those requests when they are executed. This
- * is used internally by the library to provide deduplication and caching of
- * requests.
+ * A `Cache` maintains an internal state with a mapping from requests to
+ * `Promise`s that will contain the result of those requests when they are
+ * executed. This is used internally by the library to provide deduplication and
+ * caching of requests.
  */
 trait Cache {
 
   /**
    * Looks up a request in the cache, failing with the unit value if the request
-   * is not in the cache, succeeding with `Ref(None)` if the request is in the
-   * cache but has not been executed yet, or `Ref(Some(value))` if the request
-   * has been executed.
+   * is not in the cache or succeeding with a `Promise` if the request is in the
+   * cache that will contain the result of the request when it is executed.
    */
-  def get[E, A](request: Request[E, A])(implicit trace: Trace): IO[Unit, Ref[Option[Exit[E, A]]]]
+  def get[E, A](request: Request[E, A])(implicit trace: Trace): IO[Unit, Promise[E, A]]
 
   /**
    * Looks up a request in the cache. If the request is not in the cache returns
-   * a `Left` with a `Ref` that can be set with a `Some` to complete the
-   * request. If the request is in the cache returns a `Right` with a `Ref` that
-   * either contains `Some` with a result if the request has been executed or
-   * `None` if the request has not been executed yet.
+   * a `Left` with a `Promise` that can be completed to complete the request. If
+   * the request is in the cache returns a `Right` with a `Promise` that will
+   * contain the result of the request when it is executed.
    */
   def lookup[R, E, A, B](request: A)(implicit
     ev: A <:< Request[E, B],
     trace: Trace
-  ): UIO[Either[Ref[Option[Exit[E, B]]], Ref[Option[Exit[E, B]]]]]
+  ): UIO[Either[Promise[E, B], Promise[E, B]]]
 
   /**
-   * Inserts a request and a `Ref` that will contain the result of the request
-   * when it is executed into the cache.
+   * Inserts a request and a `Promise` that will contain the result of the
+   * request when it is executed into the cache.
    */
-  def put[E, A](request: Request[E, A], result: Ref[Option[Exit[E, A]]])(implicit trace: Trace): UIO[Unit]
+  def put[E, A](request: Request[E, A], result: Promise[E, A])(implicit trace: Trace): UIO[Unit]
 
   /**
    * Removes a request from the cache.
@@ -67,31 +65,31 @@ object Cache {
   def empty(implicit trace: Trace): UIO[Cache] =
     ZIO.succeed(Cache.unsafeMake())
 
-  private final class Default(private val state: Ref[Map[Any, Any]]) extends Cache {
+  private final class Default(private val state: Ref[Map[Request[_, _], Promise[_, _]]]) extends Cache {
 
-    def get[E, A](request: Request[E, A])(implicit trace: Trace): IO[Unit, Ref[Option[Exit[E, A]]]] =
-      state.get.map(_.get(request).asInstanceOf[Option[Ref[Option[Exit[E, A]]]]]).some.orElseFail(())
+    def get[E, A](request: Request[E, A])(implicit trace: Trace): IO[Unit, Promise[E, A]] =
+      state.get.map(_.get(request).asInstanceOf[Option[Promise[E, A]]]).some.orElseFail(())
 
     def lookup[R, E, A, B](request: A)(implicit
       ev: A <:< Request[E, B],
       trace: Trace
-    ): UIO[Either[Ref[Option[Exit[E, B]]], Ref[Option[Exit[E, B]]]]] =
-      Ref.make(Option.empty[Exit[E, B]]).flatMap { ref =>
+    ): UIO[Either[Promise[E, B], Promise[E, B]]] =
+      Promise.make[E, B].flatMap { promise =>
         state.modify { map =>
           map.get(request) match {
-            case None      => (Left(ref), map + (request -> ref))
-            case Some(ref) => (Right(ref.asInstanceOf[Ref[Option[Exit[E, B]]]]), map)
+            case Some(promise) => (Right(promise.asInstanceOf[Promise[E, B]]), map)
+            case None          => (Left(promise), map + (ev(request) -> promise))
           }
         }
       }
 
-    def put[E, A](request: Request[E, A], result: Ref[Option[Exit[E, A]]])(implicit trace: Trace): UIO[Unit] =
+    def put[E, A](request: Request[E, A], result: Promise[E, A])(implicit trace: Trace): UIO[Unit] =
       state.update(_ + (request -> result))
 
     def remove[E, A](request: Request[E, A])(implicit trace: Trace): UIO[Unit] =
       state.update(_ - request)
   }
 
   private[query] def unsafeMake(): Cache =
-    new Default(Ref.unsafe.make(Map.empty[Any, Any])(Unsafe.unsafe))
+    new Default(Ref.unsafe.make(Map.empty[Request[_, _], Promise[_, _]])(Unsafe.unsafe))
 }

zio-query/shared/src/main/scala/zio/query/ZQuery.scala

@@ -1242,25 +1242,27 @@ object ZQuery {
           if (cachingEnabled) {
             ZQuery.currentCache.get.flatMap { cache =>
               cache.lookup(request).flatMap {
-                case Left(ref) =>
+                case Left(promise) =>
                   ZIO.succeed(
                     Result.blocked(
-                      BlockedRequests.single(dataSource, BlockedRequest(request, ref)),
-                      Continue(request, dataSource, ref)
+                      BlockedRequests.single(dataSource, BlockedRequest(request, promise)),
+                      Continue(request, dataSource, promise)
                     )
                   )
-                case Right(ref) =>
-                  ref.get.map {
-                    case None    => Result.blocked(BlockedRequests.empty, Continue(request, dataSource, ref))
-                    case Some(b) => Result.fromExit(b)
+                case Right(promise) =>
+                  promise.poll.flatMap {
+                    case None =>
+                      ZIO.succeed(Result.blocked(BlockedRequests.empty, Continue(request, dataSource, promise)))
+                    case Some(io) =>
+                      io.exit.map(Result.fromExit)
                   }
               }
             }
           } else {
-            Ref.make(Option.empty[Exit[E, B]]).map { ref =>
+            Promise.make[E, B].map { promise =>
               Result.blocked(
-                BlockedRequests.single(dataSource, BlockedRequest(request, ref)),
-                Continue(request, dataSource, ref)
+                BlockedRequests.single(dataSource, BlockedRequest(request, promise)),
+                Continue(request, dataSource, promise)
               )
             }
           }

zio-query/shared/src/main/scala/zio/query/internal/BlockedRequest.scala

@@ -18,30 +18,30 @@ package zio.query.internal
 
 import zio.query.Request
 import zio.stacktracer.TracingImplicits.disableAutoTrace
-import zio.{Exit, Ref}
+import zio.{Exit, Promise}
 
 /**
- * A `BlockedRequest[A]` keeps track of a request of type `A` along with a `Ref`
- * containing the result of the request, existentially hiding the result type.
- * This is used internally by the library to support data sources that return
- * different result types for different requests while guaranteeing that results
- * will be of the type requested.
+ * A `BlockedRequest[A]` keeps track of a request of type `A` along with a
+ * `Promise` containing the result of the request, existentially hiding the
+ * result type. This is used internally by the library to support data sources
+ * that return different result types for different requests while guaranteeing
+ * that results will be of the type requested.
  */
 private[query] sealed trait BlockedRequest[+A] {
   type Failure
   type Success
 
   def request: Request[Failure, Success]
 
-  def result: Ref[Option[Exit[Failure, Success]]]
+  def result: Promise[Failure, Success]
 
   override final def toString: String =
     s"BlockedRequest($request, $result)"
 }
 
 private[query] object BlockedRequest {
 
-  def apply[E, A, B](request0: A, result0: Ref[Option[Exit[E, B]]])(implicit
+  def apply[E, A, B](request0: A, result0: Promise[E, B])(implicit
     ev: A <:< Request[E, B]
   ): BlockedRequest[A] =
     new BlockedRequest[A] {

zio-query/shared/src/main/scala/zio/query/internal/BlockedRequests.scala

@@ -17,9 +17,9 @@
 package zio.query.internal
 
 import zio.query.internal.BlockedRequests._
-import zio.query.{Cache, CompletedRequestMap, DataSource, DataSourceAspect, Described, ZQuery}
+import zio.query.{Cache, CompletedRequestMap, DataSource, DataSourceAspect, Described, QueryFailure, ZQuery}
 import zio.stacktracer.TracingImplicits.disableAutoTrace
-import zio.{Exit, Ref, Trace, ZEnvironment, ZIO}
+import zio.{Exit, Promise, Trace, ZEnvironment, ZIO}
 
 import scala.annotation.tailrec
 
@@ -115,10 +115,16 @@ private[query] sealed trait BlockedRequests[-R] { self =>
             blockedRequests = sequential.flatten
             leftovers       = completedRequests.requests -- blockedRequests.map(_.request)
             _ <- ZIO.foreachDiscard(blockedRequests) { blockedRequest =>
-                   blockedRequest.result.set(completedRequests.lookup(blockedRequest.request))
+                   blockedRequest.result.done(
+                     completedRequests
+                       .lookup(blockedRequest.request)
+                       .getOrElse(Exit.die(QueryFailure(dataSource, blockedRequest.request)))
+                   )
                  }
             _ <- ZIO.foreachDiscard(leftovers) { request =>
-                   Ref.make(completedRequests.lookup(request)).flatMap(cache.put(request, _))
+                   ZIO.foreach(completedRequests.lookup(request)) { response =>
+                     Promise.make[Any, Any].tap(_.done(response)).flatMap(cache.put(request, _))
+                   }
                  }
           } yield ()
         }

zio-query/shared/src/main/scala/zio/query/internal/Continue.scala

@@ -160,19 +160,14 @@ private[query] sealed trait Continue[-R, +E, +A] { self =>
 private[query] object Continue {
 
   /**
-   * Constructs a continuation from a request, a data source, and a `Ref` that
-   * will contain the result of the request when it is executed.
+   * Constructs a continuation from a request, a data source, and a `Promise`
+   * that will contain the result of the request when it is executed.
    */
-  def apply[R, E, A, B](request: A, dataSource: DataSource[R, A], ref: Ref[Option[Exit[E, B]]])(implicit
+  def apply[R, E, A, B](request: A, dataSource: DataSource[R, A], promise: Promise[E, B])(implicit
     ev: A <:< Request[E, B],
     trace: Trace
   ): Continue[R, E, B] =
-    Continue.get {
-      ref.get.flatMap {
-        case None    => ZIO.die(QueryFailure(dataSource, request))
-        case Some(b) => ZIO.done(b)
-      }
-    }
+    Continue.get(promise.await)
 
   /**
    * Collects a collection of continuation into a continuation returning a