Refactored BSP module and use isolated classloading for implementations

bsp/src/mill/bsp/BSP.scala

@@ -1,33 +1,34 @@
 package mill.bsp
 
-import ch.epfl.scala.bsp4j.BuildClient
-
-import java.io.{InputStream, PrintStream, PrintWriter}
-import java.nio.file.FileAlreadyExistsException
-import java.util.concurrent.Executors
-import mill.{BuildInfo, MillMain, T}
-import mill.define.{Command, Discover, ExternalModule}
-import mill.eval.Evaluator
-import mill.main.{BspServerHandle, BspServerResult}
-import org.eclipse.lsp4j.jsonrpc.Launcher
-
-import scala.concurrent.{Await, CancellationException, Promise}
-import upickle.default.write
-
+import java.io.{InputStream, PrintStream}
+import scala.concurrent.{Await, Promise}
 import scala.concurrent.duration.Duration
-import scala.util.Using
-import scala.util.chaining.scalaUtilChainingOps
+import mill.api.{Ctx, DummyInputStream, Logger, PathRef, Result}
+import mill.{Agg, T, BuildInfo => MillBuildInfo}
+import mill.define.{Command, Discover, ExternalModule, Task}
+import mill.eval.Evaluator
+import mill.main.{BspServerHandle, BspServerResult, BspServerStarter}
+import mill.scalalib.{CoursierModule, Dep}
+import mill.util.PrintLogger
+import os.Path
 
