Optimize BSP initialization stuff (forwardport #4698) (#4724)

Report of #4698 in `main`

---

This optimizes the `workspaceBuildTargets` BSP request. On my machine,
this makes it go from around 15s to about 3s on the Mill repo. Which is
time saved when opening or reloading Mill in one's IDE. Note that gains
are more modest on smaller code bases, but still there.

In more detail, two things are optimized:
- the `MillBuildServer.groupList` method, whose first implementation
(that I wrote) was too naive (don't know how the code it replaced at the
time fared)
- it reduces the number of calls to `Evaluator.evaluate` when the BSP
server needs to evaluate tasks. It does so by really grouping them by
evaluator, rather than by (evaluator, BuildTargetId). Seems the former
grouping was intended, but the grouping was actually more loose.

Lastly, not an optimization per se, but this makes the BSP server less
verbose. It was printing some very large objects in particular, that
tend to flood the output / BSP server log. These messages are now only
printed if the env var `MILL_BSP_DEBUG` is `true`.

Author: alexarchambault

bsp/worker/src/mill/bsp/worker/BspWorkerImpl.scala

@@ -29,7 +29,8 @@ private class BspWorkerImpl() extends BspWorker {
         serverVersion = BuildInfo.millVersion,
         serverName = Constants.serverName,
         logStream = logStream,
-        canReload = canReload
+        canReload = canReload,
+        debugMessages = Option(System.getenv("MILL_BSP_DEBUG")).contains("true")
       ) with MillJvmBuildServer with MillJavaBuildServer with MillScalaBuildServer
 
     val executor = Executors.newCachedThreadPool()

bsp/worker/src/mill/bsp/worker/MillBuildServer.scala

@@ -17,6 +17,7 @@ import mill.given
 
 import java.io.PrintStream
 import java.util.concurrent.CompletableFuture
+import scala.collection.mutable
 import scala.concurrent.Promise
 import scala.jdk.CollectionConverters.*
 import scala.reflect.ClassTag
@@ -29,7 +30,8 @@ private class MillBuildServer(
     serverVersion: String,
     serverName: String,
     logStream: PrintStream,
-    canReload: Boolean
+    canReload: Boolean,
+    debugMessages: Boolean
 ) extends ExternalModule
     with BuildServer {
 
@@ -55,12 +57,16 @@ private class MillBuildServer(
     if (statePromise.isCompleted) statePromise = Promise[State]() // replace the promise
     evaluatorsOpt.foreach { evaluators =>
       statePromise.success(
-        new State(topLevelProjectRoot, evaluators, debug)
+        new State(topLevelProjectRoot, evaluators, s => debug(s()))
       )
     }
   }
 
-  def debug(msg: String): Unit = logStream.println(msg)
+  def print(msg: String): Unit =
+    logStream.println(msg)
+  def debug(msg: => String): Unit =
+    if (debugMessages)
+      logStream.println("[debug] " + msg)
 
   def onConnectWithClient(buildClient: BuildClient): Unit = client = buildClient
 
@@ -109,7 +115,7 @@ private class MillBuildServer(
         case d: JsonObject =>
           debug(s"extra data: ${d} of type ${d.getClass}")
           readVersion(d, "semanticdbVersion").foreach { version =>
-            debug(
+            print(
               s"Got client semanticdbVersion: ${version}. Enabling SemanticDB support."
             )
             clientWantsSemanticDb = true
@@ -126,27 +132,27 @@ private class MillBuildServer(
     }
 
   override def onBuildInitialized(): Unit = {
-    debug("Build initialized")
+    print("Build initialized")
   }
 
   override def buildShutdown(): CompletableFuture[Object] = {
-    debug(s"Entered buildShutdown")
+    print("Entered buildShutdown")
     shutdownRequested = true
     onSessionEnd match {
       case None =>
       case Some(onEnd) =>
-        debug("Shutdown build...")
+        print("Shutdown build...")
         onEnd(BspServerResult.Shutdown)
     }
     SemanticDbJavaModule.resetContext()
     CompletableFuture.completedFuture(null.asInstanceOf[Object])
   }
   override def onBuildExit(): Unit = {
-    debug("Entered onBuildExit")
+    print("Entered onBuildExit")
     onSessionEnd match {
       case None =>
       case Some(onEnd) =>
-        debug("Exiting build...")
+        print("Exiting build...")
         onEnd(BspServerResult.Shutdown)
     }
     SemanticDbJavaModule.resetContext()
@@ -204,7 +210,7 @@ private class MillBuildServer(
       onSessionEnd match {
         case None => "unsupportedWorkspaceReload".asInstanceOf[Object]
         case Some(onEnd) =>
-          debug("Reloading workspace...")
+          print("Reloading workspace...")
           onEnd(BspServerResult.ReloadWorkspace).asInstanceOf[Object]
       }
     }
@@ -405,7 +411,7 @@ private class MillBuildServer(
           (Task.Anon { () }, ev)
           val task = Task.Anon {
             Task.log.debug(
-              "Ignoring invalid compile request for test module ${m.bspBuildTarget.displayName}"
+              s"Ignoring invalid compile request for test module ${m.bspBuildTarget.displayName}"
             )
           }
           (task, ev)
@@ -669,38 +675,50 @@ private class MillBuildServer(
       }
 
       // group by evaluator (different root module)
-      val evaluated = groupList(tasksSeq.toSeq)(_._2)(_._1)
-        .map { case ((ev, id), ts) =>
-          val results = evaluate(ev, ts)
-          val failures = results.transitiveResults.collect {
-            case (_, res: ExecResult.Failing[_]) => res
+      val groups0 = groupList(tasksSeq.toSeq)(_._2._1) {
+        case (tasks, (_, id)) => (id, tasks)
+      }
+
+      val evaluated = groups0.flatMap {
+        case (ev, targetIdTasks) =>
+          val results = evaluate(ev, targetIdTasks.map(_._2))
+          val idByTasks = targetIdTasks.map { case (id, task) => (task: Task[_], id) }.toMap
+          val failures = results.transitiveResults.toSeq.collect {
+            case (task, res: ExecResult.Failing[_]) if idByTasks.contains(task) =>
+              (idByTasks(task), res)
           }
 
-          def logError(errorMsg: String): Unit = {
+          def logError(id: BuildTargetIdentifier, errorMsg: String): Unit = {
             val msg = s"Request '$prefix' failed for ${id.getUri}: ${errorMsg}"
             debug(msg)
             client.onBuildLogMessage(new LogMessageParams(MessageType.ERROR, msg))
           }
 
-          if (failures.nonEmpty) {
-            logError(failures.mkString(", "))
+          if (failures.nonEmpty)
+            for ((id, failure) <- failures)
+              logError(id, failure.toString)
+
+          val resultsById = targetIdTasks.flatMap {
+            case (id, task) =>
+              results.transitiveResults(task)
+                .asSuccess
+                .map(_.value.value.asInstanceOf[W])
+                .map((id, _))
           }
 
-          // only the successful results
-          val successes = results.values.map(_.value).asInstanceOf[Seq[W]]
-          successes
-            .flatMap { v =>
+          resultsById.flatMap {
+            case (id, values) =>
               try {
-                Seq(f(ev, state, id, state.bspModulesById(id)._1, v))
+                Seq(f(ev, state, id, state.bspModulesById(id)._1, values))
               } catch {
                 case NonFatal(e) =>
-                  logError(e.toString)
+                  logError(id, e.toString)
                   Seq()
               }
-            }
-        }
+          }
+      }
 
-      agg(evaluated.flatten.asJava, state)
+      agg(evaluated.asJava, state)
     }
   }
 
@@ -717,12 +735,12 @@ private class MillBuildServer(
       hint: String,
       checkInitialized: Boolean = true
   )(f: State => V): CompletableFuture[V] = {
-    debug(s"Entered ${hint}")
+    print(s"Entered ${hint}")
 
     val start = System.currentTimeMillis()
     val prefix = hint.split(" ").head
     def took =
-      debug(s"${prefix} took ${System.currentTimeMillis() - start} msec")
+      print(s"${prefix} took ${System.currentTimeMillis() - start} msec")
 
     val future = new CompletableFuture[V]()
     if (checkInitialized && !initialized) {
@@ -762,11 +780,11 @@ private class MillBuildServer(
       hint: String,
       checkInitialized: Boolean = true
   )(f: => V): CompletableFuture[V] = {
-    debug(s"Entered ${hint}")
+    print(s"Entered ${hint}")
     val start = System.currentTimeMillis()
     val prefix = hint.split(" ").head
     def took =
-      debug(s"${prefix} took ${System.currentTimeMillis() - start} msec")
+      print(s"${prefix} took ${System.currentTimeMillis() - start} msec")
 
     val future = new CompletableFuture[V]()
 
@@ -834,10 +852,27 @@ private object MillBuildServer {
    * the order of appearance of the keys from the input sequence
    */
   private def groupList[A, K, B](seq: Seq[A])(key: A => K)(f: A => B): Seq[(K, Seq[B])] = {
-    val keyIndices = seq.map(key).distinct.zipWithIndex.toMap
-    seq.groupMap(key)(f)
-      .toSeq
-      .sortBy { case (k, _) => keyIndices(k) }
+    val map = new mutable.HashMap[K, mutable.ListBuffer[B]]
+    val list = new mutable.ListBuffer[(K, mutable.ListBuffer[B])]
+    for (a <- seq) {
+      val k = key(a)
+      val b = f(a)
+      val l = map.getOrElseUpdate(
+        k, {
+          val buf = mutable.ListBuffer[B]()
+          list.append((k, buf))
+          buf
+        }
+      )
+      l.append(b)
+    }
+    list
+      .iterator
+      .map {
+        case (k, l) =>
+          (k, l.result())
+      }
+      .toList
   }
 
   def jvmBuildTarget(d: JvmBuildTarget): bsp4j.JvmBuildTarget =

bsp/worker/src/mill/bsp/worker/State.scala

@@ -5,7 +5,11 @@ import mill.scalalib.bsp.BspModule
 import mill.scalalib.internal.JavaModuleUtils
 import mill.define.{Evaluator, Module}
 
-private class State(workspaceDir: os.Path, evaluators: Seq[Evaluator], debug: String => Unit) {
+private class State(
+    workspaceDir: os.Path,
+    evaluators: Seq[Evaluator],
+    debug: (() => String) => Unit
+) {
   lazy val bspModulesIdList: Seq[(BuildTargetIdentifier, (BspModule, Evaluator))] = {
     val modules: Seq[(Module, Seq[Module], Evaluator)] = evaluators
       .map(ev => (ev.rootModule, JavaModuleUtils.transitiveModules(ev.rootModule), ev))
@@ -24,7 +28,7 @@ private class State(workspaceDir: os.Path, evaluators: Seq[Evaluator], debug: St
   }
   lazy val bspModulesById: Map[BuildTargetIdentifier, (BspModule, Evaluator)] = {
     val map = bspModulesIdList.toMap
-    debug(s"BspModules: ${map.view.mapValues(_._1.bspDisplayName).toMap}")
+    debug(() => s"BspModules: ${map.view.mapValues(_._1.bspDisplayName).toMap}")
     map
   }