[PLAT-33341] Add Caffeine interface to limit Sjsonnet worker cache si…

…ze (#128) - The current map-based parse cache is now DefaultParseCache - The new ParseCache abstraction can be used to supply other cache implementations
databricks · Oct 27, 2021 · 8193180 · 8193180
1 parent cbce3f0
commit 8193180
Show file tree

Hide file tree

Showing 19 changed files with 63 additions and 28 deletions.
diff --git a/bench/src/main/scala/sjsonnet/MainBenchmark.scala b/bench/src/main/scala/sjsonnet/MainBenchmark.scala
@@ -19,15 +19,17 @@ object MainBenchmark {
     val file = config.file
     val wd = os.pwd
     val path = OsPath(os.Path(file, wd))
+    val parseCache = new DefaultParseCache
     val interp = new Interpreter(
       Map.empty[String, ujson.Value],
       Map.empty[String, ujson.Value],
       OsPath(wd),
-      importer = SjsonnetMain.resolveImport(config.jpaths.map(os.Path(_, wd)).map(OsPath(_)), None)
+      importer = SjsonnetMain.resolveImport(config.jpaths.map(os.Path(_, wd)).map(OsPath(_)), None),
+      parseCache = parseCache
     )
     val renderer = new Renderer(new StringWriter, indent = 3)
     interp.interpret0(interp.resolver.read(path).get, path, renderer).getOrElse(???)
-    (interp.parseCache.keySet.toIndexedSeq, interp.evaluator)
+    (parseCache.keySet.toIndexedSeq, interp.evaluator)
   }
 }
 
