Let VizWorker extends AutoCloseable

lefou · Mar 16, 2022 · a5d1899 · a5d1899
1 parent 9de08bb
commit a5d1899
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 10 deletions.
diff --git a/main/src/mill/main/MainModule.scala b/main/src/mill/main/MainModule.scala
@@ -7,6 +7,7 @@ import mill.define.{Command, NamedTask, Task}
 import mill.eval.{Evaluator, EvaluatorPaths}
 import mill.util.{PrintLogger, Watched}
 import mill.define.SelectMode
+import mill.main.VisualizeModule.VizWorker
 import pprint.{Renderer, Truncated}
 import ujson.Value
 
@@ -380,11 +381,6 @@ trait MainModule extends mill.Module {
     System.exit(0)
   }
 
-  private type VizWorker = (
-      LinkedBlockingQueue[(scala.Seq[_], scala.Seq[_], os.Path)],
-      LinkedBlockingQueue[Result[scala.Seq[PathRef]]]
-  )
-
   private def visualize0(
       evaluator: Evaluator,
       targets: Seq[String],
@@ -396,9 +392,8 @@ trait MainModule extends mill.Module {
         rs: List[NamedTask[Any]],
         allRs: List[NamedTask[Any]]
     ): Result[Seq[PathRef]] = {
-      val (in, out) = vizWorker
-      in.put((rs, allRs, ctx.dest))
-      out.take()
+      vizWorker.in.put((rs, allRs, ctx.dest))
+      vizWorker.out.take()
     }
 
     RunScript.resolveTasks(

diff --git a/main/src/mill/main/VisualizeModule.scala b/main/src/mill/main/VisualizeModule.scala
@@ -1,15 +1,24 @@
 package mill.main
 
 import java.util.concurrent.LinkedBlockingQueue
-
 import coursier.LocalRepositories
 import coursier.core.Repository
 import coursier.maven.MavenRepository
 import mill.T
 import mill.define.{Discover, ExternalModule}
 import mill.api.{PathRef, Result}
+import mill.main.VisualizeModule.VizWorker
 
 object VisualizeModule extends ExternalModule with VisualizeModule {
+
+  private[main] class VizWorker(
+      val in: LinkedBlockingQueue[(scala.Seq[_], scala.Seq[_], os.Path)],
+      val out: LinkedBlockingQueue[Result[scala.Seq[PathRef]]],
+      closer: () => Unit = () => {}
+  ) extends AutoCloseable {
+    override def close(): Unit = closer()
+  }
+
   def repositories = Seq(
     LocalRepositories.ivy2Local,
     MavenRepository("https://repo1.maven.org/maven2"),
@@ -19,6 +28,7 @@ object VisualizeModule extends ExternalModule with VisualizeModule {
   implicit def millScoptEvaluatorReads[T] = new mill.main.EvaluatorScopt[T]()
   lazy val millDiscover = Discover[this.type]
 }
+
 trait VisualizeModule extends mill.define.TaskModule {
   def repositories: Seq[Repository]
   def defaultCommandName() = "run"
@@ -55,6 +65,12 @@ trait VisualizeModule extends mill.define.TaskModule {
     )
     visualizeThread.setDaemon(true)
     visualizeThread.start()
-    (in, out)
+    new VizWorker(
+      in,
+      out,
+      () => {
+        visualizeThread.interrupt()
+      }
+    )
   }
 }