From c78ab753b50a565e341a63999a605e87ed8acf88 Mon Sep 17 00:00:00 2001
From: Alexandre Archambault <alexandre.archambault@gmail.com>
Date: Mon, 18 Jan 2021 11:13:18 +0100
Subject: [PATCH] Add basic BSP server integration test

---
 .../interp/script/AmmoniteBuildServer.scala   | 43 ++++++++++++++++---
 .../script/DummyBuildServerImplems.scala      |  3 --
 amm/src/main/scala/ammonite/Main.scala        |  8 ++--
 .../ammonite/integration/BasicTests.scala     | 13 ++++++
 4 files changed, 55 insertions(+), 12 deletions(-)

diff --git a/amm/interp/src/main/scala/ammonite/interp/script/AmmoniteBuildServer.scala b/amm/interp/src/main/scala/ammonite/interp/script/AmmoniteBuildServer.scala
index 27b392acf..868f101b1 100644
--- a/amm/interp/src/main/scala/ammonite/interp/script/AmmoniteBuildServer.scala
+++ b/amm/interp/src/main/scala/ammonite/interp/script/AmmoniteBuildServer.scala
@@ -17,7 +17,7 @@ import coursierapi.{Dependency, Repository}
 import org.eclipse.lsp4j.jsonrpc.Launcher
 
 import scala.collection.JavaConverters._
-import scala.concurrent.{ExecutionContext, Future}
+import scala.concurrent.{ExecutionContext, Future, Promise}
 import scala.util.{Failure, Success}
 import scala.util.control.NonFatal
 
@@ -458,6 +458,15 @@ class AmmoniteBuildServer(
       new CleanCacheResult("", true)
     }
 
+  private val shutdownPromise = Promise[Unit]()
+  def buildShutdown(): CompletableFuture[Object] =
+    nonBlocking {
+      if (!shutdownPromise.isCompleted)
+        shutdownPromise.success(())
+      null
+    }
+
+  def initiateShutdown: Future[Unit] = shutdownPromise.future
 }
 
 object AmmoniteBuildServer {
@@ -537,12 +546,29 @@ object AmmoniteBuildServer {
         .map(_.split('/').toSeq)
         .toSet
 
+  private def naiveJavaFutureToScalaFuture[T](
+    f: java.util.concurrent.Future[T]
+  ): Future[T] = {
+    val p = Promise[T]()
+    val t = new Thread {
+      setDaemon(true)
+      setName("ammonite-bsp-wait-for-exit")
+      override def run(): Unit =
+        p.complete {
+          try Success(f.get())
+          catch { case t: Throwable => Failure(t) }
+        }
+    }
+    t.start()
+    p.future
+  }
+
   def start(
     server: AmmoniteBuildServer,
     input: InputStream = System.in,
     output: OutputStream = System.out
-  ): Launcher[BuildClient] = {
-    val ec = Executors.newFixedThreadPool(4) // FIXME Daemon threads
+  ): (Launcher[BuildClient], Future[Unit]) = {
+    val ec = Executors.newFixedThreadPool(4, threadFactory("ammonite-bsp-jsonrpc"))
     val launcher = new Launcher.Builder[BuildClient]()
       .setExecutorService(ec)
       .setInput(input)
@@ -552,7 +578,14 @@ object AmmoniteBuildServer {
       .create()
     val client = launcher.getRemoteProxy
     server.onConnectWithClient(client)
-    launcher
-  }
+    val f = launcher.startListening()
 
+    val scalaEc = ExecutionContext.fromExecutorService(ec)
+    val futures = Seq(
+      naiveJavaFutureToScalaFuture(f).map(_ => ())(scalaEc),
+      server.initiateShutdown
+    )
+    val shutdownFuture = Future.firstCompletedOf(futures)(scalaEc)
+    (launcher, shutdownFuture)
+  }
 }
diff --git a/amm/interp/src/main/scala/ammonite/interp/script/DummyBuildServerImplems.scala b/amm/interp/src/main/scala/ammonite/interp/script/DummyBuildServerImplems.scala
index 04fa31674..f5c15ef30 100644
--- a/amm/interp/src/main/scala/ammonite/interp/script/DummyBuildServerImplems.scala
+++ b/amm/interp/src/main/scala/ammonite/interp/script/DummyBuildServerImplems.scala
@@ -6,9 +6,6 @@ import scala.collection.JavaConverters._
 
 private[script] trait DummyBuildServerImplems extends BuildServer with ScalaBuildServer {
 
-  def buildShutdown(): CompletableFuture[Object] =
-    CompletableFuture.completedFuture(null)
-
   def buildTargetResources(params: ResourcesParams): CompletableFuture[ResourcesResult] = {
     val items = params.getTargets.asScala.toList.map { target =>
       new ResourcesItem(target, List.empty[String].asJava)
diff --git a/amm/src/main/scala/ammonite/Main.scala b/amm/src/main/scala/ammonite/Main.scala
index f8677755f..f6fd6a055 100644
--- a/amm/src/main/scala/ammonite/Main.scala
+++ b/amm/src/main/scala/ammonite/Main.scala
@@ -17,6 +17,8 @@ import ammonite.util._
 import scala.annotation.tailrec
 import ammonite.runtime.ImportHook
 import coursierapi.Dependency
+import scala.concurrent.Await
+import scala.concurrent.duration.Duration
 
 
 
@@ -309,12 +311,10 @@ object Main{
               Seq("ammonite.interp.api.InterpBridge" -> "interp")
             ) ++ AmmoniteBuildServer.defaultImports
           )
-          val launcher = AmmoniteBuildServer.start(buildServer)
           printErr.println("Starting BSP server")
-          val f = launcher.startListening()
-          f.get()
+          val (launcher, shutdownFuture) = AmmoniteBuildServer.start(buildServer)
+          Await.result(shutdownFuture, Duration.Inf)
           printErr.println("BSP server done")
-          // FIXME Doesn't exit for now
           true
         }else{
 
diff --git a/integration/src/test/scala/ammonite/integration/BasicTests.scala b/integration/src/test/scala/ammonite/integration/BasicTests.scala
index 0582685c8..75d062669 100644
--- a/integration/src/test/scala/ammonite/integration/BasicTests.scala
+++ b/integration/src/test/scala/ammonite/integration/BasicTests.scala
@@ -270,5 +270,18 @@ object BasicTests extends TestSuite{
         assert(out.contains(expected))
       }
     }
+
+    test("BSP"){
+      val jsonrpc = """{"jsonrpc": "2.0", "id": 1, "method": "build/shutdown", "params": null}"""
+        .getBytes("UTF-8")
+      val input = Array(
+        s"Content-Length: ${jsonrpc.length}",
+        "\r\n" * 2
+      ).flatMap(_.getBytes("UTF-8")) ++ jsonrpc
+      val res = os.proc(TestUtils.executable, "--bsp").call(
+        stdin = input
+      )
+      assert(res.exitCode == 0)
+    }
   }
 }