From 1b9f6470d0f18462fe670732752147186904f20b Mon Sep 17 00:00:00 2001 From: Alexandre Archambault Date: Wed, 18 Sep 2024 17:54:56 +0200 Subject: [PATCH] Clean workers too from the clean command --- main/define/src/mill/define/Segments.scala | 3 + main/eval/src/mill/eval/Evaluator.scala | 7 ++ main/src/mill/main/MainModule.scala | 21 +++- main/test/src/mill/main/MainModuleTests.scala | 111 ++++++++++++++++++ 4 files changed, 138 insertions(+), 4 deletions(-) diff --git a/main/define/src/mill/define/Segments.scala b/main/define/src/mill/define/Segments.scala index 82ab529fd32..e5cd4d14260 100644 --- 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 index 1a1d2853cb5..c694d3e02de 100644 --- a/main/eval/src/mill/eval/Evaluator.scala +++ b/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("should not happen") + } def disableCallgraphInvalidation: Boolean = false @deprecated( diff --git a/main/src/mill/main/MainModule.scala b/main/src/mill/main/MainModule.scala index 54421527574..b6b6fd3059f 100644 --- 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,25 @@ 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 + val closeableCleanedWorkers = workersToRemove + .flatMap(workerSegments => evaluator.mutableWorkerCache.remove(workerSegments)) + .collect { + case (_, Val(closeable: AutoCloseable)) => + closeable + } + closeableCleanedWorkers.foreach(_.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 index f27497680e3..312273876e6 100644 --- 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 { @@ -66,6 +68,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") { @@ -308,5 +359,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) + } + } } }