Improve performance of ZQuery#run (#497)

* Improve performance of `ZQuery#run`

* Implement using a custom scope

* Reset finalizers to avoid leaks

* Cleanups

Author: kyri-petrou

benchmarks/src/main/scala/zio/query/BenchmarkUtil.scala

@@ -14,4 +14,7 @@ object BenchmarkUtil extends Runtime[Any] { self =>
 
   def unsafeRunCache[E, A](query: ZQuery[Any, E, A], cache: Cache): A =
     Unsafe.unsafe(implicit unsafe => self.unsafe.run(query.runCache(cache)).getOrThrowFiberFailure())
+
+  def unsafeRunZIO[E, A](query: ZIO[Any, E, A]): A =
+    Unsafe.unsafe(implicit unsafe => self.unsafe.run(query).getOrThrowFiberFailure())
 }

benchmarks/src/main/scala/zio/query/DataSourceBenchmark.scala

@@ -3,7 +3,7 @@ package zio.query
 import cats.effect.IO
 import cats.effect.unsafe.implicits._
 import cats.syntax.all._
-import fetch.Fetch
+import fetch.{Fetch, fetchM}
 import org.openjdk.jmh.annotations.{Scope => JScope, _}
 import zio.query.BenchmarkUtil._
 import zio.{Chunk, ZIO}
