From b4a0bfcd0390e8729dfc94d99689165555157c8c Mon Sep 17 00:00:00 2001
From: Li Haoyi <haoyi.sg@gmail.com>
Date: Wed, 9 Oct 2024 15:05:58 +0200
Subject: [PATCH] Rebootstrap on top of master to make use of updated
 `GraphvizTools` subprocess API (#3694)

Also optimize `GraphvizTools` a bunch to re-use the javascript engine
when batch processing lots of diagrams
---
 .config/mill-version                             |  2 +-
 .github/workflows/run-tests.yml                  |  2 +-
 docs/package.mill                                | 16 +++++++++++-----
 .../src/mill/main/graphviz/GraphvizTools.scala   | 14 ++++++++++++--
 4 files changed, 25 insertions(+), 9 deletions(-)

diff --git a/.config/mill-version b/.config/mill-version
index 264ae60564b..02e50870961 100644
--- a/.config/mill-version
+++ b/.config/mill-version
@@ -1 +1 @@
-0.12.0-RC3-26-d3afbf
+0.12.0-RC3-31-9b84ae
\ No newline at end of file
diff --git a/.github/workflows/run-tests.yml b/.github/workflows/run-tests.yml
index 7ca644d904b..723d0f89354 100644
--- a/.github/workflows/run-tests.yml
+++ b/.github/workflows/run-tests.yml
@@ -44,7 +44,7 @@ jobs:
       - uses: actions/checkout@v4
         with: { fetch-depth: 0 }
 
-      - run: ./mill -i docs.githubPages
+      - run: ./mill installLocal && ./target/mill-release -i docs.githubPages
 
   linux:
     needs: build-linux
diff --git a/docs/package.mill b/docs/package.mill
index 3db8a792a7b..710b5e4315f 100644
--- a/docs/package.mill
+++ b/docs/package.mill
@@ -90,8 +90,8 @@ object `package` extends RootModule {
       createFolders = true
     )
 
-    Graphviz.useEngine(new AbstractJsGraphvizEngine(true, () => new mill.main.graphviz.V8JavascriptEngine()) {})
-
+    // Walk all files al render graphviz templates ourselves because the only Antora graphviz
+    // plugin (Kroki) relies on an online web service that is super slow and flaky
     def walkAllFiles(inputs: Map[(os.Path, Int), String]): Map[(os.Path, Int), String] = {
       val output = collection.mutable.Map.empty[(os.Path, Int), String]
       for (p <- os.walk(T.dest) if p.ext == "adoc"){
@@ -125,14 +125,20 @@ object `package` extends RootModule {
     val diagrams = walkAllFiles(Map())
     // Batch the rendering so later it can be done in one call to a single subprocess,
     // minimizing per-subprocess overhead needed to spawn them over and over
-    val renderedDiagrams = diagrams
-      .map{case (k, s) => (k, Graphviz.fromString(s).render(Format.SVG).toString)}
+    val orderedDiagrams = diagrams.toSeq.map{case ((p, i), s) => (p, i, os.temp(s), os.temp.dir())}
+
+    mill.util.Jvm.runSubprocess(
+      "mill.main.graphviz.GraphvizTools",
+      mill.main.VisualizeModule.classpath().map(_.path),
+      mainArgs = orderedDiagrams.map{case (p, i, src, dest) => s"$src;$dest;svg"}
+    )
 
-    walkAllFiles(renderedDiagrams)
+    walkAllFiles(orderedDiagrams.map{case (p, i, src, dest) => ((p, i), os.read(dest / "out.svg"))}.toMap)
 
     PathRef(T.dest)
   }
 
+
   def supplementalFiles = T.source(millSourcePath / "supplemental-ui")
 
   /**
diff --git a/main/graphviz/src/mill/main/graphviz/GraphvizTools.scala b/main/graphviz/src/mill/main/graphviz/GraphvizTools.scala
index cb416eb7fc5..d8b6b26d356 100644
--- a/main/graphviz/src/mill/main/graphviz/GraphvizTools.scala
+++ b/main/graphviz/src/mill/main/graphviz/GraphvizTools.scala
@@ -8,6 +8,7 @@ import org.slf4j.LoggerFactory
 import org.slf4j.Logger
 
 import java.util.concurrent.Executors
+import guru.nidi.graphviz.engine.{Format, Graphviz}
 import scala.concurrent.{Await, ExecutionContext, Future, duration}
 
 object GraphvizTools {
@@ -15,6 +16,17 @@ object GraphvizTools {
   def main(args: Array[String]): Unit = {
     val executor = Executors.newFixedThreadPool(Runtime.getRuntime.availableProcessors())
 
+    val threadLocalJsEngines =
+      new java.util.concurrent.ConcurrentHashMap[Thread, V8JavascriptEngine]()
+    Graphviz.useEngine(
+      new AbstractJsGraphvizEngine(
+        true,
+        () => {
+          threadLocalJsEngines.putIfAbsent(Thread.currentThread(), new V8JavascriptEngine())
+          threadLocalJsEngines.get(Thread.currentThread())
+        }
+      ) {}
+    )
     try {
       implicit val ec: ExecutionContext = ExecutionContext.fromExecutor(executor)
       val futures =
@@ -22,9 +34,7 @@ object GraphvizTools {
           val Array(src, dest0, commaSepExtensions) = arg.split(";")
           val extensions = commaSepExtensions.split(',')
           val dest = os.Path(dest0)
-          import guru.nidi.graphviz.engine.{Format, Graphviz}
 
-          Graphviz.useEngine(new AbstractJsGraphvizEngine(true, () => new V8JavascriptEngine()) {})
           val gv = Graphviz.fromFile(new java.io.File(src)).totalMemory(128 * 1024 * 1024)
 
           val outputs = extensions