Add a max size for normalize cache

This fixes #2484. Added a capacity to normalize cache. Default to 1024 and can be override by `quill.query.cacheDynamicMaxSize`. Added a new dependency Caffeine to use its cache implementation, which is recommended by Guava. For a critical use case in quill, I think it's better to have a widely tested implementation than reinventing the wheel.
zio · Sep 29, 2023 · 2b81fc5 · 2b81fc5
1 parent a6a17ae
commit 2b81fc5
Show file tree

Hide file tree

Showing 3 changed files with 21 additions and 15 deletions.
diff --git a/build.sbt b/build.sbt
@@ -206,11 +206,12 @@ lazy val `quill-engine` =
     .settings(commonSettings: _*)
     .settings(
       libraryDependencies ++= Seq(
-        "com.typesafe"                % "config"        % "1.4.2",
-        "com.typesafe.scala-logging" %% "scala-logging" % "3.9.5",
-        ("com.github.takayahilton"   %% "sql-formatter" % "1.2.1").cross(CrossVersion.for3Use2_13),
-        "io.suzaku"                  %% "boopickle"     % "1.4.0",
-        "com.lihaoyi"                %% "pprint"        % "0.8.1"
+        "com.typesafe"                  % "config"        % "1.4.2",
+        "com.typesafe.scala-logging"   %% "scala-logging" % "3.9.5",
+        ("com.github.takayahilton"     %% "sql-formatter" % "1.2.1").cross(CrossVersion.for3Use2_13),
+        "io.suzaku"                    %% "boopickle"     % "1.4.0",
+        "com.lihaoyi"                  %% "pprint"        % "0.8.1",
+        "com.github.ben-manes.caffeine" % "caffeine"      % "3.1.8"
       ),
       coverageExcludedPackages := "<empty>;.*AstPrinter;.*Using;io.getquill.Model;io.getquill.ScalarTag;io.getquill.QuotationTag"
     )

diff --git a/quill-engine/src/main/scala/io/getquill/norm/NormalizeCaching.scala b/quill-engine/src/main/scala/io/getquill/norm/NormalizeCaching.scala
@@ -1,21 +1,22 @@
 package io.getquill.norm
 
+import com.github.benmanes.caffeine.cache.{Cache, Caffeine}
 import io.getquill.ast.Ast
-import java.util.concurrent.ConcurrentHashMap
+import io.getquill.util.Messages
+
+import scala.jdk.FunctionConverters.enrichAsJavaFunction
 
 object NormalizeCaching {
-  private val cache = new ConcurrentHashMap[Ast, Ast]
+
+  private val cache: Cache[Ast, Ast] = Caffeine
+    .newBuilder()
+    .maximumSize(Messages.cacheDynamicMaxSize)
+    .recordStats()
+    .build()
 
   def apply(f: Ast => Ast): Ast => Ast = { ori =>
     val (stabilized, state) = StabilizeLifts.stabilize(ori)
-    val cachedR             = cache.get(stabilized)
-    val normalized = if (cachedR != null) {
-      cachedR
-    } else {
-      val r = f(stabilized)
-      cache.put(stabilized, r)
-      r
-    }
+    val normalized          = cache.get(stabilized, f.asJava)
     StabilizeLifts.revert(normalized, state)
   }
 

diff --git a/quill-engine/src/main/scala/io/getquill/util/Messages.scala b/quill-engine/src/main/scala/io/getquill/util/Messages.scala
@@ -49,6 +49,10 @@ object Messages {
     "quill.query.cacheDynamic",
     variable("quill.query.cacheDynamic", "query_query_cacheDynamic", "true").toBoolean
   )
+  def cacheDynamicMaxSize = cache(
+    "quill.query.cacheDynamicMaxSize",
+    variable("quill.query.cacheDynamicMaxSize", "query_query_cacheDynamicMaxSize", "1024").toLong
+  )
   def querySubexpand =
     cache("quill.query.subexpand", variable("quill.query.subexpand", "query_query_subexpand", "true").toBoolean)
   def quillLogFile = cache("quill.log.file", LogToFile(variable("quill.log.file", "quill_log_file", "false")))