Add a max size for normalize cache (#2878)

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. Co-authored-by: Jules Ivanic <[email protected]>
zio · Sep 29, 2023 · 5402d84 · 5402d84
1 parent a6a17ae
commit 5402d84
Show file tree

Hide file tree

Showing 3 changed files with 34 additions and 15 deletions.
diff --git a/build.sbt b/build.sbt
@@ -77,6 +77,16 @@ lazy val scala213Modules =
 
 lazy val scala3Modules = Seq[sbt.ClasspathDep[sbt.ProjectReference]](`quill-engine`, `quill-util`)
 
+def javaMainVersion = {
+  val javaVersion = System.getProperty("java.version")
+  val versionSplit = javaVersion.split('.')
+  if (versionSplit.head.equals("1")) {
+    versionSplit(1).toInt
+  } else {
+    versionSplit(0).toInt
+  }
+}
+
 def isScala213 = {
   val scalaVersion = sys.props.get("quill.scala.version")
   scalaVersion.map(_.startsWith("2.13")).getOrElse(false)
@@ -206,11 +216,17 @@ 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",
+        // caffeine 3.x doesn't support Java version lower than 11
+        if (javaMainVersion >= 11) {
+          "com.github.ben-manes.caffeine" % "caffeine" % "3.1.8"
+        } else {
+          "com.github.ben-manes.caffeine" % "caffeine" % "2.9.3"
+        }
       ),
       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,20 @@
 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
 
 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, ast => f(ast))
     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")))