@@ -74,7 +74,7 @@ class DataSourceBenchmark {
 
   object ZQueryImpl {
     case class Req(i: Int) extends Request[Nothing, Int]
-    val ds = DataSource.fromFunctionBatchedZIO("PlusOne") { reqs: Chunk[Req] => ZIO.succeed(reqs.map(_.i + 1)) }
+    val ds = DataSource.fromFunctionBatchedZIO("PlusOne") { (reqs: Chunk[Req]) => ZIO.succeed(reqs.map(_.i + 1)) }
   }
 
   object FetchImpl {

benchmarks/src/main/scala/zio/query/FromRequestBenchmark.scala

@@ -69,5 +69,5 @@ class FromRequestBenchmark {
   }
 
   private case class Req(i: Int) extends Request[Nothing, Int]
-  private val ds = DataSource.fromFunctionBatchedZIO("Datasource") { reqs: Chunk[Req] => ZIO.succeed(reqs.map(_.i)) }
+  private val ds = DataSource.fromFunctionBatchedZIO("Datasource") { (reqs: Chunk[Req]) => ZIO.succeed(reqs.map(_.i)) }
 }

benchmarks/src/main/scala/zio/query/ZQueryBenchmark.scala

@@ -0,0 +1,31 @@
+package zio.query
+
+import org.openjdk.jmh.annotations.{Scope => JScope, _}
+import zio.{Chunk, ZIO}
+import zio.query.BenchmarkUtil._
+
+import java.util.concurrent.TimeUnit
+
+@Measurement(iterations = 3, time = 1, timeUnit = TimeUnit.SECONDS)
+@Warmup(iterations = 3, time = 1, timeUnit = TimeUnit.SECONDS)
+@Fork(1)
+@Threads(1)
+@State(JScope.Thread)
+@BenchmarkMode(Array(Mode.Throughput))
+@OutputTimeUnit(TimeUnit.MILLISECONDS)
+class ZQueryBenchmark {
+  val cache = Cache.unsafeMake()
+
+  val qs1 = Chunk.fill(1000)(ZQuery.succeedNow("foo").runCache(cache))
+  val qs2 = Chunk.fill(1000)(ZQuery.succeed("foo").runCache(cache))
+
+  @Benchmark
+  @OperationsPerInvocation(1000)
+  def zQueryRunSucceedNowBenchmark() =
+    unsafeRunZIO(ZIO.collectAllDiscard(qs1))
+
+  @Benchmark
+  @OperationsPerInvocation(1000)
+  def zQueryRunSucceedBenchmark() =
+    unsafeRunZIO(ZIO.collectAllDiscard(qs2))
+}

build.sbt

@@ -13,7 +13,7 @@ lazy val allScalas = List("2.12", "2.13", "3.3")
 inThisBuild(
   List(
     name         := "ZIO Query",
-    zioVersion   := "2.1.4",
+    zioVersion   := "2.1.6",
     scalaVersion := scalaV,
     developers := List(
       Developer(

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

@@ -536,12 +536,17 @@ final class ZQuery[-R, +E, +A] private (private val step: ZIO[R, Nothing, Result
   def runCache(cache: => Cache)(implicit trace: Trace): ZIO[R, E, A] =
     asExitOrElse(null) match {
       case null =>
-        ZIO.acquireReleaseExitWith {
-          Scope.make
-        } { (scope: Scope.Closeable, exit: Exit[E, A]) =>
-          scope.close(exit)
-        } { scope =>
-          ZQuery.currentScope.locally(scope)(ZQuery.currentCache.locally(cache)(runToZIO))
+        ZIO.uninterruptibleMask { restore =>
+          ZIO.withFiberRuntime[R, E, A] { (state, _) =>
+            val scope = QueryScope.make()
+            state.setFiberRef(ZQuery.currentCache, cache)
+            state.setFiberRef(ZQuery.currentScope, scope)
+            restore(runToZIO).exitWith { exit =>
+              state.deleteFiberRef(ZQuery.currentCache)
+              state.deleteFiberRef(ZQuery.currentScope)
+              scope.closeAndExitWith(exit)
+            }
+          }
         }
       case exit => exit
     }
@@ -1826,8 +1831,8 @@ object ZQuery {
   val currentCache: FiberRef[Cache] =
     FiberRef.unsafe.make(Cache.unsafeMake())(Unsafe.unsafe)
 
-  val currentScope: FiberRef[Scope] =
-    FiberRef.unsafe.make[Scope](Scope.global)(Unsafe.unsafe)
+  val currentScope: FiberRef[QueryScope] =
+    FiberRef.unsafe.make[QueryScope](QueryScope.NoOp)(Unsafe.unsafe)
 
   final class Acquire[-R, +E, +A](private val acquire: () => ZIO[R, E, A]) extends AnyVal {
     def apply[R1](release: A => URIO[R1, Any]): Release[R with R1, E, A] =

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

@@ -0,0 +1,42 @@
+package zio.query.internal
+
+import zio._
+import zio.stacktracer.TracingImplicits.disableAutoTrace
+
+import java.util.concurrent.atomic.AtomicReference
+
+/**
+ * Lightweight variant of [[zio.Scope]], optimized for usage with ZQuery
+ */
+sealed trait QueryScope {
+  def addFinalizerExit(f: Exit[Any, Any] => UIO[Any])(implicit trace: Trace): UIO[Unit]
+  def closeAndExitWith[E, A](exit: Exit[E, A])(implicit trace: Trace): IO[E, A]
+}
+
+private[query] object QueryScope {
+  def make(): QueryScope = new Default
+
+  case object NoOp extends QueryScope {
+    def addFinalizerExit(f: Exit[Any, Any] => UIO[Any])(implicit trace: Trace): UIO[Unit] = ZIO.unit
+    def closeAndExitWith[E, A](exit: Exit[E, A])(implicit trace: Trace): IO[E, A]         = exit
+  }
+
+  final private class Default extends QueryScope {
+    private val ref = new AtomicReference(List.empty[Exit[Any, Any] => UIO[Any]])
+
+    def addFinalizerExit(f: Exit[Any, Any] => UIO[Any])(implicit trace: Trace): UIO[Unit] =
+      ZIO.succeed {
+        ref.updateAndGet(f :: _)
+        ()
+      }
+
+    def closeAndExitWith[E, A](exit: Exit[E, A])(implicit trace: Trace): IO[E, A] = {
+      val finalizers = ref.get
+      if (finalizers.isEmpty) exit
+      else {
+        ref.set(Nil)
+        ZIO.foreachDiscard(finalizers)(_(exit)) *> exit
+      }
+    }
+  }
+}