.

Author: lihaoyi

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

@@ -11,6 +11,62 @@ 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 => 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
@@ -76,29 +132,15 @@ 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
-        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
-      }
+    class PriorityRunnable(val priority: Int, run0: () => Unit) extends Runnable {
+      def run(): Unit = run0()
     }
 
     /**
@@ -151,25 +193,15 @@ object ExecutionContexts {
   def createExecutor(threadCount: Int): ThreadPoolExecutor = {
     val executorIndex = executorCounter.incrementAndGet()
     val threadCounter = new AtomicInteger
-    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
-      }
-    )
+    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)
   }
 }

core/exec/test/src/mill/exec/ExecutionTests.scala

@@ -6,6 +6,9 @@ import mill.testkit.{TestRootModule, UnitTester}
 import mill.{PathRef, exec}
 import utest.*
 
+import java.net.URLClassLoader
+import java.util.concurrent.{CountDownLatch, TimeUnit}
+
 object ExecutionTests extends TestSuite {
   object traverseBuild extends TestRootModule {
     trait TaskModule extends mill.Module {
@@ -88,6 +91,26 @@ object ExecutionTests extends TestSuite {
   class Checker[T <: mill.testkit.TestRootModule](module: T)
       extends exec.Checker(module)
 
+  private final class ChildFirstExecutionContextsLoader(urls: Array[java.net.URL], parent: ClassLoader)
+      extends URLClassLoader(urls, parent) {
+    override def loadClass(name: String, resolve: Boolean): Class[?] = synchronized {
+      val alreadyLoaded = findLoadedClass(name)
+      val loaded =
+        if (alreadyLoaded != null) alreadyLoaded
+        else if (name == "mill.exec.ExecutionContexts$" || name.startsWith(
+            "mill.exec.ExecutionContexts$"
+          )) {
+          try findClass(name)
+          catch {
+            case _: ClassNotFoundException => super.loadClass(name, false)
+          }
+        } else super.loadClass(name, false)
+
+      if (resolve) resolveClass(loaded)
+      loaded
+    }
+  }
+
   val tests = Tests {
     import TestGraphs.*
     import utest.*
@@ -568,6 +591,52 @@ object ExecutionTests extends TestSuite {
         assert(res.executionResults.transitiveFailing.keySet == Set(anonTaskFailure.task))
       }
     }
+
+    test("threadPoolSupportsMixedClassloaderPriorityRunnables") {
+      val executor = ExecutionContexts.createExecutor(1)
+      val localPool = new ExecutionContexts.ThreadPool(executor)
+      val codeSourceUrl = ExecutionContexts.getClass.getProtectionDomain.getCodeSource.getLocation
+      val childLoader =
+        new ChildFirstExecutionContextsLoader(Array(codeSourceUrl), ExecutionTests.getClass.getClassLoader)
+
+      val blockerStarted = new CountDownLatch(1)
+      val unblock = new CountDownLatch(1)
+      val completed = new CountDownLatch(2)
+
+      try {
+        val threadPoolClass = childLoader.loadClass("mill.exec.ExecutionContexts$ThreadPool")
+        val ctor = threadPoolClass.getDeclaredConstructors.head
+        ctor.trySetAccessible()
+        val childPool = ctor
+          .newInstance(executor)
+          .asInstanceOf[mill.api.TaskCtx.Fork.Impl]
+
+        localPool.execute(new Runnable {
+          override def run(): Unit = {
+            blockerStarted.countDown()
+            unblock.await(30, TimeUnit.SECONDS)
+          }
+        })
+
+        assert(blockerStarted.await(30, TimeUnit.SECONDS))
+
+        childPool.execute(new Runnable {
+          override def run(): Unit = completed.countDown()
+        })
+        localPool.execute(new Runnable {
+          override def run(): Unit = completed.countDown()
+        })
+
+        unblock.countDown()
+        assert(completed.await(30, TimeUnit.SECONDS))
+      } finally {
+        unblock.countDown()
+        localPool.close()
+        assert(executor.awaitTermination(30, TimeUnit.SECONDS))
+        childLoader.close()
+      }
+    }
+
     test("overloaded") {
       UnitTester(overloads, null).scoped { tester =>
         val res = tester.apply(Seq(overloads.overloaded(1)))

integration/bsp-util/src/BspServerTestUtil.scala

@@ -276,7 +276,7 @@ object BspServerTestUtil {
         try proc.stdout.close()
         catch { case _: java.io.IOException => () }
 
-        proc.join(30000L)
+        waitForProcessExit(proc, 30000L)
       } finally {
         if (!success && isCI) {
           System.err.println(" == BSP server output ==")
@@ -287,6 +287,32 @@ object BspServerTestUtil {
     }
   }
 
+  private def waitForProcessExit(proc: os.SubProcess, timeoutMillis: Long): Unit = {
+    val waitIntervalMillis = 50L
+    val deadlineNanos = System.nanoTime() + timeoutMillis * 1000000L
+
+    while (proc.isAlive() && System.nanoTime() < deadlineNanos) {
+      Thread.sleep(waitIntervalMillis)
+    }
+
+    if (proc.isAlive()) {
+      // `os.SubProcess.join(timeout)` falls back to recursive destruction on
+      // timeout, which uses `ProcessHandle.children()` and is not permitted in
+      // this macOS sandbox. BSP launchers are top-level processes for these
+      // tests, so a non-recursive destroy is sufficient here.
+      proc.destroy(recursive = false)
+
+      val shutdownDeadlineNanos = System.nanoTime() + 5000L * 1000000L
+      while (proc.isAlive() && System.nanoTime() < shutdownDeadlineNanos) {
+        Thread.sleep(waitIntervalMillis)
+      }
+
+      if (proc.isAlive()) {
+        throw new RuntimeException("BSP server did not exit within the expected timeout")
+      }
+    }
+  }
+
   lazy val millWorkspace: os.Path = {
     val value = Option(System.getenv("MILL_PROJECT_ROOT")).getOrElse(???)
     os.Path(value)

integration/dedicated/bsp-server-error/src/BspServerErrorTests.scala

@@ -45,13 +45,17 @@ object BspServerErrorTests extends UtestIntegrationTestSuite {
         bspLog = Some((bytes, len) => stderr.write(bytes, 0, len))
       ) { (buildServer, initRes) =>
 
+        val firstServerDeadline = System.nanoTime() + 5000L * 1000000L
+        while (firstServerProc.isAlive() && System.nanoTime() < firstServerDeadline) {
+          Thread.sleep(50L)
+        }
         assert(!firstServerProc.isAlive())
 
         val firstServerStderrStr = new String(firstServerStderr.toByteArray)
-        assert(firstServerStderrStr.contains("Received SIGTERM, exiting"))
+        assert(firstServerStderrStr.contains("BSP shutdown asked by client, exiting"))
 
         val currentStderrStr = new String(stderr.toByteArray)
-        assert(currentStderrStr.contains("Sent SIGTERM to process"))
+        assert(currentStderrStr.contains("Asked the active BSP session for 'Mill_Integration' to shut down"))
 
         assert(initRes.getCapabilities.getInverseSourcesProvider == true)