Cleanups for MillBuildServer (#2518)

1. Break `mill.bsp.worker.State` into a separate file

2. Fold `targetTasks` into `completableTasks`, and adjust all callsites
to use that

3. Make `completableTasks` take a `tasks: BspModule => Task[W]`
parameter, so we can perform all evaluation once up-front and then fall
back to normal code after (rather than having it all execute in the
context of the evaluator inside tasks, which obfuscates stack traces and
evaluation order)

4. Move `agg` param in `completableTasks` to a third parameter list, to
allow type inference to work better

5. Make everything in `mill.bsp.worker` private

6. Move all the remaining BSP stuff from `main/` and `runner/` to
`bsp/`, remove weird reflection hacks to just use direct instantiation
where possible since now they're in the same module

7. remove the `BspServerStarter`/`BspServerStarterImpl` layers of
indirection

8. Move `BspConfigJson.scala`, `bspConnectionJson` and
`createBspConnection` out of `bsp.worker` into `bsp` since it doesn't
depend on anything heavyweight that necessitates it being in the worker
module

9. Moved `def startBspServer` out of `mill.bsp.BSP` into
`mill.bsp.BspContext`, leaving `mill.bsp.BSP` to just serve one purpose
as `mill.define.Module` rather than additionally being a grab-bag of
helper logic

10. Simplify the concurrency story around instantiating the
`BspServerHandle`:
- Previously, the whole of `startBspServer` was put in a `Future`, with
the main result of the `Future` waiting on the server to terminate only
to do some logging and discard the result, and the creation of the
`BspServerHandle` was propagated via `Promise`s and that the main thread
would `Await` upon
- With this PR, the main thread runs the main logic of `startBspServer`,
directly returning the `BspServerHandle` when it completes (or
`Left[String]` on error). Server termination is waited for in a side
thread, that does nothing but call `listening.get()` and do logging when
it terminates. We still need a global variable to pass it to `def
startSession`, but this can just be a `@volatile var` rather than a
global `Promise`
- This maintains the existing semantics, but using 2 parallel threads
instead of 3, with the primary logic shifted to the main thread, and
without need for any `Await`s

Overall this PR tries to cut a lot of the unnecessary complexity and
messiness around the BSP logic. This results in a substantial reduction
in lines of code, while still keeping the core logic unchanged.
Hopefully this will make things easier to maintain and debug going
forward.

Tested manually using VSCode on `examples/misc/3-import-file-ivy`

Pull request: #2518

bsp/src/mill/bsp/BSP.scala

@@ -1,24 +1,16 @@
 package mill.bsp
 
-import java.io.{InputStream, PrintStream}
-import scala.concurrent.{Await, Promise}
-import scala.concurrent.duration.Duration
-import mill.api.{Ctx, DummyInputStream, Logger, PathRef, Result, SystemStreams}
+import mill.api.{Ctx, PathRef}
 import mill.{Agg, T, BuildInfo => MillBuildInfo}
-import mill.define.{Command, Discover, ExternalModule, Task}
+import mill.define.{Command, Discover, ExternalModule}
 import mill.eval.Evaluator
-import mill.main.{BspServerHandle, BspServerResult, BspServerStarter}
 import mill.util.Util.millProjectModule
-import mill.scalalib.{CoursierModule, Dep}
-import mill.util.PrintLogger
-import os.Path
+import mill.scalalib.CoursierModule
 
-object BSP extends ExternalModule with CoursierModule with BspServerStarter {
+object BSP extends ExternalModule with CoursierModule {
 
   lazy val millDiscover: Discover[this.type] = Discover[this.type]
 
-  private[this] val millServerHandle = Promise[BspServerHandle]()
-
   private def bspWorkerLibs: T[Agg[PathRef]] = T {
     millProjectModule("mill-bsp-worker", repositoriesTask())
   }
@@ -47,7 +39,7 @@ object BSP extends ExternalModule with CoursierModule with BspServerStarter {
       libUrls.mkString("\n"),
       createFolders = true
     )
-    BspWorker(T.ctx(), Some(libUrls)).map(_.createBspConnection(jobs, Constants.serverName))
+    createBspConnection(jobs, Constants.serverName)
   }
 
   /**
@@ -58,62 +50,53 @@ object BSP extends ExternalModule with CoursierModule with BspServerStarter {
    */
   def startSession(ev: Evaluator): Command[BspServerResult] = T.command {
     T.log.errorStream.println("BSP/startSession: Starting BSP session")
-    val serverHandle: BspServerHandle = Await.result(millServerHandle.future, Duration.Inf)
-    val res = serverHandle.runSession(ev)
+    val res = BspContext.bspServerHandle.runSession(ev)
     T.log.errorStream.println(s"BSP/startSession: Finished BSP session, result: ${res}")
     res
   }
 
-  override def startBspServer(
-      initialEvaluator: Option[Evaluator],
-      streams: SystemStreams,
-      logStream: Option[PrintStream],
-      workspaceDir: os.Path,
-      ammoniteHomeDir: os.Path,
-      canReload: Boolean,
-      serverHandle: Option[Promise[BspServerHandle]] = None
-  ): BspServerResult = {
-    val ctx = new Ctx.Workspace with Ctx.Home with Ctx.Log {
-      override def workspace: Path = workspaceDir
-      override def home: Path = ammoniteHomeDir
-      // This all goes to the BSP log file mill-bsp.stderr
-      override def log: Logger = new Logger {
-        override def colored: Boolean = false
-        override def systemStreams: SystemStreams = new SystemStreams(
-          out = streams.out,
-          err = streams.err,
-          in = DummyInputStream
-        )
-        override def info(s: String): Unit = streams.err.println(s)
-        override def error(s: String): Unit = streams.err.println(s)
-        override def ticker(s: String): Unit = streams.err.println(s)
-        override def debug(s: String): Unit = streams.err.println(s)
-        override def debugEnabled: Boolean = true
-      }
-    }
+  private def createBspConnection(
+      jobs: Int,
+      serverName: String
+  )(implicit ctx: Ctx): (PathRef, ujson.Value) = {
+    // we create a json connection file
+    val bspFile = ctx.workspace / Constants.bspDir / s"${serverName}.json"
+    if (os.exists(bspFile)) ctx.log.info(s"Overwriting BSP connection file: ${bspFile}")
+    else ctx.log.info(s"Creating BSP connection file: ${bspFile}")
+    val withDebug = ctx.log.debugEnabled
+    if (withDebug) ctx.log.debug(
+      "Enabled debug logging for the BSP server. If you want to disable it, you need to re-run this install command without the --debug option."
+    )
+    val connectionContent = bspConnectionJson(jobs, withDebug)
+    os.write.over(bspFile, connectionContent, createFolders = true)
+    (PathRef(bspFile), upickle.default.read[ujson.Value](connectionContent))
+  }
 
-    val worker = BspWorker(ctx)
+  private def bspConnectionJson(jobs: Int, debug: Boolean): String = {
+    val props = sys.props
+    val millPath = props
+      .get("mill.main.cli")
+      // we assume, the classpath is an executable jar here
+      .orElse(props.get("java.class.path"))
+      .getOrElse(throw new IllegalStateException("System property 'java.class.path' not set"))
 
-    worker match {
-      case Result.Success(worker) =>
-        worker.startBspServer(
-          initialEvaluator,
-          streams,
-          logStream.getOrElse(streams.err),
-          workspaceDir / Constants.bspDir,
-          canReload,
-          Seq(millServerHandle) ++ serverHandle.toSeq
-        )
-      case f: Result.Failure[_] =>
-        streams.err.println("Failed to start the BSP worker. " + f.msg)
-        BspServerResult.Failure
-      case f: Result.Exception =>
-        streams.err.println("Failed to start the BSP worker. " + f.throwable)
-        BspServerResult.Failure
-      case f =>
-        streams.err.println("Failed to start the BSP worker. " + f)
-        BspServerResult.Failure
-    }
+    upickle.default.write(
+      BspConfigJson(
+        name = "mill-bsp",
+        argv = Seq(
+          millPath,
+          "--bsp",
+          "--disable-ticker",
+          "--color",
+          "false",
+          "--jobs",
+          s"${jobs}"
+        ) ++ (if (debug) Seq("--debug") else Seq()),
+        millVersion = MillBuildInfo.millVersion,
+        bspVersion = Constants.bspProtocolVersion,
+        languages = Constants.languages
+      )
+    )
   }
 
 }

bsp/src/mill/bsp/BspConfigJson.scala

@@ -0,0 +1,15 @@
+package mill.bsp
+
+import upickle.default._
+
+private case class BspConfigJson(
+    name: String,
+    argv: Seq[String],
+    millVersion: String,
+    bspVersion: String,
+    languages: Seq[String]
+)
+
+private object BspConfigJson {
+  implicit val rw: ReadWriter[BspConfigJson] = macroRW
+}

bsp/src/mill/bsp/BspContext.scala

@@ -0,0 +1,73 @@
+package mill.bsp
+
+import mill.api.{DummyInputStream, Logger, SystemStreams, internal}
+import mill.eval.Evaluator
+
+import java.io.PrintStream
+import scala.util.control.NonFatal
+
+object BspContext {
+  @volatile var bspServerHandle: BspServerHandle = null
+}
+
+@internal
+class BspContext(streams: SystemStreams, bspLogStream: Option[PrintStream], home: os.Path) {
+  // BSP mode, run with a simple evaluator command to inject the evaluator
+  // The command returns when the server exists or the workspace should be reloaded
+  // if the `lastResult` is `ReloadWorkspace` we re-run the script in a loop
+
+  streams.err.println("Running in BSP mode with hardcoded startSession command")
+
+  streams.err.println("Trying to load BSP server...")
+  BspContext.bspServerHandle =
+    try {
+      startBspServer(
+        initialEvaluator = None,
+        streams = streams,
+        logStream = bspLogStream,
+        canReload = true
+      ) match {
+        case Left(err) => sys.error(err)
+        case Right(res) => res
+      }
+    } catch {
+      case NonFatal(e) =>
+        streams.err.println(s"Could not start BSP server. ${e.getMessage}")
+        throw e
+    }
+
+  streams.err.println("BSP server started")
+
+  def startBspServer(
+      initialEvaluator: Option[Evaluator],
+      streams: SystemStreams,
+      logStream: Option[PrintStream],
+      canReload: Boolean
+  ): Either[String, BspServerHandle] = {
+    val log: Logger = new Logger {
+      override def colored: Boolean = false
+      override def systemStreams: SystemStreams = new SystemStreams(
+        out = streams.out,
+        err = streams.err,
+        in = DummyInputStream
+      )
+
+      override def info(s: String): Unit = streams.err.println(s)
+      override def error(s: String): Unit = streams.err.println(s)
+      override def ticker(s: String): Unit = streams.err.println(s)
+      override def debug(s: String): Unit = streams.err.println(s)
+      override def debugEnabled: Boolean = true
+    }
+
+    BspWorker(os.pwd, home, log).flatMap { worker =>
+      os.makeDir.all(home / Constants.bspDir)
+      worker.startBspServer(
+        initialEvaluator,
+        streams,
+        logStream.getOrElse(streams.err),
+        home / Constants.bspDir,
+        canReload
+      )
+    }
+  }
+}

main/src/mill/main/BspServerHandle.scala → bsp/src/mill/bsp/BspServerHandle.scala

@@ -1,4 +1,4 @@
-package mill.main
+package mill.bsp
 
 import mill.eval.Evaluator
 

main/src/mill/main/BspServerResult.scala → bsp/src/mill/bsp/BspServerResult.scala

@@ -1,4 +1,4 @@
-package mill.main
+package mill.bsp
 
 import mill.api.internal
 

bsp/src/mill/bsp/BspWorker.scala

@@ -1,34 +1,21 @@
 package mill.bsp
 
-import mill.Agg
-import mill.api.{Ctx, PathRef, Result, internal}
-import mill.define.Task
+import mill.api.{Ctx, Logger, SystemStreams, internal}
 import mill.eval.Evaluator
-import mill.main.{BspServerHandle, BspServerResult}
-import mill.api.SystemStreams
+import os.Path
 
 import java.io.PrintStream
 import java.net.URL
-import scala.concurrent.Promise
-import scala.util.{Failure, Success, Try}
 
 @internal
 trait BspWorker {
-
-  def createBspConnection(
-      jobs: Int,
-      serverName: String
-  )(implicit ctx: Ctx): (PathRef, ujson.Value)
-
   def startBspServer(
       initialEvaluator: Option[Evaluator],
       streams: SystemStreams,
       logStream: PrintStream,
       logDir: os.Path,
-      canReload: Boolean,
-      serverHandles: Seq[Promise[BspServerHandle]]
-  ): BspServerResult
-
+      canReload: Boolean
+  ): Either[String, BspServerHandle]
 }
 
 @internal
@@ -37,57 +24,42 @@ object BspWorker {
   private[this] var worker: Option[BspWorker] = None
 
   def apply(
-      millCtx: Ctx.Workspace with Ctx.Home with Ctx.Log,
+      workspace: os.Path,
+      home0: os.Path,
+      log: Logger,
       workerLibs: Option[Seq[URL]] = None
-  ): Result[BspWorker] = {
+  ): Either[String, BspWorker] = {
     worker match {
-      case Some(x) => Result.Success(x)
+      case Some(x) => Right(x)
       case None =>
         val urls = workerLibs.map { urls =>
-          millCtx.log.debug("Using direct submitted worker libs")
+          log.debug("Using direct submitted worker libs")
           urls
         }.getOrElse {
           // load extra classpath entries from file
           val cpFile =
-            millCtx.workspace / Constants.bspDir / s"${Constants.serverName}-${mill.BuildInfo.millVersion}.resources"
-          if (!os.exists(cpFile)) return Result.Failure(
+            workspace / Constants.bspDir / s"${Constants.serverName}-${mill.BuildInfo.millVersion}.resources"
+          if (!os.exists(cpFile)) return Left(
             "You need to run `mill mill.bsp.BSP/install` before you can use the BSP server"
           )
 
           // TODO: if outdated, we could regenerate the resource file and re-load the worker
 
           // read the classpath from resource file
-          millCtx.log.debug(s"Reading worker classpath from file: ${cpFile}")
+          log.debug(s"Reading worker classpath from file: ${cpFile}")
           os.read(cpFile).linesIterator.map(u => new URL(u)).toSeq
         }
 
         // create classloader with bsp.worker and deps
-        val cl = mill.api.ClassLoader.create(urls, getClass().getClassLoader())(millCtx)
-
-        // check the worker version
-        Try {
-          val workerBuildInfo = cl.loadClass(Constants.bspWorkerBuildInfoClass)
-          workerBuildInfo.getMethod("millBspWorkerVersion").invoke(null)
-        } match {
-          case Success(mill.BuildInfo.millVersion) => // same as Mill, everything is good
-          case Success(workerVersion) =>
-            millCtx.log.error(
-              s"""BSP worker version ($workerVersion) does not match Mill version (${mill.BuildInfo.millVersion}).
-                 |You need to run `mill mill.bsp.BSP/install` again.""".stripMargin
-            )
-          case Failure(e) =>
-            millCtx.log.error(
-              s"""Could not validate worker version number.
-                 |Error message: ${e.getMessage}
-                 |""".stripMargin
-            )
-        }
+        val cl = mill.api.ClassLoader.create(urls, getClass().getClassLoader())(
+          new Ctx.Home { override def home: Path = home0 }
+        )
 
         val workerCls = cl.loadClass(Constants.bspWorkerImplClass)
         val ctr = workerCls.getConstructor()
         val workerImpl = ctr.newInstance().asInstanceOf[BspWorker]
         worker = Some(workerImpl)
-        Result.Success(workerImpl)
+        Right(workerImpl)
     }
   }
 

bsp/worker/src/mill/bsp/worker/BspCompileProblemReporter.scala

@@ -27,7 +27,7 @@ import scala.language.implicitConversions
  *                            back as part of the published diagnostics
  *                            as well as compile report
  */
-class BspCompileProblemReporter(
+private class BspCompileProblemReporter(
     client: bsp.BuildClient,
     targetId: BuildTargetIdentifier,
     targetDisplayName: String,

bsp/worker/src/mill/bsp/worker/BspTestReporter.scala

@@ -39,7 +39,7 @@ import java.io.{PrintWriter, StringWriter}
  *                  in case special arguments need to be passed to
  *                  the compiler before running the test task.
  */
-class BspTestReporter(
+private class BspTestReporter(
     client: BuildClient,
     targetId: BuildTargetIdentifier,
     taskId: TaskId,