refactor: use SharedWorker (closes #13)

aaronshaf · aaronshaf · commit d64da157f749 · 2024-11-25T15:23:42.000-07:00
diff --git a/packages/idb-cache/src/encryptionTasks.ts b/packages/idb-cache/src/encryptionTasks.ts
@@ -8,7 +8,7 @@ import {
 } from "./errors";
 
 /**
- * Encrypts a chunk of data using the worker.
+ * Encrypts a chunk of data using the SharedWorker.
  * @param port - The MessagePort instance.
  * @param value - The plaintext string to encrypt.
  * @param pendingRequests - Map of pending requests awaiting responses.
@@ -50,7 +50,7 @@ export async function encryptChunk(
 }
 
 /**
- * Decrypts a chunk of data using the worker.
+ * Decrypts a chunk of data using the SharedWorker.
  * @param port - The MessagePort instance.
  * @param iv - The Initialization Vector used during encryption.
  * @param ciphertext - The encrypted data.
diff --git a/packages/idb-cache/src/encryptionWorkerFn.ts b/packages/idb-cache/src/encryptionWorkerFn.ts
@@ -3,11 +3,12 @@
 
 import type { WorkerMessage } from "./types";
 
+declare const self: SharedWorkerGlobalScope;
+
 export function encryptionWorkerFunction() {
   let cacheKey: Uint8Array | null = null;
   const derivedKeyCache: Map<string, CryptoKey> = new Map();
   let pbkdf2Iterations = 100000;
-  let port: MessagePort | null = null;
   let fixedSalt: ArrayBuffer | null = null;
 
   async function getKeyFromCacheKey(
@@ -59,12 +60,16 @@ export function encryptionWorkerFunction() {
       throw new Error("Cache key not provided for encryption worker");
     }
     try {
-      port?.postMessage({ type: "ready" });
+      for (const port of ports) {
+        port.postMessage({ type: "ready" });
+      }
     } catch (error: unknown) {
       console.error("Worker: Failed to initialize AES key:", error);
       const errorMessage =
         error instanceof Error ? error.message : "Unknown initialization error";
-      port?.postMessage({ type: "initError", error: errorMessage });
+      for (const port of ports) {
+        port.postMessage({ type: "initError", error: errorMessage });
+      }
     }
   }
 
@@ -158,7 +163,7 @@ export function encryptionWorkerFunction() {
     }
   }
 
-  function handleEncrypt(requestId: string, value: string) {
+  function handleEncrypt(requestId: string, value: string, port: MessagePort) {
     enqueueTask(async () => {
       try {
         const encrypted = await encrypt(value);
@@ -186,7 +191,8 @@ export function encryptionWorkerFunction() {
   function handleDecrypt(
     requestId: string,
     iv: ArrayBuffer,
-    ciphertext: ArrayBuffer
+    ciphertext: ArrayBuffer,
+    port: MessagePort
   ) {
     enqueueTask(async () => {
       try {
@@ -209,73 +215,71 @@ export function encryptionWorkerFunction() {
     });
   }
 
-  async function onMessage(e: MessageEvent<WorkerMessage>) {
-    const { type, payload, requestId } = e.data;
-
-    switch (type) {
-      case "initialize":
-        {
-          const {
-            cacheKey: incomingCacheKey,
-            pbkdf2Iterations: incomingIterations,
-            cacheBuster,
-          } = payload;
-          cacheKey = new TextEncoder().encode(incomingCacheKey);
-          pbkdf2Iterations = incomingIterations || 100000;
-          fixedSalt = new TextEncoder().encode(cacheBuster).buffer;
-          await initializeKey();
-        }
-        break;
-
-      case "encrypt":
-        {
-          const { value } = payload;
-          await handleEncrypt(requestId, value);
-        }
-        break;
+  const ports: MessagePort[] = [];
 
-      case "decrypt":
-        {
-          const { iv, ciphertext } = payload;
-          await handleDecrypt(requestId, iv, ciphertext);
-        }
-        break;
+  function onConnect(e: MessageEvent) {
+    const port = e.ports[0];
+    ports.push(port);
+    port.onmessage = (event: MessageEvent<WorkerMessage>) => {
+      const { type, payload, requestId } = event.data;
 
-      case "destroy":
-        {
-          if (cacheKey) {
-            cacheKey.fill(0);
-            cacheKey = null;
-          }
-          if (fixedSalt) {
-            const saltArray = new Uint8Array(fixedSalt);
-            saltArray.fill(0);
-            fixedSalt = null;
+      switch (type) {
+        case "initialize":
+          {
+            const {
+              cacheKey: incomingCacheKey,
+              pbkdf2Iterations: incomingIterations,
+              cacheBuster,
+            } = payload;
+            cacheKey = new TextEncoder().encode(incomingCacheKey);
+            pbkdf2Iterations = incomingIterations || 100000;
+            fixedSalt = new TextEncoder().encode(cacheBuster).buffer;
+            initializeKey().catch((error) => {
+              console.error("Worker: Initialization failed:", error);
+            });
           }
-          if (port) {
-            port.close();
-            port = null;
+          break;
+
+        case "encrypt":
+          {
+            const { value } = payload;
+            handleEncrypt(requestId, value, port);
           }
-          self.close();
-        }
-        break;
+          break;
 
-      default:
-        console.warn(
-          `Worker: Unknown message type received: ${type}. Ignoring the message.`
-        );
-    }
-  }
+        case "decrypt":
+          {
+            const { iv, ciphertext } = payload;
+            handleDecrypt(requestId, iv, ciphertext, port);
+          }
+          break;
 
-  function handleInit(e: MessageEvent) {
-    const { type } = e.data;
+        case "destroy":
+          {
+            if (cacheKey) {
+              cacheKey.fill(0);
+              cacheKey = null;
+            }
+            if (fixedSalt) {
+              const saltArray = new Uint8Array(fixedSalt);
+              saltArray.fill(0);
+              fixedSalt = null;
+            }
+            for (const p of ports) {
+              p.close();
+            }
+            self.close();
+          }
+          break;
 
-    if (type === "init" && e.ports && e.ports.length > 0) {
-      port = e.ports[0];
-      port.onmessage = onMessage;
-      port.start();
-    }
+        default:
+          console.warn(
+            `Worker: Unknown message type received: ${type}. Ignoring the message.`
+          );
+      }
+    };
+    port.start();
   }
 
-  self.onmessage = handleInit;
+  self.onconnect = onConnect;
 }
diff --git a/packages/idb-cache/src/index.ts b/packages/idb-cache/src/index.ts
@@ -104,7 +104,7 @@ const isSubtleCryptoSupported = crypto?.subtle;
 export class IDBCache implements IDBCacheInterface {
   dbReadyPromise: Promise<import("idb").IDBPDatabase<IDBCacheSchema>>;
   private storeName: STORE;
-  private worker: Worker | null = null;
+  private worker: SharedWorker | null = null;
   private port: MessagePort | null = null;
   private pendingRequests: Map<
     string,
@@ -183,7 +183,7 @@ export class IDBCache implements IDBCacheInterface {
   }
 
   /**
-   * Initializes the worker by creating it, setting up communication, and handling initialization.
+   * Initializes the SharedWorker by creating it, setting up communication, and handling initialization.
    * @param cacheKey - The cache key used for encryption/decryption.
    * @param cacheBuster - The cacheBuster used as a fixed salt.
    * @throws {WorkerInitializationError} If the worker fails to initialize.
@@ -324,7 +324,7 @@ export class IDBCache implements IDBCacheInterface {
 
             // Define key range for this baseKey
             const lowerBound = `${baseKey}-chunk-000000-`;
-            const upperBound = `${baseKey}-chunk-999999\uffff`;
+            const upperBound = `${baseKey}-chunk-999999￿`;
             const range = IDBKeyRange.bound(
               lowerBound,
               upperBound,
@@ -800,7 +800,7 @@ export class IDBCache implements IDBCacheInterface {
   }
 
   /**
-   * Destroys the IDBCache instance by clearing data (optional), releasing resources, and terminating the worker.
+   * Destroys the IDBCache instance by clearing data (optional), releasing resources, and terminating the SharedWorker.
    * @param options - Configuration options for destruction.
    * @param options.clearData - Whether to clear all cached data before destruction.
    * @throws {DatabaseError} If there is an issue accessing the database during data clearing.
@@ -831,7 +831,7 @@ export class IDBCache implements IDBCacheInterface {
       }
 
       if (this.worker) {
-        this.worker.terminate();
+        this.worker.port.close();
         this.worker = null;
       }
 
diff --git a/packages/idb-cache/src/workerUtils.ts b/packages/idb-cache/src/workerUtils.ts
@@ -14,6 +14,9 @@ import {
   IDBCacheError,
 } from "./errors";
 
+/**
+ * Utility type guards for Worker responses
+ */
 function isReadyResponse(
   message: WorkerResponse
 ): message is { type: "ready" } {
@@ -51,7 +54,7 @@ function isErrorResponse(
 }
 
 /**
- * Creates a worker from a given function and sets up initial communication.
+ * Creates a SharedWorker from a given function and sets up initial communication.
  * @param fn - The worker function to execute.
  * @param rejectAll - Function to call to reject all pending requests in case of failure.
  * @returns An object containing the worker instance and its message port.
@@ -60,38 +63,34 @@ export function createWorkerFromFunction(
   fn: () => void,
   rejectAll: (errorMessage: string) => void
 ): {
-  worker: Worker;
+  worker: SharedWorker;
   port: MessagePort;
 } {
   const blob = new Blob([`(${fn.toString()})()`], {
     type: "application/javascript",
   });
   const url = URL.createObjectURL(blob);
-  const worker = new Worker(url);
-
-  const channel = new MessageChannel();
+  const worker = new SharedWorker(url);
 
-  worker.postMessage({ type: "init" }, [channel.port2]);
+  const port = worker.port;
 
-  worker.onmessage = () => {
-    URL.revokeObjectURL(url);
-  };
+  port.start();
 
   worker.onerror = (event) => {
-    console.error("Worker encountered an error:", event.message);
-    rejectAll("Worker encountered an error and was terminated.");
-    worker.terminate();
+    console.error("SharedWorker encountered an error:", event.message);
+    rejectAll("SharedWorker encountered an error and was terminated.");
+    worker.port.close();
   };
 
-  channel.port1.onmessageerror = () => {
+  port.onmessageerror = () => {
     console.warn(
-      "MessagePort encountered a message error. Worker may have been terminated."
+      "MessagePort encountered a message error. SharedWorker may have been terminated."
     );
-    rejectAll("Worker was terminated unexpectedly.");
-    channel.port1.close();
+    rejectAll("SharedWorker was terminated unexpectedly.");
+    port.close();
   };
 
-  return { worker, port: channel.port1 };
+  return { worker, port };
 }
 
 /**
@@ -176,12 +175,14 @@ export function initializeWorker(
   };
 
   port.onmessageerror = (e: MessageEvent) => {
-    console.error("Worker encountered a message error:", e);
-    const error = new WorkerInitializationError("Worker failed to initialize");
+    console.error("SharedWorker encountered a message error:", e);
+    const error = new WorkerInitializationError(
+      "SharedWorker failed to communicate properly."
+    );
     rejectReady(error);
     rejectAllPendingRequests(
       pendingRequests,
-      "Worker encountered an error and was terminated."
+      "SharedWorker encountered an error and was terminated."
     );
     port.close();
   };
@@ -229,13 +230,13 @@ export async function sendMessageToWorker<T extends WorkerMessage["type"]>(
         port.postMessage(message);
       }
     } catch (error) {
-      console.error("Failed to post message to worker:", error);
+      console.error("Failed to post message to SharedWorker:", error);
       const pending = pendingRequests.get(requestId);
       if (pending) {
         clearTimeout(pending.timer);
         pending.reject(
           new WorkerInitializationError(
-            "Failed to communicate with the worker."
+            "Failed to communicate with the SharedWorker."
           )
         );
         pendingRequests.delete(requestId);
