Run worksheets asynchronously

A worksheet in an infinite loop shouldn't block the whole language
server from responding to queries. We just need to be careful to lock
the language server when operating on trees since these operations may
not be thread-safe (for example, when a symbol gets forced). We could
also allow running multiple worksheets at the same time but this commit
doesn't implement that.

Author: smarter

language-server/src/dotty/tools/languageserver/DottyLanguageServer.scala

@@ -107,7 +107,7 @@ class DottyLanguageServer extends LanguageServer
     if (Memory.isCritical()) CompletableFutures.computeAsync { _ => restart() }
 
   /** The driver instance responsible for compiling `uri` */
-  def driverFor(uri: URI): InteractiveDriver = {
+  def driverFor(uri: URI): InteractiveDriver = thisServer.synchronized {
     val matchingConfig =
       drivers.keys.find(config => config.sourceDirectories.exists(sourceDir =>
         new File(uri.getPath).getCanonicalPath.startsWith(sourceDir.getCanonicalPath)))
@@ -133,10 +133,10 @@ class DottyLanguageServer extends LanguageServer
     CompletableFuture.completedFuture(new Object)
   }
 
-  def computeAsync[R](fun: CancelChecker => R): CompletableFuture[R] =
+  def computeAsync[R](fun: CancelChecker => R, synchronize: Boolean = true): CompletableFuture[R] =
     CompletableFutures.computeAsync { cancelToken =>
       // We do not support any concurrent use of the compiler currently.
-      thisServer.synchronized {
+      def computation(): R = {
         cancelToken.checkCanceled()
         checkMemory()
         try {
@@ -147,6 +147,10 @@ class DottyLanguageServer extends LanguageServer
             throw ex
         }
       }
+      if (synchronize)
+        synchronized { computation() }
+      else
+        computation()
     }
 
   override def initialize(params: InitializeParams) = computeAsync { cancelToken =>

language-server/src/dotty/tools/languageserver/Main.scala

@@ -73,7 +73,7 @@ object Main {
         .setInput(in)
         .setOutput(out)
         // For debugging JSON messages:
-        // .traceMessages(new java.io.PrintWriter(System.err, true))
+        //.traceMessages(new java.io.PrintWriter(System.err, true))
         .create();
 
     val client = launcher.getRemoteProxy()

language-server/src/dotty/tools/languageserver/worksheet/Evaluator.scala

@@ -28,7 +28,7 @@ private object Evaluator {
    * @param cancelChecker The token that indicates whether evaluation has been cancelled.
    * @return A JVM running the REPL.
    */
-  def get(cancelChecker: CancelChecker)(implicit ctx: Context): Option[Evaluator] = {
+  def get(cancelChecker: CancelChecker)(implicit ctx: Context): Option[Evaluator] = synchronized {
     val classpath = ctx.settings.classpath.value
     previousEvaluator match {
       case Some(cp, evaluator) if evaluator.isAlive() && cp == classpath =>

language-server/src/dotty/tools/languageserver/worksheet/Worksheet.scala

@@ -10,60 +10,68 @@ import dotty.tools.dotc.core.Flags.Synthetic
 
 import org.eclipse.lsp4j.jsonrpc.CancelChecker
 
-import java.util.concurrent.CancellationException
+import java.util.concurrent.{CompletableFuture, CancellationException}
 
 object Worksheet {
 
   /**
    * Run `tree` as a worksheet using the REPL.
    *
    * @param tree          The top level object wrapping the worksheet.
+   * @param treeLock      Object on which to lock when doing operations on trees.
    * @param sendMessage   A mean of communicating the results of evaluation back.
    * @param cancelChecker A token to check whether execution should be cancelled.
    */
   def run(tree: SourceTree,
+          treeLock: Object,
           sendMessage: (Int, String) => Unit,
           cancelChecker: CancelChecker)(
-      implicit ctx: Context): Unit = synchronized {
+    implicit ctx: Context): Unit = {
+    // For now, don't try to run multiple evaluators in parallel, this would require
+    // changes to the logic of Evaluator.get among other things.
+    Evaluator.synchronized {
+      Evaluator.get(cancelChecker) match {
+        case None =>
+          sendMessage(1, "Couldn't start the JVM.")
+        case Some(evaluator) =>
+          val queries = treeLock.synchronized {
+            tree.tree match {
+              case td @ TypeDef(_, template: Template) =>
+                val seen = collection.mutable.Set.empty[(Int, Int)]
 
-    Evaluator.get(cancelChecker) match {
-      case None =>
-        sendMessage(1, "Couldn't start JVM.")
-      case Some(evaluator) =>
-        tree.tree match {
-          case td @ TypeDef(_, template: Template) =>
-            val executed = collection.mutable.Set.empty[(Int, Int)]
-
-            template.body.foreach {
-              case statement: DefTree if statement.symbol.is(Synthetic) =>
-                ()
-
-              case statement if evaluator.isAlive() && executed.add(bounds(statement.pos)) =>
-                try {
-                  cancelChecker.checkCanceled()
-                  val (line, result) = execute(evaluator, statement, tree.source)
-                  if (result.nonEmpty) sendMessage(line, result)
-                } catch { case _: CancellationException => () }
-
-              case _ =>
-                ()
+                template.body.flatMap {
+                  case statement: DefTree if statement.symbol.is(Synthetic) =>
+                    None
+                  case statement if seen.add(bounds(statement.pos)) =>
+                    Some(query(statement, tree.source))
+                  case _ =>
+                    None
+                }
             }
-        }
+          }
+          queries.foreach { (line, code) =>
+            cancelChecker.checkCanceled()
+            val res = evaluator.eval(code).getOrElse("")
+            cancelChecker.checkCanceled()
+            if (res.nonEmpty)
+              sendMessage(line, res)
+          }
+      }
     }
   }
 
   /**
-   * Extract `tree` from the source and evaluate it in the REPL.
+   * Extract the line number and source code corresponding to this tree
    *
    * @param evaluator  The JVM that runs the REPL.
    * @param tree       The compiled tree to evaluate.
    * @param sourcefile The sourcefile of the worksheet.
-   * @return The line in the sourcefile that corresponds to `tree`, and the result.
    */
-  private def execute(evaluator: Evaluator, tree: Tree, sourcefile: SourceFile): (Int, String) = {
-    val source = sourcefile.content.slice(tree.pos.start, tree.pos.end).mkString
+  private def query(tree: Tree, sourcefile: SourceFile): (Int, String) = {
     val line = sourcefile.offsetToLine(tree.pos.end)
-    (line, evaluator.eval(source).getOrElse(""))
+    val source = sourcefile.content.slice(tree.pos.start, tree.pos.end).mkString
+    (line, source)
+    // CompletableFuture.supplyAsync(() => (line, evaluator.eval(source).getOrElse("")))
   }
 
   private def bounds(pos: Position): (Int, Int) = (pos.start, pos.end)

language-server/src/dotty/tools/languageserver/worksheet/WorksheetService.scala

@@ -13,28 +13,23 @@ import java.util.concurrent.{CompletableFuture, ConcurrentHashMap}
 @JsonSegment("worksheet")
 trait WorksheetService { thisServer: DottyLanguageServer =>
 
-  val worksheets: ConcurrentHashMap[URI, CompletableFuture[_]] = new ConcurrentHashMap()
-
   @JsonRequest
-  def run(params: WorksheetRunParams): CompletableFuture[WorksheetRunResult] = thisServer.synchronized {
-    val uri = new URI(params.textDocument.getUri)
-    val future =
-      computeAsync { cancelChecker =>
-        try {
-          val driver = driverFor(uri)
-          val sendMessage = (line: Int, msg: String) => client.publishOutput(WorksheetRunOutput(params.textDocument, line, msg))
-          runWorksheet(driver, uri, sendMessage, cancelChecker)(driver.currentCtx)
-          WorksheetRunResult(success = true)
-        } catch {
-          case _: Throwable =>
-            WorksheetRunResult(success = false)
-        } finally {
-          worksheets.remove(uri)
-        }
+  def run(params: WorksheetRunParams): CompletableFuture[WorksheetRunResult] =
+    computeAsync(synchronize = false, fun = { cancelChecker =>
+      val uri = new URI(params.textDocument.getUri)
+      try {
+        val driver = driverFor(uri)
+        val sendMessage =
+          (line: Int, msg: String) => client.publishOutput(WorksheetRunOutput(params.textDocument, line, msg))
+
+        runWorksheet(driver, uri, sendMessage, cancelChecker)(driver.currentCtx)
+        cancelChecker.checkCanceled()
+        WorksheetRunResult(success = true)
+      } catch {
+        case _: Throwable =>
+          WorksheetRunResult(success = false)
       }
-    worksheets.put(uri, future)
-    future
-  }
+    })
 
   /**
    * Run the worksheet at `uri`.
@@ -45,13 +40,13 @@ trait WorksheetService { thisServer: DottyLanguageServer =>
    * @param cancelChecker Token to check whether evaluation was cancelled
    */
   private def runWorksheet(driver: InteractiveDriver,
-                                uri: URI,
-                                sendMessage: (Int, String) => Unit,
-                                cancelChecker: CancelChecker)(
+                           uri: URI,
+                           sendMessage: (Int, String) => Unit,
+                           cancelChecker: CancelChecker)(
       implicit ctx: Context): Unit = {
-    val trees = driver.openedTrees(uri)
-    trees.headOption.foreach { tree =>
-      Worksheet.run(tree, sendMessage, cancelChecker)
+    val treeOpt = thisServer.synchronized {
+      driver.openedTrees(uri).headOption
     }
+    treeOpt.foreach(tree => Worksheet.run(tree, thisServer, sendMessage, cancelChecker))
   }
 }