-object BSP extends ExternalModule {
+object BSP extends ExternalModule with CoursierModule with BspServerStarter {
   implicit def millScoptEvaluatorReads[T] = new mill.main.EvaluatorScopt[T]()
 
   lazy val millDiscover: Discover[this.type] = Discover[this.type]
-  val bspProtocolVersion = _root_.mill.bsp.BuildInfo.bsp4jVersion
-  val languages = Seq("scala", "java")
-  val serverName = "mill-bsp"
 
   private[this] val millServerHandle = Promise[BspServerHandle]()
 
+  private def bspWorkerIvyDeps: T[Agg[Dep]] = T {
+    Agg(Dep.parse(BuildInfo.millBspWorkerDep))
+  }
+
+  private def bspWorkerLibs: T[Agg[PathRef]] = T {
+    resolveDeps(T.task {
+      bspWorkerIvyDeps().map(bindDependency())
+    })()
+  }
+
   /**
    * Installs the mill-bsp server. It creates a json file
    * with connection details in the ./.bsp directory for
@@ -42,69 +43,15 @@ object BSP extends ExternalModule {
    * reason, the message and stacktrace of the exception will be
    * printed to stdout.
    */
-  def install(evaluator: Evaluator, jobs: Int = 1): Command[Unit] = T.command {
-    val bspDirectory = evaluator.rootModule.millSourcePath / ".bsp"
-    val bspFile = bspDirectory / s"${serverName}.json"
-    if (os.exists(bspFile)) T.log.info(s"Overwriting BSP connection file: ${bspFile}")
-    else T.log.info(s"Creating BSP connection file: ${bspFile}")
-    val withDebug = T.log.debugEnabled
-    if (withDebug) T.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."
-    )
-    os.write.over(bspFile, createBspConnectionJson(jobs, withDebug), createFolders = true)
-  }
-
-  @deprecated("Use other overload instead.", "Mill after 0.10.7")
-  def createBspConnectionJson(jobs: Int): String =
-    BSP.createBspConnectionJson(jobs: Int, debug = false)
-
-  // creates a Json with the BSP connection details
-  def createBspConnectionJson(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"))
-
-    write(
-      BspConfigJson(
-        name = "mill-bsp",
-        argv = Seq(
-          millPath,
-          "--bsp",
-          "--disable-ticker",
-          "--color",
-          "false",
-          "--jobs",
-          s"${jobs}"
-        ) ++ (if (debug) Seq("--debug") else Seq()),
-        millVersion = BuildInfo.millVersion,
-        bspVersion = bspProtocolVersion,
-        languages = languages
-      )
-    )
-  }
-
-  /**
-   * Computes a mill command which starts the mill-bsp
-   * server and establishes connection to client. Waits
-   * until a client connects and ends the connection
-   * after the client sent an "exit" notification
-   *
-   * @param ev Environment, used by mill to evaluate commands
-   * @return: mill.Command which executes the starting of the
-   *          server
-   */
-  def start(ev: Evaluator): Command[BspServerResult] = T.command {
-    startBspServer(
-      initialEvaluator = Some(ev),
-      outStream = MillMain.initialSystemStreams.out,
-      errStream = T.log.errorStream,
-      inStream = MillMain.initialSystemStreams.in,
-      logDir = ev.rootModule.millSourcePath / ".bsp",
-      canReload = false
+  def install(jobs: Int = 1): Command[Unit] = T.command {
+    // we create a file containing the additional jars to load
+    val cpFile = T.workspace / Constants.bspDir / s"${Constants.serverName}.resources"
+    os.write.over(
+      cpFile,
+      bspWorkerLibs().iterator.map(_.path.toNIO.toUri.toURL).mkString("\n"),
+      createFolders = true
     )
+    BspWorker(T.ctx()).map(_.createBspConnection(jobs, Constants.serverName))
   }
 
   /**
@@ -121,103 +68,49 @@ object BSP extends ExternalModule {
     res
   }
 
-  def startBspServer(
+  override def startBspServer(
       initialEvaluator: Option[Evaluator],
       outStream: PrintStream,
       errStream: PrintStream,
       inStream: InputStream,
-      logDir: os.Path,
+      workspaceDir: os.Path,
+      ammoniteHomeDir: os.Path,
       canReload: Boolean,
       serverHandle: Option[Promise[BspServerHandle]] = None
-  ) = {
-    val evaluator = initialEvaluator.map(_.withFailFast(false))
-
-    val millServer =
-      new MillBuildServer(
-        initialEvaluator = evaluator,
-        bspVersion = bspProtocolVersion,
-        serverVersion = BuildInfo.millVersion,
-        serverName = serverName,
-        logStream = errStream,
-        canReload = canReload
-      ) with MillJvmBuildServer with MillJavaBuildServer with MillScalaBuildServer
-
-    val executor = Executors.newCachedThreadPool()
-
-    var shutdownRequestedBeforeExit = false
-
-    try {
-      val launcher = new Launcher.Builder[BuildClient]()
-        .setOutput(outStream)
-        .setInput(inStream)
-        .setLocalService(millServer)
-        .setRemoteInterface(classOf[BuildClient])
-        .traceMessages(new PrintWriter(
-          (logDir / s"${serverName}.trace").toIO
-        ))
-        .setExecutorService(executor)
-        .create()
-      millServer.onConnectWithClient(launcher.getRemoteProxy)
-      val listening = launcher.startListening()
-      millServer.cancellator = shutdownBefore => {
-        shutdownRequestedBeforeExit = shutdownBefore
-        listening.cancel(true)
+  ): 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 errorStream: PrintStream = errStream
+        override def outputStream: PrintStream = errStream
+        override def inStream: InputStream = DummyInputStream
+        override def info(s: String): Unit = errStream.println(s)
+        override def error(s: String): Unit = errStream.println(s)
+        override def ticker(s: String): Unit = errStream.println(s)
+        override def debug(s: String): Unit = errStream.println(s)
+        override def debugEnabled: Boolean = true
       }
+    }
 
-      val bspServerHandle = new BspServerHandle {
-        private[this] var _lastResult: Option[BspServerResult] = None
-
-        override def runSession(evaluator: Evaluator): BspServerResult = {
-          _lastResult = None
-          millServer.updateEvaluator(Option(evaluator))
-          val onReload = Promise[BspServerResult]()
-          millServer.onSessionEnd = Some { serverResult =>
-            if (!onReload.isCompleted) {
-              errStream.println("Unsetting evaluator on session end")
-              millServer.updateEvaluator(None)
-              _lastResult = Some(serverResult)
-              onReload.success(serverResult)
-            }
-          }
-          Await.result(onReload.future, Duration.Inf).tap { r =>
-            errStream.println(s"Reload finished, result: ${r}")
-            _lastResult = Some(r)
-          }
-        }
-
-        override def lastResult: Option[BspServerResult] = _lastResult
-
-        override def stop(): Unit = {
-          errStream.println("Stopping server via handle...")
-          listening.cancel(true)
-        }
-      }
-      millServerHandle.success(bspServerHandle)
-      serverHandle.foreach(_.success(bspServerHandle))
-
-      listening.get()
-      ()
-    } catch {
-      case _: CancellationException =>
-        errStream.println("The mill server was shut down.")
-      case e: Exception =>
-        errStream.println(
-          s"""An exception occurred while connecting to the client.
-             |Cause: ${e.getCause}
-             |Message: ${e.getMessage}
-             |Exception class: ${e.getClass}
-             |Stack Trace: ${e.getStackTrace}""".stripMargin
+    val worker = BspWorker(ctx)
+
+    worker match {
+      case Result.Success(worker) =>
+        worker.startBspServer(
+          initialEvaluator,
+          outStream,
+          errStream,
+          inStream,
+          workspaceDir / Constants.bspDir,
+          canReload,
+          Seq(millServerHandle) ++ serverHandle.toSeq
         )
-    } finally {
-      errStream.println("Shutting down executor")
-      executor.shutdown()
-
+      case _ => BspServerResult.Failure
     }
-
-    val finalReuslt =
-      if (shutdownRequestedBeforeExit) BspServerResult.Shutdown
-      else BspServerResult.Failure
-
-    finalReuslt
   }
+
 }

bsp/src/mill/bsp/BspServerStarterImpl.scala

@@ -0,0 +1,7 @@
+package mill.bsp
+
+import mill.main.BspServerStarter
+
+object BspServerStarterImpl {
+  def get: BspServerStarter = BSP
+}

bsp/src/mill/bsp/BspWorker.scala

@@ -0,0 +1,80 @@
+package mill.bsp
+
+import mill.Agg
+import mill.api.{Ctx, PathRef, Result, internal}
+import mill.define.Task
+import mill.eval.Evaluator
+import mill.main.{BspServerHandle, BspServerResult}
+
+import java.io.{InputStream, 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): Unit
+
+  def startBspServer(
+      initialEvaluator: Option[Evaluator],
+      outStream: PrintStream,
+      errStream: PrintStream,
+      inStream: InputStream,
+      logDir: os.Path,
+      canReload: Boolean,
+      serverHandles: Seq[Promise[BspServerHandle]]
+  ): BspServerResult
+
+}
+
+@internal
+object BspWorker {
+
+  private[this] var worker: Option[BspWorker] = None
+
+  def apply(millCtx: Ctx.Workspace with Ctx.Home with Ctx.Log): Result[BspWorker] = {
+    worker match {
+      case Some(x) => Result.Success(x)
+      case None =>
+        // load extra classpath entries from file
+        val cpFile = millCtx.workspace / Constants.bspDir / s"${Constants.serverName}.resources"
+        if (!os.exists(cpFile)) return Result.Failure(
+          "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
+
+        val urls = 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)
+        val workerVersion = 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 workerCls = cl.loadClass(Constants.bspWorkerImplClass)
+        val ctr = workerCls.getConstructor()
+        val workerImpl = ctr.newInstance().asInstanceOf[BspWorker]
+        worker = Some(workerImpl)
+        Result.Success(workerImpl)
+    }
+  }
+
+}

bsp/src/mill/bsp/Constants.scala

@@ -0,0 +1,13 @@
+package mill.bsp
+
+import mill.api.internal
+
+@internal
+object Constants {
+  val bspDir = os.sub / ".bsp"
+  val bspProtocolVersion = BuildInfo.bsp4jVersion
+  val bspWorkerImplClass = "mill.bsp.worker.BspWorkerImpl"
+  val bspWorkerBuildInfoClass = "mill.bsp.worker.BuildInfo"
+  val languages = Seq("scala", "java")
+  val serverName = "mill-bsp"
+}

bsp/test/src/BspInstallDebugTests.scala

@@ -14,7 +14,7 @@ object BspInstallDebugTests extends ScriptTestSuite(false) {
     test("BSP install forwards --debug option to server") {
       val workspacePath = initWorkspace()
       eval("mill.bsp.BSP/install") ==> true
-      val jsonFile = workspacePath / ".bsp" / s"${BSP.serverName}.json"
+      val jsonFile = workspacePath / Constants.bspDir / s"${Constants.serverName}.json"
       os.exists(jsonFile) ==> true
       os.read(jsonFile).contains("--debug") ==> true
     }

bsp/test/src/BspInstallTests.scala

@@ -11,7 +11,7 @@ object BspInstallTests extends ScriptTestSuite(false) {
     test("BSP install") {
       val workspacePath = initWorkspace()
       eval("mill.bsp.BSP/install") ==> true
-      val jsonFile = workspacePath / ".bsp" / s"${BSP.serverName}.json"
+      val jsonFile = workspacePath / Constants.bspDir / s"${Constants.serverName}.json"
       os.exists(jsonFile) ==> true
       os.read(jsonFile).contains("--debug") ==> false
     }