Add Support for Post-interrupt Hooks (#1186)

Co-authored-by: Peter Christensen <pchristensen2@apple.com>
almond-sh · Jul 7, 2023 · 4ccc348 · 4ccc348
1 parent 92aa8a9
commit 4ccc348
Show file tree

Hide file tree

Showing 3 changed files with 33 additions and 1 deletion.
diff --git a/docs/pages/api-jupyter.md b/docs/pages/api-jupyter.md
@@ -154,6 +154,29 @@ kernel.publish.updateHtml("Got all items", id)
 
 ![](/demo/updatable.gif)
 
+### After Interrupt Hooks
+
+After Interrupt Hooks allow clean-up code to be run after cell
+execution is interrupted.
+```scala
+val sparkAmmonite = {
+  new AmmoniteSparkSessionBuilder()
+    .getOrCreate()
+}
+lazy val spark = {
+  NotebookSparkSession.builder()
+    .progress(enable=true, keep=false)
+    .logsInDeveloperConsole(false)
+    .getOrCreate()
+}
+lazy val sc = {
+  spark.sparkContext
+}
+kernel.afterInterruptHooks += { _ =>
+  sc.cancelAllJobs()
+}
+```
+
 ### Hooks
 
 Hooks allow to pre-process code right before it's executed. Use like

diff --git a/modules/scala/scala-interpreter/src/main/scala/almond/JupyterApiImpl.scala b/modules/scala/scala-interpreter/src/main/scala/almond/JupyterApiImpl.scala
@@ -89,4 +89,6 @@ final class JupyterApiImpl(
       true
     }
   }
+
+  val afterInterruptHooks = mutable.Buffer.empty[Any => Any]
 }
diff --git a/modules/scala/scala-interpreter/src/main/scala/almond/ScalaInterpreter.scala b/modules/scala/scala-interpreter/src/main/scala/almond/ScalaInterpreter.scala
@@ -147,9 +147,16 @@ final class ScalaInterpreter(
 
   override def interruptSupported: Boolean =
     true
-  override def interrupt(): Unit =
+  override def interrupt(): Unit = {
     execute0.interrupt()
 
+    try Function.chain(jupyterApi.afterInterruptHooks).apply(())
+    catch {
+      case NonFatal(e) =>
+        log.warn("Caught exception while trying to run after Interrupt hooks", e)
+    }
+  }
+
   override def supportComm: Boolean = true
   override def setCommHandler(commHandler0: CommHandler): Unit =
     commHandlerOpt = Some(commHandler0)