From cea2984ad648c48161ce4a3fe0ccb63f6e6c6444 Mon Sep 17 00:00:00 2001
From: purocean <purocean@gmail.com>
Date: Sun, 21 Jul 2024 22:04:48 +0800
Subject: [PATCH] refactor: improve JavaScript code execution performance and
 abort functionality

---
 src/renderer/plugins/code-runners.tsx     | 31 +++++++++++++++++------
 src/renderer/plugins/markdown-code-run.ts |  6 +++--
 2 files changed, 27 insertions(+), 10 deletions(-)

diff --git a/src/renderer/plugins/code-runners.tsx b/src/renderer/plugins/code-runners.tsx
index 2274171c..0f6453df 100644
--- a/src/renderer/plugins/code-runners.tsx
+++ b/src/renderer/plugins/code-runners.tsx
@@ -35,13 +35,15 @@ let javascriptWorker: Worker | null = null
 
 class JavascriptIframeExecutor implements ReadableStreamDefaultReader<string> {
   private code: string
+  private signal?: AbortSignal
   private _readResolve: (value: string) => void = () => 0
 
   closed: Promise<undefined>
   _state: 'pending' | 'done' | 'error'
 
-  constructor (code: string) {
+  constructor (code: string, opts?: { signal?: AbortSignal }) {
     this.code = code
+    this.signal = opts?.signal
     this._state = 'pending'
     this.closed = this.runCode()
 
@@ -54,7 +56,7 @@ class JavascriptIframeExecutor implements ReadableStreamDefaultReader<string> {
     })
   }
 
-  private getIframeWindow () {
+  private getIframe () {
     const id = 'code-runner-javascript-vm'
 
     // clean up
@@ -64,15 +66,21 @@ class JavascriptIframeExecutor implements ReadableStreamDefaultReader<string> {
     iframe.id = id
     iframe.style.display = 'none'
     document.body.appendChild(iframe)
-    const iframeWindow = iframe.contentWindow! as Window & typeof globalThis
 
+    const iframeWindow = iframe.contentWindow! as Window & typeof globalThis
     iframeWindow.ctx = window.ctx
 
-    return iframeWindow
+    return iframe
   }
 
   private async runCode (): Promise<undefined> {
-    const iframeWindow = this.getIframeWindow()
+    const iframe = this.getIframe()
+    const iframeWindow = iframe.contentWindow! as Window & typeof globalThis
+
+    this.signal?.addEventListener('abort', () => {
+      iframe.remove()
+    })
+
     const xConsole = getConsole(iframeWindow.console, val => this._readResolve(val))
 
     const AsyncFunction = iframeWindow.eval('(async function(){}).constructor')
@@ -113,6 +121,7 @@ class JavascriptIframeExecutor implements ReadableStreamDefaultReader<string> {
 
 class JavascriptWorkerExecutor implements ReadableStreamDefaultReader<string> {
   private code: string
+  private signal?: AbortSignal
   private _readResolve: (value: string) => void = () => 0
   private workerScript = `
     const getConsole = ${getConsole.toString()}
@@ -134,8 +143,9 @@ class JavascriptWorkerExecutor implements ReadableStreamDefaultReader<string> {
   closed: Promise<undefined>
   _state: 'pending' | 'done' | 'error'
 
-  constructor (code: string) {
+  constructor (code: string, opts?: { signal?: AbortSignal }) {
     this.code = code
+    this.signal = opts?.signal
     this._state = 'pending'
     this.closed = this.runCode()
 
@@ -177,6 +187,11 @@ class JavascriptWorkerExecutor implements ReadableStreamDefaultReader<string> {
       }
 
       worker.postMessage({ code: this.code })
+
+      this.signal?.addEventListener('abort', () => {
+        worker.terminate()
+        resolve()
+      })
     })
 
     await sleep(0)
@@ -224,13 +239,13 @@ export default {
       getTerminalCmd () {
         return null
       },
-      async run (_, code) {
+      async run (_, code, opts) {
         const firstLine = code.split('\n')[0].trim()
         const noWorker = firstLine.includes('--no-worker--')
         const outputHtml = firstLine.includes('--output-html--')
         return {
           type: outputHtml ? 'html' : 'plain',
-          value: noWorker ? new JavascriptIframeExecutor(code) : new JavascriptWorkerExecutor(code)
+          value: noWorker ? new JavascriptIframeExecutor(code, opts) : new JavascriptWorkerExecutor(code, opts)
         }
       },
     })
diff --git a/src/renderer/plugins/markdown-code-run.ts b/src/renderer/plugins/markdown-code-run.ts
index fbdfe4e8..b64f8e96 100644
--- a/src/renderer/plugins/markdown-code-run.ts
+++ b/src/renderer/plugins/markdown-code-run.ts
@@ -91,10 +91,12 @@ const RunCode = defineComponent({
       result.value = t('code-run.running')
 
       try {
-        if (!runner.value.nonInterruptible) {
-          abortController.value = new AbortController()
+        if (abortController.value) {
+          abort()
         }
 
+        abortController.value = new AbortController()
+
         const { type, value: val } = await runner.value.run(language!, code, { signal: abortController.value?.signal })
 
         if (typeof val === 'string') {