From af3816e7d0f1257419e59637d42ecc868148ceb5 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 | 2 +- .../src/mill/resolve/ResolveCore.scala | 6 +- main/src/mill/main/MainModule.scala | 19 +++- main/test/src/mill/main/MainModuleTests.scala | 107 ++++++++++++++++++ 5 files changed, 129 insertions(+), 8 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..0dcd0f076f4 100644 --- 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 index 31a7d632f39..6b05d5971ba 100644 --- 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 index 54421527574..9c63e11d33c 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,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 index 6f5b552ccdf..c3cd19aba4c 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 { @@ -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) + } + } } }