@@ -50,7 +52,7 @@ class MainBenchmark {
   def main(bh: Blackhole): Unit = {
     bh.consume(SjsonnetMain.main0(
       MainBenchmark.mainArgs,
-      collection.mutable.HashMap.empty,
+      new DefaultParseCache,
       System.in,
       dummyOut,
       System.err,

diff --git a/bench/src/main/scala/sjsonnet/MaterializerBenchmark.scala b/bench/src/main/scala/sjsonnet/MaterializerBenchmark.scala
@@ -32,6 +32,7 @@ class MaterializerBenchmark {
       Map.empty[String, ujson.Value],
       OsPath(wd),
       importer = SjsonnetMain.resolveImport(config.jpaths.map(os.Path(_, wd)).map(OsPath(_)), None),
+      parseCache = new DefaultParseCache
     )
     value = interp.evaluate(os.read(path), OsPath(path)).getOrElse(???)
     assert(renderYaml() == oldRenderYaml())

diff --git a/bench/src/main/scala/sjsonnet/RunProfiler.scala b/bench/src/main/scala/sjsonnet/RunProfiler.scala
@@ -8,11 +8,13 @@ object RunProfiler extends App {
   val file = config.file
   val wd = os.pwd
   val path = OsPath(os.Path(file, wd))
+  val parseCache = new DefaultParseCache
   val interp = new Interpreter(
     Map.empty[String, ujson.Value],
     Map.empty[String, ujson.Value],
     OsPath(wd),
     importer = SjsonnetMain.resolveImport(config.jpaths.map(os.Path(_, wd)).map(OsPath(_)), None),
+    parseCache = parseCache
   ) {
     override def createEvaluator(resolver: CachedResolver, extVars: Map[String, ujson.Value], wd: Path,
                                  settings: Settings, warn: Error => Unit): Evaluator =
@@ -36,7 +38,7 @@ object RunProfiler extends App {
   profiler.clear()
   val total = (for(i <- 1 to 5) yield run()).sum
 
-  val roots = interp.parseCache.valuesIterator.map(_.getOrElse(???)).map(_._1).toSeq
+  val roots = parseCache.valuesIterator.map(_.getOrElse(???)).map(_._1).toSeq
   roots.foreach(profiler.accumulate)
 
   println(s"\nTop 20 by time:")

diff --git a/readme.md b/readme.md
@@ -18,7 +18,7 @@ Sjsonnet can be used from Java:
 ```java
 sjsonnet.SjsonnetMain.main0(
     new String[]{"foo.jsonnet"},
-    sjsonnet.SjsonnetMain.createParseCache(),
+    new DefaultParseCache,
     System.in,
     System.out,
     System.err,
@@ -37,7 +37,7 @@ ivy"com.databricks::sjsonnet:0.4.1" // Mill
 ```scala
 sjsonnet.SjsonnetMain.main0(
     Array("foo.jsonnet"),
-    sjsonnet.SjsonnetMain.createParseCache(),
+    new DefaultParseCache,
     System.in,
     System.out,
     System.err,

diff --git a/sjsonnet/bench/src-jvm-native/sjsonnet/SjsonnetTestMain.scala b/sjsonnet/bench/src-jvm-native/sjsonnet/SjsonnetTestMain.scala
@@ -45,7 +45,7 @@ object SjsonnetTestMain {
 
     val start = System.currentTimeMillis()
     var count = 0
-    val parseCache = sjsonnet.SjsonnetMain.createParseCache()
+    val parseCache = new DefaultParseCache
     while(System.currentTimeMillis() - start < 20000){
       count += 1
       for(name <- names/*Seq(
@@ -62,11 +62,11 @@ object SjsonnetTestMain {
         val path = testSuiteRoot / s"$name.jsonnet"
         var currentPos: Position = null
         val interp = new Interpreter(
-          parseCache,
           Map("var1" -> "test", "var2" -> ujson.Obj("x" -> 1, "y" -> 2), "isKubecfg" -> true),
           Map("var1" -> "test", "var2" -> ujson.Obj("x" -> 1, "y" -> 2)),
           OsPath(os.pwd),
           SjsonnetMain.resolveImport(Seq(), None),
+          parseCache,
           storePos = currentPos = _
         )
         val writer = new java.io.StringWriter

diff --git a/sjsonnet/client/src/sjsonnet/client/Util.java b/sjsonnet/client/src/sjsonnet/client/Util.java
@@ -70,11 +70,11 @@ public static String readString(InputStream inputStream) throws IOException {
                 total += res;
             }
         }
-        return new String(arr);
+        return new String(arr, "UTF-8");
     }
 
     public static void writeString(OutputStream outputStream, String string) throws IOException {
-        byte[] bytes = string.getBytes();
+        byte[] bytes = string.getBytes("UTF-8");
         writeInt(outputStream, bytes.length);
         outputStream.write(bytes);
     }

diff --git a/sjsonnet/server/src/sjsonnet/SjsonnetServerMain.scala b/sjsonnet/server/src/sjsonnet/SjsonnetServerMain.scala
@@ -24,7 +24,7 @@ trait SjsonnetServerMain[T]{
             wd: os.Path): (Boolean, Option[T])
 }
 
-object SjsonnetServerMain extends SjsonnetServerMain[collection.mutable.HashMap[(Path, String), Either[Error, (Expr, FileScope)]]]{
+object SjsonnetServerMain extends SjsonnetServerMain[DefaultParseCache]{
   def main(args0: Array[String]): Unit = {
     // Disable SIGINT interrupt signal in the Mill server.
     //
@@ -45,7 +45,7 @@ object SjsonnetServerMain extends SjsonnetServerMain[collection.mutable.HashMap[
     ).run()
   }
   def main0(args: Array[String],
-            stateCache: Option[collection.mutable.HashMap[(Path, String), Either[Error, (Expr, FileScope)]]],
+            stateCache: Option[DefaultParseCache],
             mainInteractive: Boolean,
             stdin: InputStream,
             stdout: PrintStream,
@@ -55,7 +55,7 @@ object SjsonnetServerMain extends SjsonnetServerMain[collection.mutable.HashMap[
             wd: os.Path) = {
 
     val stateCache2 = stateCache.getOrElse{
-      val p = collection.mutable.HashMap[(Path, String), Either[Error, (Expr, FileScope)]]()
+      val p = new DefaultParseCache
       this.stateCache = Some(p)
       p
     }

diff --git a/sjsonnet/src-js/sjsonnet/SjsonnetMain.scala b/sjsonnet/src-js/sjsonnet/SjsonnetMain.scala
@@ -6,7 +6,6 @@ import scala.scalajs.js.annotation.{JSExport, JSExportTopLevel}
 
 @JSExportTopLevel("SjsonnetMain")
 object SjsonnetMain {
-  def createParseCache() = collection.mutable.HashMap[(Path, String), Either[Error, (Expr, FileScope)]]()
   @JSExport
   def interpret(text: String,
                 extVars: js.Any,
@@ -28,8 +27,8 @@ object SjsonnetMain {
         def read(path: Path): Option[String] =
           Option(importLoader(path.asInstanceOf[JsVirtualPath].path))
       },
+      parseCache = new DefaultParseCache,
       new Settings(preserveOrder = preserveOrder),
-      parseCache = createParseCache()
     )
     interp.interpret0(text, JsVirtualPath("(memory)"), ujson.WebJson.Builder) match{
       case Left(msg) => throw new js.JavaScriptException(msg)

diff --git a/sjsonnet/src-jvm-native/sjsonnet/SjsonnetMain.scala b/sjsonnet/src-jvm-native/sjsonnet/SjsonnetMain.scala
@@ -8,8 +8,6 @@ import scala.util.Try
 import scala.util.control.NonFatal
 
 object SjsonnetMain {
-  def createParseCache() = collection.mutable.HashMap[(Path, String), Either[Error, (Expr, FileScope)]]()
-
   def resolveImport(searchRoots0: Seq[Path], allowedInputs: Option[Set[os.Path]] = None) = new Importer {
     def resolve(docBase: Path, importName: String): Option[Path] =
       (docBase +: searchRoots0)
@@ -33,7 +31,7 @@ object SjsonnetMain {
         case Array(s, _*) if s == "-i" || s == "--interactive" => args.tail
         case _ => args
       },
-      collection.mutable.HashMap.empty,
+      new DefaultParseCache,
       System.in,
       System.out,
       System.err,
@@ -44,7 +42,7 @@ object SjsonnetMain {
   }
 
   def main0(args: Array[String],
-            parseCache: collection.mutable.HashMap[(Path, String), Either[Error, (Expr, FileScope)]],
+            parseCache: ParseCache,
             stdin: InputStream,
             stdout: PrintStream,
             stderr: PrintStream,
@@ -137,7 +135,7 @@ object SjsonnetMain {
 
   def mainConfigured(file: String,
                      config: Config,
-                     parseCache: collection.mutable.HashMap[(Path, String), Either[Error, (Expr, FileScope)]],
+                     parseCache: ParseCache,
                      wd: os.Path,
                      allowedInputs: Option[Set[os.Path]] = None,
                      importer: Option[(Path, String) => Option[os.Path]] = None,
@@ -193,13 +191,13 @@ object SjsonnetMain {
         }
         case None => resolveImport(config.jpaths.map(os.Path(_, wd)).map(OsPath(_)), allowedInputs)
       },
+      parseCache,
       settings = new Settings(
         preserveOrder = config.preserveOrder.value,
         strict = config.strict.value,
         noStaticErrors = config.noStaticErrors.value,
       ),
       storePos = if (config.yamlDebug.value) currentPos = _ else null,
-      parseCache,
       warnLogger
     )
 

diff --git a/sjsonnet/src/sjsonnet/Importer.scala b/sjsonnet/src/sjsonnet/Importer.scala
@@ -43,8 +43,8 @@ class CachedImporter(parent: Importer) extends Importer {
 
 class CachedResolver(
   parentImporter: Importer,
-  val parseCache: mutable.HashMap[(Path, String), Either[Error, (Expr, FileScope)]] = new mutable.HashMap
-) extends CachedImporter(parentImporter) {
+  val parseCache: ParseCache
+                    ) extends CachedImporter(parentImporter) {
 
   def parse(path: Path, txt: String)(implicit ev: EvalErrorScope): Either[Error, (Expr, FileScope)] = {
     parseCache.getOrElseUpdate((path, txt), {

diff --git a/sjsonnet/src/sjsonnet/Interpreter.scala b/sjsonnet/src/sjsonnet/Interpreter.scala
@@ -15,9 +15,9 @@ class Interpreter(extVars: Map[String, ujson.Value],
                   tlaVars: Map[String, ujson.Value],
                   wd: Path,
                   importer: Importer,
+                  val parseCache: ParseCache,
                   settings: Settings = Settings.default,
                   storePos: Position => Unit = null,
-                  val parseCache: mutable.HashMap[(Path, String), Either[Error, (Expr, FileScope)]] = new mutable.HashMap,
                   warnLogger: (String => Unit) = null,
                   ) { self =>
 

diff --git a/sjsonnet/src/sjsonnet/ParseCache.scala b/sjsonnet/src/sjsonnet/ParseCache.scala
@@ -0,0 +1,28 @@
+package sjsonnet
+
+import scala.collection.Iterator
+
+// Trait extended by JsonnetWorker (in universe) so that it can pass the cache based on Caffeine to main0 here
+trait ParseCache {
+  def getOrElseUpdate(key: (Path, String), defaultValue:  => Either[Error, (Expr, FileScope)]): Either[Error, (Expr, FileScope)]
+}
+
+// A default implementation based on a mutable HashMap. This implementation is not thread-safe.
+class DefaultParseCache extends ParseCache {
+  val cache = new scala.collection.mutable.HashMap[(Path, String), Either[Error, (Expr, FileScope)]]()
+
+  // parseCache.getOrElseUpdate((path, txt), {...})
+  override def getOrElseUpdate(key: (Path, String), defaultValue:  => Either[Error, (Expr, FileScope)]): Either[Error, (Expr, FileScope)] = {
+    cache.getOrElseUpdate(key, defaultValue)
+  }
+
+  // parseCache.valuesIterator.map(_.getOrElse(???)).map(_._1).toSeq
+  def valuesIterator: Iterator[Either[Error, (Expr, FileScope)]] = {
+    cache.valuesIterator
+  }
+
+  // parseCache.keySet.toIndexedSeq
+  def keySet: scala.collection.Set[(Path, String)] = {
+    cache.keySet
+  }
+}
diff --git a/sjsonnet/test/src-jvm-native/sjsonnet/ErrorTests.scala b/sjsonnet/test/src-jvm-native/sjsonnet/ErrorTests.scala
@@ -11,6 +11,7 @@ object ErrorTests extends TestSuite{
       Map(),
       OsPath(os.pwd),
       importer = sjsonnet.SjsonnetMain.resolveImport(Array.empty[Path]),
+      parseCache = new DefaultParseCache,
       warnLogger = (msg: String) => out.append(msg).append('\n'),
       settings = new Settings(noStaticErrors = noStaticErrors),
     )

diff --git a/sjsonnet/test/src-jvm-native/sjsonnet/FileTests.scala b/sjsonnet/test/src-jvm-native/sjsonnet/FileTests.scala
@@ -10,6 +10,7 @@ object FileTests extends TestSuite{
       Map("var1" -> "test", "var2" -> ujson.Obj("x" -> 1, "y" -> 2)),
       OsPath(testSuiteRoot),
       importer = sjsonnet.SjsonnetMain.resolveImport(Array(OsPath(testSuiteRoot))),
+      parseCache = new DefaultParseCache
     )
     interp.interpret(os.read(p), OsPath(p))
   }

diff --git a/sjsonnet/test/src-jvm-native/sjsonnet/MainTests.scala b/sjsonnet/test/src-jvm-native/sjsonnet/MainTests.scala
@@ -18,9 +18,9 @@ object MainTests extends TestSuite {
       val outF = File.createTempFile("sjsonnet", ".json")
       val out = new ByteArrayOutputStream()
       val pout = new PrintStream(out)
-      SjsonnetMain.main0(Array(source), collection.mutable.HashMap.empty, System.in, pout, System.err, os.pwd, None)
+      SjsonnetMain.main0(Array(source), new DefaultParseCache, System.in, pout, System.err, os.pwd, None)
       pout.flush()
-      SjsonnetMain.main0(Array("-o", outF.getAbsolutePath, source), collection.mutable.HashMap.empty, System.in, System.out, System.err, os.pwd, None)
+      SjsonnetMain.main0(Array("-o", outF.getAbsolutePath, source), new DefaultParseCache, System.in, System.out, System.err, os.pwd, None)
       val stdoutBytes = out.toByteArray
       val fileBytes = os.read(os.Path(outF)).getBytes
 
@@ -45,7 +45,7 @@ object MainTests extends TestSuite {
     val perr = new PrintStream(err, true, "UTF-8")
     val out = new ByteArrayOutputStream()
     val pout = new PrintStream(out, true, "UTF-8")
-    val res = SjsonnetMain.main0(args.toArray, collection.mutable.HashMap.empty, System.in, pout, perr, os.pwd, None)
+    val res = SjsonnetMain.main0(args.toArray, new DefaultParseCache, System.in, pout, perr, os.pwd, None)
     (res, new String(out.toByteArray, "UTF-8"), new String(err.toByteArray, "UTF-8"))
   }
 }
diff --git a/sjsonnet/test/src-jvm-native/sjsonnet/PrettyYamlRendererTests.scala b/sjsonnet/test/src-jvm-native/sjsonnet/PrettyYamlRendererTests.scala
@@ -11,6 +11,7 @@ object PrettyYamlRendererTests extends TestSuite{
       Map(),
       OsPath(testSuiteRoot),
       importer = sjsonnet.SjsonnetMain.resolveImport(Array(OsPath(testSuiteRoot))),
+      parseCache = new DefaultParseCache,
       storePos = if (comments) currentPos = _ else null
     )
     val res = interp.interpret0(

diff --git a/sjsonnet/test/src-jvm/sjsonnet/ErrorTestsJvmOnly.scala b/sjsonnet/test/src-jvm/sjsonnet/ErrorTestsJvmOnly.scala
@@ -10,6 +10,7 @@ object ErrorTestsJvmOnly extends TestSuite {
       Map(),
       OsPath(os.pwd),
       importer = sjsonnet.SjsonnetMain.resolveImport(Array.empty[Path]),
+      parseCache = new DefaultParseCache
     )
     interp.interpret(os.read(p), OsPath(p))
   }

diff --git a/sjsonnet/test/src-jvm/sjsonnet/Example.java b/sjsonnet/test/src-jvm/sjsonnet/Example.java
@@ -4,7 +4,7 @@ public class Example {
     public void example(){
         sjsonnet.SjsonnetMain.main0(
             new String[]{"foo.jsonnet"},
-            sjsonnet.SjsonnetMain.createParseCache(),
+            new DefaultParseCache(),
             System.in,
             System.out,
             System.err,

diff --git a/sjsonnet/test/src/sjsonnet/TestUtils.scala b/sjsonnet/test/src/sjsonnet/TestUtils.scala
@@ -7,6 +7,7 @@ object TestUtils {
       Map(),
       DummyPath(),
       Importer.empty,
+      parseCache = new DefaultParseCache,
       new Settings(preserveOrder = preserveOrder, strict = strict)
     ).interpret(s, DummyPath("(memory)")) match {
       case Right(x) => x