Clean workers too from the clean command

com-lihaoyi · Sep 18, 2024 · af3816e · af3816e
1 parent b5e04d6
commit af3816e
Show file tree

Hide file tree

Showing 5 changed files with 129 additions and 8 deletions.
diff --git a/main/define/src/mill/define/Segments.scala b/main/define/src/mill/define/Segments.scala
@@ -14,6 +14,9 @@ case class Segments private (value: Seq[Segment]) {
   def ++(other: Seq[Segment]): Segments = Segments(value ++ other)
   def ++(other: Segments): Segments = Segments(value ++ other.value)
 
+  def startsWith(prefix: Segments): Boolean =
+    value.startsWith(prefix.value)
+
   def parts: List[String] = value.toList match {
     case Nil => Nil
     case Segment.Label(head) :: rest =>

diff --git a/main/eval/src/mill/eval/Evaluator.scala b/main/eval/src/mill/eval/Evaluator.scala
@@ -21,7 +21,7 @@ trait Evaluator {
   def outPath: os.Path
   def externalOutPath: os.Path
   def pathsResolver: EvaluatorPathsResolver
-  def workerCache: collection.Map[Segments, (Int, Val)]
+  def workerCache: collection.mutable.Map[Segments, (Int, Val)]
   def disableCallgraphInvalidation: Boolean = false
 
   @deprecated(

diff --git a/main/resolve/src/mill/resolve/ResolveCore.scala b/main/resolve/src/mill/resolve/ResolveCore.scala
@@ -376,8 +376,8 @@ private object ResolveCore {
       }
     }
 
-    val targets = Reflect
-      .reflect(cls, classOf[Target[_]], namePred, noParams = true)
+    val namedTasks = Reflect
+      .reflect(cls, classOf[NamedTask[_]], namePred, noParams = true)
       .map { m =>
         Resolved.Target(Segments.labels(decode(m.getName))) ->
           None
@@ -388,7 +388,7 @@ private object ResolveCore {
       .map(m => decode(m.getName))
       .map { name => Resolved.Command(Segments.labels(name)) -> None }
 
-    modulesOrErr.map(_ ++ targets ++ commands)
+    modulesOrErr.map(_ ++ namedTasks ++ commands)
   }
 
   def notFoundResult(

diff --git a/main/src/mill/main/MainModule.scala b/main/src/mill/main/MainModule.scala
@@ -2,7 +2,7 @@ package mill.main
 
 import java.util.concurrent.LinkedBlockingQueue
 import mill.define.{BaseModule0, Command, NamedTask, Segments, Target, Task}
-import mill.api.{Ctx, Logger, PathRef, Result}
+import mill.api.{Ctx, Logger, PathRef, Result, Val}
 import mill.eval.{Evaluator, EvaluatorPaths, Terminal}
 import mill.resolve.{Resolve, SelectMode}
 import mill.resolve.SelectMode.Separated
@@ -328,14 +328,14 @@ trait MainModule extends BaseModule0 {
 
     val pathsToRemove =
       if (targets.isEmpty)
-        Right(os.list(rootDir).filterNot(keepPath))
+        Right((os.list(rootDir).filterNot(keepPath), List(mill.define.Segments())))
       else
         mill.resolve.Resolve.Segments.resolve(
           evaluator.rootModule,
           targets,
           SelectMode.Multi
         ).map { ts =>
-          ts.flatMap { segments =>
+          val allPaths = ts.flatMap { segments =>
             val evPaths = EvaluatorPaths.resolveDestPaths(rootDir, segments)
             val paths = Seq(evPaths.dest, evPaths.meta, evPaths.log)
             val potentialModulePath = rootDir / EvaluatorPaths.makeSegmentStrings(segments)
@@ -348,12 +348,23 @@ trait MainModule extends BaseModule0 {
               paths :+ potentialModulePath
             } else paths
           }
+          (allPaths, ts)
         }
 
     pathsToRemove match {
       case Left(err) =>
         Result.Failure(err)
-      case Right(paths) =>
+      case Right((paths, allSegments)) =>
+        val workersToRemove = evaluator.workerCache
+          .keysIterator
+          .filter(workerSegments => allSegments.exists(workerSegments.startsWith))
+          .toVector
+        for {
+          workerSegments <- workersToRemove
+          (_, Val(closable: AutoCloseable)) <- evaluator.workerCache.remove(workerSegments)
+        }
+          closable.close()
+
         val existing = paths.filter(p => os.exists(p))
         Target.log.debug(s"Cleaning ${existing.size} paths ...")
         existing.foreach(os.remove.all)

diff --git a/main/test/src/mill/main/MainModuleTests.scala b/main/test/src/mill/main/MainModuleTests.scala
@@ -9,6 +9,8 @@ import utest.{TestSuite, Tests, assert, test}
 
 import java.io.{ByteArrayOutputStream, PrintStream}
 
+import scala.collection.mutable
+
 object MainModuleTests extends TestSuite {
 
   object mainModule extends TestBaseModule with MainModule {
@@ -60,6 +62,55 @@ object MainModuleTests extends TestSuite {
     }
   }
 
+  class TestWorker(val name: String, workers: mutable.HashSet[TestWorker]) extends AutoCloseable {
+
+    workers.synchronized {
+      workers.add(this)
+    }
+
+    var closed = false
+    def close(): Unit =
+      if (!closed) {
+        workers.synchronized {
+          workers.remove(this)
+        }
+        closed = true
+      }
+
+    override def toString(): String =
+      s"TestWorker($name)@${Integer.toHexString(System.identityHashCode(this))}"
+  }
+
+  class WorkerModule(workers: mutable.HashSet[TestWorker]) extends TestBaseModule with MainModule {
+
+    trait Cleanable extends Module {
+      def theWorker = Task.Worker {
+        new TestWorker("shared", workers)
+      }
+    }
+
+    object foo extends Cleanable {
+      object sub extends Cleanable
+    }
+    object bar extends Cleanable {
+      def theWorker = Task.Worker {
+        new TestWorker("bar", workers)
+      }
+    }
+    object bazz extends Cross[Bazz]("1", "2", "3")
+    trait Bazz extends Cleanable with Cross.Module[String]
+
+    def all = Task {
+      foo.theWorker()
+      bar.theWorker()
+      bazz("1").theWorker()
+      bazz("2").theWorker()
+      bazz("3").theWorker()
+
+      ()
+    }
+  }
+
   override def tests: Tests = Tests {
 
     test("inspect") {
@@ -274,5 +325,61 @@ object MainModuleTests extends TestSuite {
         )
       }
     }
+
+    test("cleanWorker") {
+      test("all") {
+        val workers = new mutable.HashSet[TestWorker]
+        val workerModule = new WorkerModule(workers)
+        val ev = UnitTester(workerModule, null)
+
+        val r1 = ev.evaluator.evaluate(Agg(workerModule.all))
+        assert(r1.failing.keyCount == 0)
+        assert(workers.size == 5)
+
+        val r2 = ev.evaluator.evaluate(Agg(workerModule.clean(ev.evaluator)))
+        assert(r2.failing.keyCount == 0)
+        assert(workers.isEmpty)
+      }
+
+      test("single-target") {
+        val workers = new mutable.HashSet[TestWorker]
+        val workerModule = new WorkerModule(workers)
+        val ev = UnitTester(workerModule, null)
+
+        val r1 = ev.evaluator.evaluate(Agg(workerModule.all))
+        assert(r1.failing.keyCount == 0)
+        assert(workers.size == 5)
+
+        val r2 = ev.evaluator.evaluate(Agg(workerModule.clean(ev.evaluator, "foo.theWorker")))
+        assert(r2.failing.keyCount == 0)
+        assert(workers.size == 4)
+
+        val r3 = ev.evaluator.evaluate(Agg(workerModule.clean(ev.evaluator, "bar.theWorker")))
+        assert(r3.failing.keyCount == 0)
+        assert(workers.size == 3)
+      }
+
+      test("single-module") {
+        val workers = new mutable.HashSet[TestWorker]
+        val workerModule = new WorkerModule(workers)
+        val ev = UnitTester(workerModule, null)
+
+        val r1 = ev.evaluator.evaluate(Agg(workerModule.all))
+        assert(r1.failing.keyCount == 0)
+        assert(workers.size == 5)
+
+        val r2 = ev.evaluator.evaluate(Agg(workerModule.clean(ev.evaluator, "foo")))
+        assert(r2.failing.keyCount == 0)
+        assert(workers.size == 4)
+
+        val r3 = ev.evaluator.evaluate(Agg(workerModule.clean(ev.evaluator, "bar")))
+        assert(r3.failing.keyCount == 0)
+        assert(workers.size == 3)
+
+        val r4 = ev.evaluator.evaluate(Agg(workerModule.clean(ev.evaluator, "bazz[1]")))
+        assert(r4.failing.keyCount == 0)
+        assert(workers.size == 2)
+      }
+    }
   }
 }