move graphviz interactions onto a single worker thread to satisfy J2V8 limitation

3 files changed, 52 insertions, 23 deletions

diff --git a/main/src/mill/main/MainModule.scala b/main/src/mill/main/MainModule.scala
index 00ddcde1..367f241c 100644
--- a/main/src/mill/main/MainModule.scala
+++ b/main/src/mill/main/MainModule.scala
@@ -221,6 +221,15 @@ trait MainModule extends mill.Module{
   }
 
   def visualize(evaluator: Evaluator[Any], targets: String*) = mill.T.command{
-    VisualizeModule.run(evaluator, targets:_*)
+    val resolved = RunScript.resolveTasks(
+      mill.main.ResolveTasks, evaluator, targets, multiSelect = true
+    )
+    resolved match{
+      case Left(err) => Result.Failure(err)
+      case Right(rs) =>
+        val (in, out) = mill.main.VisualizeModule.worker()
+        in.put((rs, T.ctx().dest))
+        out.take()
+    }
   }
 }
diff --git a/main/src/mill/main/RunScript.scala b/main/src/mill/main/RunScript.scala
index 50cbb213..648f0447 100644
--- a/main/src/mill/main/RunScript.scala
+++ b/main/src/mill/main/RunScript.scala
@@ -217,8 +217,16 @@ object RunScript{
         }
         val fss = fs.map{
           case Result.Exception(t, outerStack) =>
-            t.toString +
-            t.getStackTrace.dropRight(outerStack.value.length).map("\n    " + _).mkString
+            var current = List(t)
+            while(current.head.getCause != null){
+              current = current.head.getCause :: current
+            }
+            current.reverse
+              .flatMap( ex =>
+                Seq(ex.toString) ++
+                ex.getStackTrace.dropRight(outerStack.value.length).map("    " + _)
+              )
+              .mkString("\n")
           case Result.Failure(t, _) => t
         }
         s"$ks ${fss.mkString(", ")}"
diff --git a/main/src/mill/main/VisualizeModule.scala b/main/src/mill/main/VisualizeModule.scala
index e939bd86..8e022a33 100644
--- a/main/src/mill/main/VisualizeModule.scala
+++ b/main/src/mill/main/VisualizeModule.scala
@@ -1,13 +1,16 @@
 package mill.main
 
+import java.util.concurrent.LinkedBlockingQueue
+
 import ammonite.ops.Path
 import coursier.Cache
 import coursier.core.Repository
 import coursier.maven.MavenRepository
-import mill.T
+import mill.{T, eval}
 import mill.define.{Discover, ExternalModule}
 import mill.eval.{Evaluator, PathRef, Result}
 
+
 object VisualizeModule extends ExternalModule with VisualizeModule {
   def repositories = Seq(
     Cache.ivy2Local,
@@ -24,27 +27,36 @@ trait VisualizeModule extends mill.define.TaskModule{
   def classpath = T{
     mill.modules.Util.millProjectModule("MILL_GRAPHVIZ", "mill-main-graphviz", repositories)
   }
+
   /**
-    * Given a set of tasks, prints out the execution plan of what tasks will be
-    * executed in what order, without actually executing them.
+    * The J2V8-based Graphviz library has a limitation that it can only ever
+    * be called from a single thread. Since Mill forks off a new thread every
+    * time you execute something, we need to keep around a worker thread that
+    * everyone can use to call into Graphviz, which the Mill execution threads
+    * can communicate via in/out queues.
     */
-  def run(evaluator: Evaluator[Any], targets: String*) = mill.T.command{
-    val resolved = RunScript.resolveTasks(
-      mill.main.ResolveTasks, evaluator, targets, multiSelect = true
-    )
-    resolved match{
-      case Left(err) => Result.Failure(err)
-      case Right(rs) =>
-        Result.Success(
-          mill.modules.Jvm.inprocess(classpath().map(_.path), false, isolated = false, cl => {
-            cl.loadClass("mill.main.graphviz.GraphvizTools")
-              .getMethod("apply", classOf[Seq[_]], classOf[Path])
-              .invoke(null, rs, T.ctx().dest)
-              .asInstanceOf[Seq[PathRef]]
-          })
-        )
+  def worker = T.worker{
+    val in = new LinkedBlockingQueue[(Seq[_], Path)]()
+    val out = new LinkedBlockingQueue[Result[Seq[PathRef]]]()
 
-    }
+    val cl = mill.util.ClassLoader.create(
+      classpath().map(_.path.toNIO.toUri.toURL).toVector,
+      getClass.getClassLoader
+    )
+    val visualizeThread = new java.lang.Thread(() =>
+      while(true){
+        val res = Result.create{
+          val (tasks, dest) = in.take()
+          cl.loadClass("mill.main.graphviz.GraphvizTools")
+            .getMethod("apply", classOf[Seq[_]], classOf[Path])
+            .invoke(null, tasks, dest)
+            .asInstanceOf[Seq[PathRef]]
+        }
+        out.put(res)
+      }
+    )
+    visualizeThread.setDaemon(true)
+    visualizeThread.start()
+    (in, out)
   }
-
 }