Support cancelling worksheet evaluation

Author: Duhemm

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

@@ -54,6 +54,7 @@ class DottyLanguageServer extends LanguageServer
 
   private[this] var rootUri: String = _
   private[this] var client: LanguageClient = _
+  private[this] val worksheets: mutable.Map[URI, CompletableFuture[_]] = mutable.Map.empty
 
   private[this] var myDrivers: mutable.Map[ProjectConfig, InteractiveDriver] = _
 
@@ -227,13 +228,25 @@ class DottyLanguageServer extends LanguageServer
     /*thisServer.synchronized*/ {}
 
   override def didSave(params: DidSaveTextDocumentParams): Unit = {
-    thisServer.synchronized {
-      val uri = new URI(params.getTextDocument.getUri)
-      if (isWorksheet(uri)) {
-        val driver = driverFor(uri)
+    val uri = new URI(params.getTextDocument.getUri)
+    if (isWorksheet(uri)) {
+      if (uri.getScheme == "cancel") {
+        val fileURI = new URI("file", uri.getUserInfo, uri.getHost, uri.getPort, uri.getPath, uri.getQuery, uri.getFragment)
+        worksheets.get(fileURI).foreach(_.cancel(true))
+      } else {
         val sendMessage = (msg: String) => client.logMessage(new MessageParams(MessageType.Info, uri + msg))
-        evaluateWorksheet(driver, uri, sendMessage)(driver.currentCtx)
-        sendMessage("FINISHED")
+        worksheets.put(
+          uri,
+          computeAsync { cancelChecker =>
+            try {
+              val driver = driverFor(uri)
+              evaluateWorksheet(driver, uri, sendMessage, cancelChecker)(driver.currentCtx)
+            } finally {
+              worksheets.remove(uri)
+              sendMessage("FINISHED")
+            }
+          }
+        )
       }
     }
   }
@@ -430,14 +443,19 @@ class DottyLanguageServer extends LanguageServer
   /**
    * Evaluate the worksheet at `uri`.
    *
-   * @param driver      The driver for the project that contains the worksheet.
-   * @param uri         The URI of the worksheet.
-   * @param sendMessage A mean of communicating the results of evaluation back.
+   * @param driver        The driver for the project that contains the worksheet.
+   * @param uri           The URI of the worksheet.
+   * @param sendMessage   A mean of communicating the results of evaluation back.
+   * @param cancelChecker Token to check whether evaluation was cancelled
    */
-  private def evaluateWorksheet(driver: InteractiveDriver, uri: URI, sendMessage: String => Unit)(implicit ctx: Context): Unit = {
+  private def evaluateWorksheet(driver: InteractiveDriver,
+                                uri: URI,
+                                sendMessage: String => Unit,
+                                cancelChecker: CancelChecker)(
+      implicit ctx: Context): Unit = {
     val trees = driver.openedTrees(uri)
     trees.headOption.foreach { tree =>
-      Worksheet.evaluate(tree, sendMessage)
+      Worksheet.evaluate(tree, sendMessage, cancelChecker)
     }
   }
 }

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

@@ -10,17 +10,27 @@ import dotty.tools.dotc.core.Flags.Synthetic
 
 import dotty.tools.repl.{ReplDriver, State}
 
+import org.eclipse.lsp4j.jsonrpc.CancelChecker
+
 import java.io.{ByteArrayOutputStream, PrintStream}
+import java.util.concurrent.{Callable, CancellationException, Executors, TimeUnit}
 
 object Worksheet {
 
+  private val executor = Executors.newFixedThreadPool(1)
+  private final val checkCancelledDelayMs = 50
+
   /**
    * Evaluate `tree` as a worksheet using the REPL.
    *
-   * @param tree        The top level object wrapping the worksheet.
-   * @param sendMessage A mean of communicating the results of evaluation back.
+   * @param tree          The top level object wrapping the worksheet.
+   * @param sendMessage   A mean of communicating the results of evaluation back.
+   * @param cancelChecker A token to check whether execution should be cancelled.
    */
-  def evaluate(tree: SourceTree, sendMessage: String => Unit)(implicit ctx: Context): Unit = {
+  def evaluate(tree: SourceTree,
+               sendMessage: String => Unit,
+               cancelChecker: CancelChecker)(
+      implicit ctx: Context): Unit = {
     tree.tree match {
       case td @ TypeDef(_, template: Template) =>
         val replOut = new ByteArrayOutputStream
@@ -32,11 +42,28 @@ object Worksheet {
             state
 
           case (state, statement) if executed.add(bounds(statement.pos)) =>
-            val (line, newState) = execute(repl, state, statement, tree.source)
-            val result = new String(replOut.toByteArray())
-            if (result.trim.nonEmpty) sendMessage(encode(result, line))
-            replOut.reset()
-            newState
+            val task = executor.submit(new Callable[State] {
+              override def call(): State = {
+                  val (line, newState) = execute(repl, state, statement, tree.source)
+                  val result = new String(replOut.toByteArray())
+                  if (result.trim.nonEmpty) sendMessage(encode(result, line))
+                  replOut.reset()
+                  newState
+              }
+            })
+
+            while (!task.isDone()) {
+              try {
+                cancelChecker.checkCanceled()
+                Thread.sleep(checkCancelledDelayMs)
+              } catch {
+                case _: CancellationException =>
+                  task.cancel(true)
+              }
+            }
+
+            try task.get(0, TimeUnit.SECONDS)
+            catch { case _: Throwable => state }
 
           case (state, statement) =>
             state

vscode-dotty/package.json

@@ -44,6 +44,10 @@
       {
         "command": "worksheet.evaluate",
         "title": "Run worksheet"
+      },
+      {
+        "command": "worksheet.cancel",
+        "title": "Cancel worksheet evaluation"
       }
     ],
     "configurationDefaults": {

vscode-dotty/src/extension.ts

@@ -209,6 +209,10 @@ function run(serverOptions: ServerOptions) {
     worksheet.worksheetSave(client)
   })
 
+  vscode.commands.registerCommand(worksheet.worksheetCancelEvaluationKey, () => {
+    worksheet.cancelExecution(client)
+  })
+
   // Push the disposable to the context's subscriptions so that the
   // client can be deactivated on extension deactivation
   extensionContext.subscriptions.push(client.start());

vscode-dotty/src/worksheet.ts

@@ -29,6 +29,11 @@ export const worksheetEvaluateKey = "worksheet.evaluate"
  */
 export const worksheetEvaluateAfterSaveKey = "worksheet.evaluateAfterSave"
 
+/**
+ * The command key for cancelling worksheet evaluation.
+ */
+export const worksheetCancelEvaluationKey = "worksheet.cancel"
+
 /** Is this document a worksheet? */
 export function isWorksheet(document: vscode.TextDocument): boolean {
   return document.fileName.endsWith(".sc")
@@ -112,16 +117,38 @@ function _prepareWorksheet(document: vscode.TextDocument) {
 
 function showWorksheetProgress(document: vscode.TextDocument) {
   return vscode.window.withProgress({
-    location: vscode.ProgressLocation.Window,
-    title: "Evaluating worksheet"
-  }, _ => {
+    location: vscode.ProgressLocation.Notification,
+    title: "Evaluating worksheet",
+    cancellable: true
+  }, (_, token) => {
+    token.onCancellationRequested(_ => {
+      vscode.commands.executeCommand(worksheetCancelEvaluationKey)
+    })
+
     function isFinished() {
       return worksheetFinished.get(document) || false
     }
     return wait(isFinished, 500)
   })
 }
 
+/**
+ * Cancel the execution of the worksheet.
+ *
+ * This sends a `textDocument/didSave` message to the language server
+ * with the scheme of the URI set to `cancel`. The language server
+ * interprets that as a request to cancel the execution of the specified worksheet.
+ */
+export function cancelExecution(client: LanguageClient) {
+  const editor = vscode.window.activeTextEditor
+  if (editor) {
+    const document = editor.document
+    client.sendNotification("textDocument/didSave", {
+      textDocument: { uri: document.uri.with({ scheme: "cancel" }).toString() }
+    })
+  }
+}
+
 /** Wait until `cond` evaluates to true; test every `delay` ms. */
 function wait(cond: () => boolean, delayMs: number): Promise<boolean> {
   const isFinished = cond()