Worker pool

src/hooks/usePython.ts

@@ -7,7 +7,7 @@ import {
   useState
 } from 'react'
 import { PythonContext, suppressedMessages } from '../providers/PythonProvider'
-import { proxy, Remote, wrap } from 'comlink'
+import { Remote } from 'comlink'
 import useFilesystem from './useFilesystem'
 
 import { Packages } from '../types/Packages'
@@ -18,25 +18,25 @@ interface UsePythonProps {
 }
 
 export default function usePython(props?: UsePythonProps) {
-  const { packages = {} } = props ?? {}
+  const { packages } = props ?? {}
 
-  const [isLoading, setIsLoading] = useState(false)
-  const [pyodideVersion, setPyodideVersion] = useState<string | undefined>()
+  const [runnerId, setRunnerId] = useState<string>()
   const [isRunning, setIsRunning] = useState(false)
   const [output, setOutput] = useState<string[]>([])
   const [stdout, setStdout] = useState('')
   const [stderr, setStderr] = useState('')
-  const [pendingCode, setPendingCode] = useState<string | undefined>()
   const [hasRun, setHasRun] = useState(false)
 
   const {
-    packages: globalPackages,
     timeout,
     lazy,
-    terminateOnCompletion
+    terminateOnCompletion,
+    loading,
+    getRunner,
+    run,
+    terminate
   } = useContext(PythonContext)
 
-  const workerRef = useRef<Worker>()
   const runnerRef = useRef<Remote<PythonRunner>>()
 
   const {
@@ -49,103 +49,47 @@ export default function usePython(props?: UsePythonProps) {
     watchedModules
   } = useFilesystem({ runner: runnerRef?.current })
 
-  const createWorker = () => {
-    const worker = new Worker(
-      new URL('../workers/python-worker', import.meta.url)
-    )
-    workerRef.current = worker
-  }
-
   useEffect(() => {
-    if (!lazy) {
-      // Spawn worker on mount
-      createWorker()
-    }
-
     // Cleanup worker on unmount
     return () => {
       cleanup()
     }
   }, [])
 
-  const allPackages = useMemo(() => {
-    const official = [
-      ...new Set([
-        ...(globalPackages.official ?? []),
-        ...(packages.official ?? [])
-      ])
-    ]
-    const micropip = [
-      ...new Set([
-        ...(globalPackages.micropip ?? []),
-        ...(packages.micropip ?? [])
-      ])
-    ]
-    return [official, micropip]
-  }, [globalPackages, packages])
-
-  const isReady = !isLoading && !!pyodideVersion
-
-  useEffect(() => {
-    if (workerRef.current && !isReady) {
-      const init = async () => {
-        try {
-          setIsLoading(true)
-          const runner: Remote<PythonRunner> = wrap(workerRef.current as Worker)
-          runnerRef.current = runner
-
-          await runner.init(
-            proxy((msg: string) => {
-              // Suppress messages that are not useful for the user
-              if (suppressedMessages.includes(msg)) {
-                return
-              }
-              setOutput((prev) => [...prev, msg])
-            }),
-            proxy(({ version }) => {
-              // The runner is ready once the Pyodide version has been set
-              setPyodideVersion(version)
-              console.debug('Loaded pyodide version:', version)
-            }),
-            allPackages
-          )
-        } catch (error) {
-          console.error('Error loading Pyodide:', error)
-        } finally {
-          setIsLoading(false)
-        }
-      }
-      init()
-    }
-  }, [workerRef.current])
-
   // Immediately set stdout upon receiving new input
   useEffect(() => {
     if (output.length > 0) {
       setStdout(output.join('\n'))
     }
   }, [output])
 
-  // React to ready state and run delayed code if pending
-  useEffect(() => {
-    if (pendingCode && isReady) {
-      const delayedRun = async () => {
-        await runPython(pendingCode)
-        setPendingCode(undefined)
-      }
-      delayedRun()
+  const cleanup = () => {
+    if (runnerId) {
+      terminate(runnerId)
     }
-  }, [pendingCode, isReady])
+    setIsRunning(false)
+    setRunnerId(undefined)
+    setHasRun(false)
+    setOutput([])
+    setStdout('')
+    setStderr('')
+  }
 
   // React to run completion and run cleanup if worker should terminate on completion
   useEffect(() => {
     if (terminateOnCompletion && hasRun && !isRunning) {
       cleanup()
-      setIsRunning(false)
-      setPyodideVersion(undefined)
     }
   }, [terminateOnCompletion, hasRun, isRunning])
 
+  const interruptExecution = () => {
+    cleanup()
+  }
+
+  const isReady = !!runnerId
+
+  const isLoading = loading && !isReady
+
   const pythonRunnerCode = `
 import sys
 
@@ -172,6 +116,15 @@ def run(code, preamble=''):
         print()
 `
 
+  // send a callback to the worker to handle the output
+  const handleOutput = (msg: string) => {
+    // Suppress messages that are not useful for the user
+    if (suppressedMessages.includes(msg)) {
+      return
+    }
+    setOutput((prev) => [...prev, msg])
+  }
+
   // prettier-ignore
   const moduleReloadCode = (modules: Set<string>) => `
 import importlib
@@ -190,29 +143,36 @@ del sys
       setStdout('')
       setStderr('')
 
-      if (lazy && !isReady) {
-        // Spawn worker and set pending code
-        createWorker()
-        setPendingCode(code)
+      let newRunnerId
+      if (!runnerId || terminateOnCompletion || packages) {
+        console.log('no runnerId, getting runner')
+        newRunnerId = await getRunner(handleOutput, packages)
+        setRunnerId(newRunnerId)
+      }
+
+      const r = runnerId || newRunnerId
+
+      if (!r) {
+        console.log('no runnerId, returning')
         return
       }
 
       code = `${pythonRunnerCode}\n\nrun(${JSON.stringify(
         code
       )}, ${JSON.stringify(preamble)})`
 
-      if (!isReady) {
-        throw new Error('Pyodide is not loaded yet')
-      }
+      // if (!isReady) {
+      //   throw new Error('Pyodide is not loaded yet')
+      // }
       let timeoutTimer
       try {
         setIsRunning(true)
         setHasRun(true)
         // Clear output
         setOutput([])
-        if (!isReady || !runnerRef.current) {
-          throw new Error('Pyodide is not loaded yet')
-        }
+        // if (!isReady || !runnerRef.current) {
+        //   throw new Error('Pyodide is not loaded yet')
+        // }
         if (timeout > 0) {
           timeoutTimer = setTimeout(() => {
             setStdout('')
@@ -221,9 +181,9 @@ del sys
           }, timeout)
         }
         if (watchedModules.size > 0) {
-          await runnerRef.current.run(moduleReloadCode(watchedModules))
+          await run(r, moduleReloadCode(watchedModules))
         }
-        await runnerRef.current.run(code)
+        await run(r, code)
         // eslint-disable-next-line
       } catch (error: any) {
         setStderr('Traceback (most recent call last):\n' + error.message)
@@ -232,27 +192,9 @@ del sys
         clearTimeout(timeoutTimer)
       }
     },
-    [lazy, isReady, timeout, watchedModules]
+    [runnerId, lazy, timeout, watchedModules]
   )
 
-  const interruptExecution = () => {
-    cleanup()
-    setIsRunning(false)
-    setPyodideVersion(undefined)
-    setOutput([])
-
-    // Spawn new worker
-    createWorker()
-  }
-
-  const cleanup = () => {
-    if (!workerRef.current) {
-      return
-    }
-    console.debug('Terminating worker')
-    workerRef.current.terminate()
-  }
-
   return {
     runPython,
     stdout,
@@ -269,3 +211,19 @@ del sys
     unwatchModules
   }
 }
+
+// const allPackages = useMemo(() => {
+//   const official = [
+//     ...new Set([
+//       ...(globalPackages.official ?? []),
+//       ...(packages.official ?? [])
+//     ])
+//   ]
+//   const micropip = [
+//     ...new Set([
+//       ...(globalPackages.micropip ?? []),
+//       ...(packages.micropip ?? [])
+//     ])
+//   ]
+//   return [official, micropip]
+// }, [globalPackages, packages])