Add thread-based interruption implementation

Author: ileasile

src/main/kotlin/org/jetbrains/kotlinx/jupyter/connection.kt

@@ -9,13 +9,15 @@ import kotlinx.serialization.json.JsonObject
 import kotlinx.serialization.json.decodeFromJsonElement
 import kotlinx.serialization.json.encodeToJsonElement
 import kotlinx.serialization.json.jsonObject
+import org.jetbrains.kotlinx.jupyter.exceptions.ReplException
 import org.zeromq.SocketType
 import org.zeromq.ZMQ
 import java.io.Closeable
 import java.io.IOException
 import java.security.SignatureException
 import javax.crypto.Mac
 import javax.crypto.spec.SecretKeySpec
+import kotlin.concurrent.thread
 import kotlin.math.min
 
 class JupyterConnection(val config: KernelConfig) : Closeable {
@@ -144,6 +146,49 @@ class JupyterConnection(val config: KernelConfig) : Closeable {
 
     var contextMessage: Message? = null
 
+    private val currentExecutions = HashSet<Thread>()
+
+    data class ConnectionExecutionResult<T>(
+        val result: T?,
+        val throwable: Throwable?,
+        val isInterrupted: Boolean,
+    )
+
+    fun <T> runExecution(body: () -> T): ConnectionExecutionResult<T> {
+        var execRes: T? = null
+        var execException: Throwable? = null
+        val execThread = thread {
+            try {
+                execRes = body()
+            } catch (e: Throwable) {
+                execException = e
+            }
+        }
+        currentExecutions.add(execThread)
+        execThread.join()
+        currentExecutions.remove(execThread)
+
+        val exception = execException
+        val isInterrupted = exception is ThreadDeath ||
+            (exception is ReplException && exception.cause is ThreadDeath)
+        return ConnectionExecutionResult(execRes, exception, isInterrupted)
+    }
+
+    /**
+     * We cannot use [Thread.interrupt] here because we have no way
+     * to control the code user executes. [Thread.interrupt] will do nothing for
+     * the simple calculation (like `while (true) 1`). Consider replacing with
+     * something more smart in the future.
+     */
+    fun interruptExecution() {
+        @Suppress("deprecation")
+        while (currentExecutions.isNotEmpty()) {
+            val execution = currentExecutions.firstOrNull()
+            execution?.stop()
+            currentExecutions.remove(execution)
+        }
+    }
+
     override fun close() {
         heartbeat.close()
         shell.close()

src/main/kotlin/org/jetbrains/kotlinx/jupyter/protocol.kt

@@ -192,7 +192,7 @@ class ErrorResponseWithMessage(
 fun JupyterConnection.Socket.controlMessagesHandler(msg: Message, repl: ReplForJupyter?) {
     when (msg.content) {
         is InterruptRequest -> {
-            log.warn("Interruption is not yet supported!")
+            connection.interruptExecution()
             send(makeReplyMessage(msg, MessageType.INTERRUPT_REPLY, content = msg.content))
         }
         is ShutdownRequest -> {
@@ -439,6 +439,12 @@ fun JupyterConnection.evalWithIO(repl: ReplForJupyter, srcMessage: Message, body
     val forkedError = getCapturingStream(err, JupyterOutType.STDERR, false)
     val userError = getCapturingStream(null, JupyterOutType.STDERR, true)
 
+    fun flushStreams() {
+        forkedOut.flush()
+        forkedError.flush()
+        userError.flush()
+    }
+
     val printForkedOut = PrintStream(forkedOut, false, "UTF-8")
     val printForkedErr = PrintStream(forkedError, false, "UTF-8")
     val printUserError = PrintStream(userError, false, "UTF-8")
@@ -453,26 +459,30 @@ fun JupyterConnection.evalWithIO(repl: ReplForJupyter, srcMessage: Message, body
     System.setIn(if (allowStdIn) stdinIn else DisabledStdinInputStream)
     try {
         return try {
-            val exec = body()
-            if (exec == null) {
-                AbortResponseWithMessage("NO REPL!")
-            } else {
-                forkedOut.flush()
-                forkedError.flush()
-                userError.flush()
-
-                try {
-                    val result = exec.resultValue?.toDisplayResult(repl.notebook)
-                    OkResponseWithMessage(result, exec.metadata)
-                } catch (e: Exception) {
-                    AbortResponseWithMessage("error:  Unable to convert result to a string: $e")
+            val (exec, execException, executionInterrupted) = runExecution(body)
+            when {
+                executionInterrupted -> {
+                    flushStreams()
+                    AbortResponseWithMessage("The execution was interrupted")
+                }
+                execException != null -> {
+                    throw execException
+                }
+                exec == null -> {
+                    AbortResponseWithMessage("NO REPL!")
+                }
+                else -> {
+                    flushStreams()
+                    try {
+                        val result = exec.resultValue?.toDisplayResult(repl.notebook)
+                        OkResponseWithMessage(result, exec.metadata)
+                    } catch (e: Exception) {
+                        AbortResponseWithMessage("error:  Unable to convert result to a string: $e")
+                    }
                 }
             }
         } catch (ex: ReplException) {
-            forkedOut.flush()
-            forkedError.flush()
-            userError.flush()
-
+            flushStreams()
             ErrorResponseWithMessage(
                 ex.render(),
                 ex.javaClass.canonicalName,
@@ -482,9 +492,7 @@ fun JupyterConnection.evalWithIO(repl: ReplForJupyter, srcMessage: Message, body
             )
         }
     } finally {
-        forkedOut.close()
-        forkedError.close()
-        userError.close()
+        flushStreams()
         System.setIn(`in`)
         System.setErr(err)
         System.setOut(out)