.

Author: lihaoyi

core/api/daemon/src/mill/api/daemon/internal/EvaluatorApi.scala

@@ -56,7 +56,7 @@ trait EvaluatorApi extends AutoCloseable {
   private[mill] def withIsFinalDepth(isFinalDepth: Boolean): EvaluatorApi = this
 }
 object EvaluatorApi {
-  final case class SelectiveReuseDecision(reusable: Boolean, nextMetadata: String)
+  private[mill] final case class SelectiveReuseDecision(reusable: Boolean, nextMetadata: String)
 
   trait Result[T] {
     def watchable: Seq[Watchable]

core/constants/src/mill/constants/EnvVars.java

@@ -34,7 +34,7 @@ public class EnvVars {
   public static final String MILL_BSP_OUTPUT_DIR = "MILL_BSP_OUTPUT_DIR";
 
   /**
-   * If set to "1", Mill will re-use the regular out/ folder instead of
+   * If set to "1", Mill will re-use the regular @{Link OutFiles#out} folder instead of
    * using a separate one for BSP output.
    */
   public static final String MILL_NO_SEPARATE_BSP_OUTPUT_DIR = "MILL_NO_SEPARATE_BSP_OUTPUT_DIR";

core/eval/src/mill/eval/EvaluatorImpl.scala

@@ -319,6 +319,7 @@ final class EvaluatorImpl(
       selectiveExecution: Boolean = false
   ): Evaluator.Result[T] = {
     val selectiveExecutionEnabled = selectiveExecution && !tasks.exists(_.isExclusiveCommand)
+
     val (selectedTasks, selectiveResults, maybeNewMetadata) =
       if (!selectiveExecutionEnabled) (tasks, Map.empty, None)
       else {

core/exec/src/mill/exec/Execution.scala

@@ -45,7 +45,7 @@ case class Execution(
 ) extends GroupExecution with AutoCloseable {
 
   // Track nesting depth of executeTasks calls to only show final status on outermost call
-  private val executionNestingDepth = new AtomicInteger(0)
+  val executionNestingDepth = new AtomicInteger(0)
 
   // Lazily computed worker dependency graph, cached for the duration of the execution. It's
   // ok to take a snapshot of the cache, since the workerCache entries we may want to remove
@@ -126,70 +126,6 @@ case class Execution(
 
   def withIsFinalDepth(newIsFinalDepth: Boolean) = this.copy(isFinalDepth = newIsFinalDepth)
 
-  private def newDownstreamTracker(
-      indexToTerminal: Array[Task[?]],
-      interGroupDeps: Map[Task[?], Seq[Task[?]]]
-  ): Execution.LeaseTracker = {
-    val retainedLeasesByTask =
-      new ConcurrentHashMap[Task[?], java.util.concurrent.ConcurrentLinkedQueue[
-        LauncherLocking.Lease
-      ]]()
-    val pendingCount = new ConcurrentHashMap[Task[?], AtomicInteger]()
-    for (t <- indexToTerminal) {
-      pendingCount.put(t, new AtomicInteger(0))
-      retainedLeasesByTask.put(
-        t,
-        new java.util.concurrent.ConcurrentLinkedQueue[LauncherLocking.Lease]()
-      )
-    }
-    for ((_, deps) <- interGroupDeps; dep <- deps) {
-      val c = pendingCount.get(dep)
-      if (c != null) c.incrementAndGet()
-    }
-
-    def releaseLeasesFor(task: Task[?]): Unit = {
-      val q = retainedLeasesByTask.remove(task)
-      if (q != null) {
-        var lease = q.poll()
-        while (lease != null) {
-          try lease.close()
-          catch { case _: Throwable => () }
-          lease = q.poll()
-        }
-      }
-    }
-
-    new Execution.LeaseTracker {
-      override def retain(task: Task[?], lease: LauncherLocking.Lease): Unit = {
-        val q = retainedLeasesByTask.get(task)
-        if (q != null) q.add(lease)
-        else
-          try lease.close()
-          catch { case _: Throwable => () }
-      }
-
-      override def onCompleted(terminal: Task[?]): Unit = {
-        for (dep <- interGroupDeps.getOrElse(terminal, Nil)) {
-          val c = pendingCount.get(dep)
-          if (c != null && c.decrementAndGet() == 0) releaseLeasesFor(dep)
-        }
-        val ownCount = pendingCount.get(terminal)
-        if (ownCount != null && ownCount.get() == 0) releaseLeasesFor(terminal)
-      }
-
-      override def drain(): Unit = {
-        import scala.jdk.CollectionConverters.*
-        retainedLeasesByTask.values().asScala
-          .flatMap(q => Iterator.continually(q.poll()).takeWhile(_ != null))
-          .foreach { lease =>
-            try lease.close()
-            catch { case _: Throwable => () }
-          }
-        retainedLeasesByTask.clear()
-      }
-    }
-  }
-
   /**
    * @param goals The tasks that need to be evaluated
    * @param reporter A function that will accept a module id and provide a listener for build problems in that module
@@ -208,7 +144,9 @@ case class Execution(
       PathRef.validatedPaths.withValue(new PathRef.ValidatedPaths()) {
         execute0(goals, logger, reporter, testReporter, serialCommandExec)
       }
-    } finally executionNestingDepth.decrementAndGet()
+    } finally {
+      executionNestingDepth.decrementAndGet()
+    }
   }
 
   private def execute0(
@@ -283,7 +221,7 @@ case class Execution(
           downstreamEdges.getOrElse(t, Set())
         )
 
-      val tracker = newDownstreamTracker(indexToTerminal, interGroupDeps)
+      val tracker = new Execution.LeaseTracker(indexToTerminal, interGroupDeps)
       def onTerminalCompleted(t: Task[?]): Unit = tracker.onCompleted(t)
       try {
 
@@ -379,13 +317,13 @@ case class Execution(
                         leaseTracker = tracker
                       )
 
+                    // Count new failures - if there are upstream failures, tasks should be skipped, not failed
                       val newFailures = res.newResults.values.count(r => r.asFailing.isDefined)
 
-                      // Count new failures: tasks with upstream failures should be skipped, not failed.
                       rootFailedCount.addAndGet(newFailures)
                       completedCount.incrementAndGet()
 
-                      // Always show the completed count in the header after a task finishes.
+                    // Always show completed count in header after task finishes
                       logger.prompt.setPromptHeaderPrefix(formatHeaderPrefix())
 
                       if (failFast && res.newResults.values.exists(_.asSuccess.isEmpty))
@@ -411,14 +349,14 @@ case class Execution(
                     }
                   }
                 } catch {
-                  case e: mill.api.daemon.StopWithResponse[?] =>
-                    // Let StopWithResponse propagate: it is a controlled shutdown signal.
-                    throw e
+                // Let StopWithResponse propagate - it's a controlled shutdown signal
+                case e: mill.api.daemon.StopWithResponse[?] => throw e
+                // Wrapping the fatal error in a non-fatal exception, so it would be caught by Scala's Future
+                // infrastructure, rather than silently terminating the future and leaving downstream Awaits hanging.
                   case e: Throwable if !mill.api.daemon.internal.NonFatal(e) =>
-                    // Wrap fatal errors so Scala's Future machinery reports them instead of
-                    // silently terminating the future and leaving downstream Awaits hanging.
                     val nonFatal = new Exception(s"fatal exception occurred: $e", e)
-                    // Preserve the original stack trace, since that points at the real failure.
+                  // Set the stack trace of the non-fatal exception to the original exception's stack trace
+                  // as it actually indicates the location of the error.
                     nonFatal.setStackTrace(e.getStackTrace)
                     throw nonFatal
                 } finally {
@@ -429,17 +367,17 @@ case class Execution(
           }
 
           // Make sure we wait for all tasks from this batch to finish before starting the next
-          // one, so we don't mix up exclusive and non-exclusive tasks running at the same time.
+        // one, so we don't mix up exclusive and non-exclusive tasks running at the same time
           terminals.map(t => (t, Await.result(futures(t), duration.Duration.Inf)))
         }
 
-        // Run all non-command tasks according to the configured thread count,
-        // but run exclusive commands in linear order.
         val (nonExclusiveTasks, leafExclusiveCommands) = indexToTerminal.partition {
           case t: Task.Named[_] => !downstreamOfExclusive.contains(t)
           case _ => !serialCommandExec
         }
 
+      // Run all non-command tasks according to the threads
+      // given but run the commands in linear order
         val nonExclusiveResults = evaluateTerminals(nonExclusiveTasks, exclusive = false)
 
         val exclusiveResults = evaluateTerminals(leafExclusiveCommands, exclusive = true)
@@ -511,10 +449,61 @@ case class Execution(
 
 object Execution {
 
-  trait LeaseTracker {
-    def retain(task: Task[?], lease: LauncherLocking.Lease): Unit
-    def onCompleted(terminal: Task[?]): Unit
-    def drain(): Unit
+
+  /**
+   * Tracks per-task read leases on the workspace lock and releases them once
+   * every downstream task that depends on the holder has completed.
+   */
+  final class LeaseTracker(
+      indexToTerminal: Array[Task[?]],
+      interGroupDeps: Map[Task[?], Seq[Task[?]]]
+  ) {
+    final class State {
+      val pending = new AtomicInteger(0)
+      val leases = new java.util.concurrent.ConcurrentLinkedQueue[LauncherLocking.Lease]()
+    }
+
+    val states = new ConcurrentHashMap[Task[?], State]()
+    for (t <- indexToTerminal) states.put(t, new State)
+    for ((_, deps) <- interGroupDeps; dep <- deps) {
+      val s = states.get(dep)
+      if (s != null) s.pending.incrementAndGet()
+    }
+
+    def closeQuietly(lease: LauncherLocking.Lease): Unit =
+      try lease.close()
+      catch { case _: Throwable => () }
+
+    def releaseLeasesFor(task: Task[?]): Unit = {
+      val s = states.remove(task)
+      if (s != null) {
+        var lease = s.leases.poll()
+        while (lease != null) {
+          closeQuietly(lease)
+          lease = s.leases.poll()
+        }
+      }
+    }
+
+    def retain(task: Task[?], lease: LauncherLocking.Lease): Unit = {
+      val s = states.get(task)
+      if (s != null) s.leases.add(lease)
+      else closeQuietly(lease)
+    }
+
+    def onCompleted(terminal: Task[?]): Unit = {
+      for (dep <- interGroupDeps.getOrElse(terminal, Nil)) {
+        val s = states.get(dep)
+        if (s != null && s.pending.decrementAndGet() == 0) releaseLeasesFor(dep)
+      }
+      val own = states.get(terminal)
+      if (own != null && own.pending.get() == 0) releaseLeasesFor(terminal)
+    }
+
+    def drain(): Unit = {
+      import scala.jdk.CollectionConverters.*
+      states.keys().asScala.toList.foreach(releaseLeasesFor)
+    }
   }
 
   /**

core/exec/src/mill/exec/ExecutionContexts.scala

@@ -11,64 +11,6 @@ import mill.api.Logger
 import mill.api.daemon.internal.NonFatal
 
 object ExecutionContexts {
-  private object RunnablePriority {
-    private val priorityMethod = new ClassValue[Option[java.lang.reflect.Method]] {
-      override def computeValue(clazz: Class[?]): Option[java.lang.reflect.Method] =
-        try {
-          val method = clazz.getMethod("priority")
-          method.trySetAccessible()
-          Some(method)
-        } catch {
-          case _: ReflectiveOperationException => None
-        }
-    }
-
-    def apply(runnable: Runnable): Int =
-      priorityMethod.get(runnable.getClass) match {
-        case Some(method) => method.invoke(runnable).asInstanceOf[Int]
-        case None => 0
-      }
-  }
-
-  private final class QueuedRunnable(
-      runnable: Runnable,
-      val priority: Int,
-      val submissionIndex: Long
-  ) extends Runnable
-      with Comparable[QueuedRunnable] {
-    def run(): Unit = runnable.run()
-
-    override def compareTo(other: QueuedRunnable): Int =
-      priority.compareTo(other.priority) match {
-        case 0 =>
-          // Comparable needs a total ordering, so break ties using submission order.
-          submissionIndex.compareTo(other.submissionIndex)
-        case n => n
-      }
-  }
-
-  private final class PriorityThreadPoolExecutor(
-      threadCount: Int,
-      threadFactory: ThreadFactory
-  ) extends ThreadPoolExecutor(
-        threadCount,
-        threadCount,
-        60 * 1000,
-        TimeUnit.SECONDS,
-        new PriorityBlockingQueue[Runnable](),
-        threadFactory
-      ) {
-    private val submissionCount = new java.util.concurrent.atomic.AtomicLong()
-
-    override def execute(command: Runnable): Unit =
-      super.execute(
-        new QueuedRunnable(
-          runnable = command,
-          priority = RunnablePriority(command),
-          submissionIndex = submissionCount.getAndIncrement()
-        )
-      )
-  }
 
   /**
    * Execution context that runs code immediately when scheduled, without
@@ -134,15 +76,29 @@ object ExecutionContexts {
     def reportFailure(t: Throwable): Unit = {}
     def close(): Unit = executor.shutdown()
 
+    val priorityRunnableCount = java.util.concurrent.atomic.AtomicLong()
+
     /**
      * Subclass of [[java.lang.Runnable]] that assigns a priority to execute it
      *
      * Priority 0 is the default priority of all Mill task, priorities <0 can be used to
      * prioritize this runnable over most other tasks, while priorities >0 can be used to
      * de-prioritize it.
      */
-    class PriorityRunnable(val priority: Int, run0: () => Unit) extends Runnable {
-      def run(): Unit = run0()
+    class PriorityRunnable(val priority: Int, run0: () => Unit) extends Runnable
+        with Comparable[PriorityRunnable] {
+      def run() = run0()
+      val priorityRunnableIndex: Long = priorityRunnableCount.getAndIncrement()
+      override def compareTo(o: PriorityRunnable): Int = priority.compareTo(o.priority) match {
+        case 0 =>
+          // `Comparable` wants a *total* ordering, so we need to use `priorityRunnableIndex`
+          // to break ties between instances with the same priority. This index is assigned
+          // when a task is submitted, so it should more or less follow insertion order,
+          // and is a `Long` which should be big enough never to overflow
+          assert(this == o || this.priorityRunnableIndex != o.priorityRunnableIndex)
+          this.priorityRunnableIndex.compareTo(o.priorityRunnableIndex)
+        case n => n
+      }
     }
 
     /**
@@ -195,15 +151,25 @@ object ExecutionContexts {
   def createExecutor(threadCount: Int): ThreadPoolExecutor = {
     val executorIndex = executorCounter.incrementAndGet()
     val threadCounter = new AtomicInteger
-    val threadFactory: ThreadFactory = runnable => {
-      val threadIndex = threadCounter.incrementAndGet()
-      val t = new Thread(
-        runnable,
-        s"execution-contexts-threadpool-$executorIndex-thread-$threadIndex"
-      )
-      t.setDaemon(true)
-      t
-    }
-    new PriorityThreadPoolExecutor(threadCount, threadFactory)
+    new ThreadPoolExecutor(
+      threadCount,
+      threadCount,
+      60 * 1000,
+      TimeUnit.SECONDS,
+      // Use a `Deque` rather than a normal `Queue`, with the various `poll`/`take`
+      // operations reversed, providing elements in a LIFO order. This ensures that
+      // child `fork.async` tasks always take priority over parent tasks, avoiding
+      // large numbers of blocked parent tasks from piling up
+      new PriorityBlockingQueue[Runnable](),
+      runnable => {
+        val threadIndex = threadCounter.incrementAndGet()
+        val t = new Thread(
+          runnable,
+          s"execution-contexts-threadpool-$executorIndex-thread-$threadIndex"
+        )
+        t.setDaemon(true)
+        t
+      }
+    )
   }
 }