Clean workers too from the clean command

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 =>

main/eval/src/mill/eval/Evaluator.scala

@@ -21,7 +21,14 @@ trait Evaluator {
   def outPath: os.Path
   def externalOutPath: os.Path
   def pathsResolver: EvaluatorPathsResolver
+  // TODO In 0.13.0, workerCache should have the type of mutableWorkerCache,
+  // while the latter should be removed
   def workerCache: collection.Map[Segments, (Int, Val)]
+  private[mill] final def mutableWorkerCache: collection.mutable.Map[Segments, (Int, Val)] =
+    workerCache match {
+      case mut: collection.mutable.Map[Segments, (Int, Val)] => mut
+      case _ => sys.error("Evaluator#workerCache must be a mutable map")
+    }
   def disableCallgraphInvalidation: Boolean = false
 
   @deprecated(

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,21 @@ 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)) =>
+        for {
+          workerSegments <- evaluator.workerCache.keys.toList
+          if allSegments.exists(workerSegments.startsWith)
+          (_, Val(closeable: AutoCloseable)) <- evaluator.mutableWorkerCache.remove(workerSegments)
+        } {
+          closeable.close()
+        }
+
         val existing = paths.filter(p => os.exists(p))
         Target.log.debug(s"Cleaning ${existing.size} paths ...")
         existing.foreach(os.remove.all)

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 {
@@ -70,6 +72,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") {
@@ -313,5 +364,65 @@ 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)
+
+        val r4 = ev.evaluator.evaluate(Agg(workerModule.clean(ev.evaluator, "bazz[1].theWorker")))
+        assert(r4.failing.keyCount == 0)
+        assert(workers.size == 2)
+      }
+
+      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)
+      }
+    }
   }
 }