From a1eb82c87e9dd734ca02fcc1a57c08ec5ba8126f Mon Sep 17 00:00:00 2001 From: Alexandre Archambault Date: Thu, 26 Sep 2024 03:08:59 +0200 Subject: [PATCH] Clean workers too from the clean command (#3579) This PR ensures the `clean` command works for worker tasks, and allows users to dispose of workers by removing their corresponding metadata file on disk (`out/foo/theWorker.json` for worker `foo.theWorker`) - in that case, the worker instance is dropped upon first access to it. Fixes https://github.com/com-lihaoyi/mill/issues/3276 --------- Co-authored-by: Li Haoyi --- main/define/src/mill/define/Segments.scala | 3 + main/eval/src/mill/eval/Evaluator.scala | 7 + main/eval/src/mill/eval/GroupEvaluator.scala | 14 +- main/src/mill/main/MainModule.scala | 17 +- main/test/src/mill/main/MainModuleTests.scala | 161 ++++++++++++++++++ 5 files changed, 195 insertions(+), 7 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..42829650626 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("Evaluator#workerCache must be a mutable map") + } def disableCallgraphInvalidation: Boolean = false @deprecated( diff --git a/main/eval/src/mill/eval/GroupEvaluator.scala b/main/eval/src/mill/eval/GroupEvaluator.scala index dcdcf2e7695..b2566465007 100644 --- a/main/eval/src/mill/eval/GroupEvaluator.scala +++ b/main/eval/src/mill/eval/GroupEvaluator.scala @@ -149,7 +149,14 @@ private[mill] trait GroupEvaluator { val cached = loadCachedJson(logger, inputsHash, labelled, paths) - val upToDateWorker = loadUpToDateWorker(logger, inputsHash, labelled) + val upToDateWorker = loadUpToDateWorker( + logger, + inputsHash, + labelled, + forceDiscard = + // worker metadata file removed by user, let's recompute the worker + cached.isEmpty + ) upToDateWorker.map((_, inputsHash)) orElse cached.flatMap(_._2) match { case Some((v, hashCode)) => @@ -444,7 +451,8 @@ private[mill] trait GroupEvaluator { private def loadUpToDateWorker( logger: ColorLogger, inputsHash: Int, - labelled: Terminal.Labelled[_] + labelled: Terminal.Labelled[_], + forceDiscard: Boolean ): Option[Val] = { labelled.task.asWorker .flatMap { w => @@ -454,7 +462,7 @@ private[mill] trait GroupEvaluator { } .flatMap { case (cachedHash, upToDate) - if cachedHash == workerCacheHash(inputsHash) => + if cachedHash == workerCacheHash(inputsHash) && !forceDiscard => Some(upToDate) // worker cached and up-to-date case (_, Val(obsolete: AutoCloseable)) => diff --git a/main/src/mill/main/MainModule.scala b/main/src/mill/main/MainModule.scala index 54421527574..109c37786d1 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,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) diff --git a/main/test/src/mill/main/MainModuleTests.scala b/main/test/src/mill/main/MainModuleTests.scala index 82d0790c7bc..8e7a05d9c0d 100644 --- a/main/test/src/mill/main/MainModuleTests.scala +++ b/main/test/src/mill/main/MainModuleTests.scala @@ -3,12 +3,15 @@ package mill.main import mill.api.{PathRef, Result, Val} import mill.{Agg, T, Task} import mill.define.{Cross, Discover, Module} +import mill.main.client.OutFiles import mill.testkit.UnitTester import mill.testkit.TestBaseModule 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 { @@ -72,6 +75,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") { @@ -317,5 +369,114 @@ 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-target via rm") { + val workers = new mutable.HashSet[TestWorker] + val workerModule = new WorkerModule(workers) + val ev = UnitTester(workerModule, null) + + ev.evaluator.evaluate(Agg(workerModule.foo.theWorker)) + .ensuring(_.failing.keyCount == 0) + assert(workers.size == 1) + + val originalFooWorker = workers.head + + ev.evaluator.evaluate(Agg(workerModule.bar.theWorker)) + .ensuring(_.failing.keyCount == 0) + assert(workers.size == 2) + assert(workers.exists(_ eq originalFooWorker)) + + val originalBarWorker = workers.filter(_ ne originalFooWorker).head + + ev.evaluator.evaluate(Agg(workerModule.foo.theWorker)) + .ensuring(_.failing.keyCount == 0) + assert(workers.size == 2) + assert(workers.exists(_ eq originalFooWorker)) + + ev.evaluator.evaluate(Agg(workerModule.bar.theWorker)) + .ensuring(_.failing.keyCount == 0) + assert(workers.size == 2) + assert(workers.exists(_ eq originalBarWorker)) + + val outDir = os.Path(OutFiles.out, workerModule.millSourcePath) + + assert(!originalFooWorker.closed) + os.remove(outDir / "foo/theWorker.json") + + ev.evaluator.evaluate(Agg(workerModule.foo.theWorker)) + .ensuring(_.failing.keyCount == 0) + assert(workers.size == 2) + assert(!workers.exists(_ eq originalFooWorker)) + assert(originalFooWorker.closed) + + assert(!originalBarWorker.closed) + os.remove(outDir / "bar/theWorker.json") + + ev.evaluator.evaluate(Agg(workerModule.bar.theWorker)) + .ensuring(_.failing.keyCount == 0) + assert(workers.size == 2) + assert(!workers.exists(_ eq originalBarWorker)) + assert(originalBarWorker.closed) + } + + 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) + } + } } }