Execute test classes concurrently

.bazelrc_shared

@@ -6,7 +6,7 @@ build --tool_java_language_version="21"
 build --tool_java_runtime_version="remotejdk_21"
 
 # Other options
-build --experimental_use_hermetic_linux_sandbox
+# build --experimental_use_hermetic_linux_sandbox
 build --experimental_worker_cancellation
 build --experimental_worker_multiplex_sandboxing
 build --experimental_worker_sandbox_hardening

rules/private/phases/phase_test_launcher.bzl

@@ -61,6 +61,8 @@ def phase_test_launcher(ctx, g):
         files.append(subprocess_executable)
         args.add("--isolation", "process")
         args.add("--subprocess_exec", subprocess_executable.short_path)
+    if ctx.attr.sequential:
+        args.add("--sequential")
     args.add_all("--", test_jars, map_each = _short_path)
     args.set_param_file_format("multiline")
     args_file = ctx.actions.declare_file("{}/test.params".format(ctx.label.name))

rules/scala.bzl

@@ -382,6 +382,10 @@ def make_scala_test(*extras):
                     cfg = _scala_outgoing_transition,
                     default = "@rules_scala_annex//src/main/scala/higherkindness/rules_scala/workers/zinc/test",
                 ),
+                "sequential": attr.bool(
+                    default = False,
+                    doc = "Whether to run test classes sequentially. If false, they'll be run concurrently.",
+                ),
                 "subprocess_runner": attr.label(
                     cfg = _scala_outgoing_transition,
                     default = "@rules_scala_annex//src/main/scala/higherkindness/rules_scala/common/sbt-testing:subprocess",

src/main/scala/higherkindness/rules_scala/common/sbt-testing/BufferedLogger.scala

@@ -0,0 +1,37 @@
+package higherkindness.rules_scala.common.sbt_testing
+
+import sbt.testing.Logger
+
+import scala.collection.mutable
+
+private sealed trait SbtLogEntry
+private object SbtLogEntry {
+  case class Error(message: String) extends SbtLogEntry
+  case class Warn(message: String) extends SbtLogEntry
+  case class Info(message: String) extends SbtLogEntry
+  case class Debug(message: String) extends SbtLogEntry
+  case class Trace(throwable: Throwable) extends SbtLogEntry
+}
+
+class BufferedLogger(underlying: Logger) extends Logger {
+  private val buffer = mutable.ArrayBuffer.empty[SbtLogEntry]
+
+  override def ansiCodesSupported(): Boolean = underlying.ansiCodesSupported()
+  override def error(message: String): Unit = buffer.addOne(SbtLogEntry.Error(message))
+  override def warn(message: String): Unit = buffer.addOne(SbtLogEntry.Warn(message))
+  override def info(message: String): Unit = buffer.addOne(SbtLogEntry.Info(message))
+  override def debug(message: String): Unit = buffer.addOne(SbtLogEntry.Debug(message))
+  override def trace(throwable: Throwable): Unit = buffer.addOne(SbtLogEntry.Trace(throwable))
+
+  def flush(): Unit = {
+    buffer.foreach {
+      case SbtLogEntry.Error(message)   => underlying.error(message)
+      case SbtLogEntry.Warn(message)    => underlying.warn(message)
+      case SbtLogEntry.Info(message)    => underlying.info(message)
+      case SbtLogEntry.Debug(message)   => underlying.debug(message)
+      case SbtLogEntry.Trace(throwable) => underlying.trace(throwable)
+    }
+
+    buffer.clear()
+  }
+}

src/main/scala/higherkindness/rules_scala/common/sbt-testing/JUnitXmlReporter.scala

@@ -2,70 +2,71 @@ package higherkindness.rules_scala.common.sbt_testing
 
 import java.io.{PrintWriter, StringWriter}
 import sbt.testing.Status.{Canceled, Error, Failure, Ignored, Pending, Skipped}
-import sbt.testing.{Event, Status, TestSelector}
-import scala.collection.mutable.ListBuffer
+import sbt.testing.{Event, Status, Task, TestSelector}
 import scala.xml.{Elem, Utility, XML}
 
-class JUnitXmlReporter(tasksAndEvents: ListBuffer[(String, ListBuffer[Event])]) {
+class JUnitXmlReporter(taskEvents: Map[Task, Array[Event]]) {
   private def escape(info: String): String = info match {
     case str: String => Utility.escape(str)
     case null        => ""
   }
 
   def result: Elem =
     XML.loadString(s"""<testsuites>
-      ${(for ((name, events) <- tasksAndEvents)
-        yield s"""<testsuite
+      ${(for {
+        (task, events) <- taskEvents
+        name = task.taskDef.fullyQualifiedName
+      } yield s"""<testsuite
         hostname=""
         name="${escape(name)}"
         tests="${events.size.toString}"
         errors="${events.count(_.status == Error).toString}"
         failures="${events.count(_.status == Failure).toString}"
         skipped="${events
-            .count(e => e.status == Ignored || e.status == Skipped || e.status == Pending || e.status == Canceled)
-            .toString}"
+          .count(e => e.status == Ignored || e.status == Skipped || e.status == Pending || e.status == Canceled)
+          .toString}"
         time="${(events.map(_.duration).sum / 1000d).toString}">
           ${(for (e <- events)
-            yield s"""<testcase
+          yield s"""<testcase
             classname="${escape(name)}"
             name="${e.selector match {
-                case selector: TestSelector => escape(selector.testName)
-                case _                      => "Error occurred outside of a test case."
-              }}"
+              case selector: TestSelector => escape(selector.testName)
+              case _                      => "Error occurred outside of a test case."
+            }}"
             time="${(e.duration / 1000d).toString}">
             ${val stringWriter = new StringWriter()
-              if (e.throwable.isDefined) {
-                val writer = new PrintWriter(stringWriter)
-                e.throwable.get.printStackTrace(writer)
-                writer.flush()
-              }
-              val trace: String = stringWriter.toString
-              e.status match {
-                case Status.Error if e.throwable.isDefined =>
-                  val t = e.throwable.get
-                  s"""<error message="${escape(t.getMessage)}" type="${escape(t.getClass.getName)}">${escape(
-                      trace,
-                    )}</error>"""
-                case Status.Error =>
-                  s"""<error message="No Exception or message provided"/>"""
-                case Status.Failure if e.throwable.isDefined =>
-                  val t = e.throwable.get
-                  s"""<failure message="${escape(t.getMessage)}" type="${escape(t.getClass.getName)}">${escape(
-                      trace,
-                    )}</failure>"""
-                case Status.Failure =>
-                  s"""<failure message="No Exception or message provided"/>"""
-                case Status.Canceled if e.throwable.isDefined =>
-                  val t = e.throwable.get
-                  s"""<skipped message="${escape(t.getMessage)}" type="${escape(t.getClass.getName)}">${escape(
-                      trace,
-                    )}</skipped>"""
-                case Status.Canceled =>
-                  s"""<skipped message="No Exception or message provided"/>"""
-                case Status.Ignored | Status.Skipped | Status.Pending =>
-                  "<skipped/>"
-                case _ =>
-              }}
+            if (e.throwable.isDefined) {
+              val writer = new PrintWriter(stringWriter)
+              e.throwable.get.printStackTrace(writer)
+              writer.flush()
+            }
+            val trace: String = stringWriter.toString
+            e.status match {
+              case Status.Error if e.throwable.isDefined =>
+                val t = e.throwable.get
+                s"""<error message="${escape(t.getMessage)}" type="${escape(t.getClass.getName)}">${escape(
+                    trace,
+                  )}</error>"""
+              case Status.Error =>
+                s"""<error message="No Exception or message provided"/>"""
+              case Status.Failure if e.throwable.isDefined =>
+                val t = e.throwable.get
+                s"""<failure message="${escape(t.getMessage)}" type="${escape(t.getClass.getName)}">${escape(
+                    trace,
+                  )}</failure>"""
+              case Status.Failure =>
+                s"""<failure message="No Exception or message provided"/>"""
+              case Status.Canceled if e.throwable.isDefined =>
+                val t = e.throwable.get
+                s"""<skipped message="${escape(t.getMessage)}" type="${escape(t.getClass.getName)}">${escape(
+                    trace,
+                  )}</skipped>"""
+              case Status.Canceled =>
+                s"""<skipped message="No Exception or message provided"/>"""
+              case Status.Ignored | Status.Skipped | Status.Pending =>
+                "<skipped/>"
+              case _ =>
+            }}
           </testcase>""").mkString("")}
         <system-out><![CDATA[]]></system-out>
         <system-err><![CDATA[]]></system-err>

src/main/scala/higherkindness/rules_scala/common/sbt-testing/SubprocessRunner.scala

@@ -3,7 +3,8 @@ package higherkindness.rules_scala.common.sbt_testing
 import higherkindness.rules_scala.common.classloaders.ClassLoaders
 import java.io.ObjectInputStream
 import java.nio.file.Paths
-import scala.collection.mutable
+import scala.concurrent.Await
+import scala.concurrent.duration.Duration
 
 object SubprocessTestRunner {
 
@@ -20,14 +21,16 @@ object SubprocessTestRunner {
         val tasks = runner.tasks(Array(TestHelper.taskDef(request.test, request.scopeAndTestName)))
         tasks.length == 0 || {
           val reporter = new TestReporter(request.logger)
-          val taskExecutor = new TestTaskExecutor(request.logger)
-          val failures = mutable.Set[String]()
-          tasks.foreach { task =>
-            reporter.preTask(task)
-            taskExecutor.execute(task, failures)
-            reporter.postTask()
-          }
-          !failures.nonEmpty
+
+          // We're only running a single test class, so there's not much of a point in using the
+          // `ConcurrentTestTaskExecutor`
+          val taskExecutor = new SequentialTestTaskExecutor(request.logger)
+
+          tasks.foreach(taskExecutor.submitTask)
+
+          val result = Await.result(taskExecutor.waitForTasks(), Duration.Inf)
+
+          !result.failures.nonEmpty
         }
       }
     }

src/main/scala/higherkindness/rules_scala/common/sbt-testing/Test.scala

@@ -2,7 +2,7 @@ package higherkindness.rules_scala.common.sbt_testing
 
 import java.nio.file.{Path, Paths}
 import play.api.libs.json.{Format, Json}
-import sbt.testing.{Event, Framework, Logger, Runner, Status, Task, TaskDef, TestWildcardSelector}
+import sbt.testing.{Framework, Logger, Runner, Task, TaskDef, TestWildcardSelector}
 import scala.collection.mutable
 import scala.util.control.NonFatal
 
@@ -101,25 +101,3 @@ class TestReporter(logger: Logger) {
 
   def preTask(task: Task) = logger.info(task.taskDef.fullyQualifiedName)
 }
-
-class TestTaskExecutor(logger: Logger) {
-  def execute(task: Task, failures: mutable.Set[String]): mutable.ListBuffer[Event] = {
-    var events = new mutable.ListBuffer[Event]()
-    def execute(task: Task): Unit = {
-      val tasks = task.execute(
-        event => {
-          events += event
-          event.status match {
-            case Status.Failure | Status.Error =>
-              failures += task.taskDef.fullyQualifiedName
-            case _ =>
-          }
-        },
-        Array(new PrefixedTestingLogger(logger, "    ")),
-      )
-      tasks.foreach(execute)
-    }
-    execute(task)
-    events
-  }
-}

src/main/scala/higherkindness/rules_scala/common/sbt-testing/TestTaskExecutor.scala

@@ -0,0 +1,119 @@
+package higherkindness.rules_scala.common.sbt_testing
+
+import java.util.concurrent.ConcurrentLinkedQueue
+import sbt.testing.{Event, Logger, Status, Task}
+import scala.collection.{concurrent, mutable}
+import scala.concurrent.{blocking, ExecutionContext, Future}
+import scala.jdk.CollectionConverters.*
+
+case class TaskExecutorResult(taskEvents: Map[Task, Array[Event]], failures: Array[String])
+
+private object TaskExecutorResult {
+  private[sbt_testing] case class Mutable(
+    taskEvents: concurrent.TrieMap[Task, ConcurrentLinkedQueue[Event]],
+    failures: ConcurrentLinkedQueue[String],
+  ) {
+    def clear(): Unit = {
+      taskEvents.clear()
+      failures.clear()
+    }
+
+    def toTaskExecutorResult: TaskExecutorResult = TaskExecutorResult(
+      taskEvents.view.map { case task -> events => task -> events.asScala.toArray }.toMap,
+      failures.asScala.toArray,
+    )
+  }
+
+  private[sbt_testing] object Mutable {
+    def empty: Mutable = apply(concurrent.TrieMap.empty, new ConcurrentLinkedQueue())
+  }
+}
+
+trait TestTaskExecutor {
+  def submitTask(task: Task): Unit
+  def waitForTasks(): Future[TaskExecutorResult]
+}
+
+class ConcurrentTestTaskExecutor(logger: Logger) extends TestTaskExecutor {
+  private val activeTasks = new ConcurrentLinkedQueue[Future[Unit]]()
+  private val currentResult = TaskExecutorResult.Mutable.empty
+
+  override def submitTask(task: Task): Unit = activeTasks.add(
+    Future {
+      blocking {
+        val bufferedLogger = new BufferedLogger(logger)
+        val reporter = new TestReporter(bufferedLogger)
+
+        reporter.preTask(task)
+
+        val additionalTasks = task.execute(
+          event => {
+            currentResult.synchronized {
+              currentResult.taskEvents.getOrElseUpdate(task, new ConcurrentLinkedQueue()).add(event)
+
+              event.status match {
+                case Status.Failure | Status.Error => currentResult.failures.add(task.taskDef.fullyQualifiedName)
+                case _                             =>
+              }
+            }
+          },
+          Array(new PrefixedTestingLogger(bufferedLogger, "    ")),
+        )
+
+        additionalTasks.foreach(submitTask)
+
+        reporter.postTask()
+
+        // Only one task should write to stderr/stdout at a time. Of course, the task implementation could write to
+        // stdout/stderr directly, but that's out of our control.
+        synchronized {
+          bufferedLogger.flush()
+        }
+      }
+    }(ExecutionContext.global),
+  )
+
+  override def waitForTasks(): Future[TaskExecutorResult] = {
+    given ExecutionContext = ExecutionContext.global
+
+    Future
+      .sequence(activeTasks.asScala)
+      .map { _ =>
+        activeTasks.clear()
+        currentResult.toTaskExecutorResult
+      }(ExecutionContext.global)
+  }
+}
+
+class SequentialTestTaskExecutor(logger: Logger) extends TestTaskExecutor {
+  private val currentResult = TaskExecutorResult.Mutable.empty
+
+  override def submitTask(task: Task): Unit = {
+    val reporter = new TestReporter(logger)
+
+    reporter.preTask(task)
+
+    val additionalTasks = task.execute(
+      event => {
+        currentResult.taskEvents.getOrElseUpdate(task, new ConcurrentLinkedQueue()).add(event)
+
+        event.status match {
+          case Status.Failure | Status.Error => currentResult.failures.add(task.taskDef.fullyQualifiedName)
+          case _                             =>
+        }
+      },
+      Array(new PrefixedTestingLogger(logger, "    ")),
+    )
+
+    additionalTasks.foreach(submitTask)
+    reporter.postTask()
+  }
+
+  override def waitForTasks(): Future[TaskExecutorResult] = {
+    val result = currentResult.toTaskExecutorResult
+
+    currentResult.clear()
+
+    Future.successful(result)
+  }
+}