Custom React Hook to run a task in a WebWorker and get the result back in a Promise.

useWebWorker.ts

import { useState, useEffect, useCallback, useRef } from 'react';

/**
 * WebWorker status flags
 * - idle: The worker is not running any task
 * - working: The worker is running a task
 */
type WebWorkerStatus = 'idle' | 'working';

/**
 * Custom hook to run a task in a WebWorker and get the result back in a Promise.
 * @param taskFunction - The function to run in the WebWorker thread
 * @returns An object with runTask function, terminate function, and status flag
 */
export const useWebWorker = <T,>(taskFunction: (...args: any[]) => T) => {
  const [status, setStatus] = useState<WebWorkerStatus>('idle');
  const workerRef = useRef<Worker | null>(null);

  // Cleanup the worker when the component unmounts
  useEffect(() => {
    return () => {
      if (workerRef.current) {
        workerRef.current.terminate();
        workerRef.current = null;
      }
    };
  }, []);

  // Run the task in the WebWorker and return the result in a Promise
  const runTask = useCallback(
    (...args: any[]): Promise<T> => {
      return new Promise((resolve, reject) => {
        setStatus('working');

        // Convert the task function to a string for the worker
        const taskFunctionString = taskFunction.toString();
        const blob = new Blob(
          [
            `
              onmessage = function(event) {
                const { func, args } = event.data;
                const taskFunction = new Function('return ' + func)();
                try {
                  const result = taskFunction(...args);
                  postMessage({ result });
                } catch (error) {
                  postMessage({ error: error.message });
                }
              }
            `
          ],
          { type: 'application/javascript' }
        );

        // Create a new worker with the blob URL and store it in a ref
        const worker = new Worker(URL.createObjectURL(blob));
        workerRef.current = worker;

        // Handle worker messages
        worker.onmessage = (event) => {
          const { result, error } = event.data;
          worker.terminate();
          workerRef.current = null;
          setStatus('idle');

          if (error) {
            reject(new Error(error));
          } else {
            resolve(result);
          }
        };

        // Handle worker errors
        worker.onerror = (e) => {
          worker.terminate();
          workerRef.current = null;
          setStatus('idle');
          reject(new Error(e.message));
        };

        // Send task function and arguments to the worker
        worker.postMessage({ func: taskFunctionString, args });
      });
    },
    [taskFunction]
  );

  // Terminate the worker
  const terminate = useCallback(() => {
    if (workerRef.current) {
      workerRef.current.terminate();
      workerRef.current = null;
      setStatus('idle');
    }
  }, []);

  return { runTask, terminate, status };
};