worker: improve typing

Signed-off-by: Varun Patil <radialapps@gmail.com>

src/components/frame/XImgCache.ts

@@ -1,11 +1,10 @@
 import { API } from '@services/API';
 import { onDOMLoaded } from '@services/utils';
-import { workerImporter } from '@services/worker';
-import type * as w from './XImgWorker';
+import { importWorker } from '@services/worker';
+import type XImgWorker from './XImgWorker';
 
 // Global web worker to fetch images
-let worker: Worker;
-let importer: ReturnType<typeof workerImporter>;
+let worker: typeof XImgWorker;
 
 // Memcache for blob URLs
 const BLOB_CACHE = new Map<string, object>() as Map<string, [number, string]>;
@@ -17,12 +16,11 @@ const BLOB_STICKY = new Map<string, number>();
 function startWorker() {
   if (worker || _m.mode !== 'user') return;
 
-  // Start worker
-  worker = new Worker(new URL('./XImgWorkerStub.ts', import.meta.url));
-  importer = workerImporter(worker);
+  // Get typed worker
+  worker = importWorker(new Worker(new URL('./XImgWorkerStub.ts', import.meta.url)));
 
   // Configure worker
-  importer<typeof w.configure>('configure')({
+  worker.configure({
     multiUrl: API.IMAGE_MULTIPREVIEW(),
   });
 }
@@ -69,7 +67,7 @@ export async function fetchImage(url: string) {
   if (entry) return entry[1];
 
   // Fetch image
-  const blobUrl = await importer<typeof w.fetchImageSrc>('fetchImageSrc')(url);
+  const blobUrl = await worker.fetchImageSrc(url);
 
   // Check memcache entry again and revoke if it was added in the meantime
   if ((entry = BLOB_CACHE.get(url))) {

src/components/frame/XImgWorker.ts

@@ -1,5 +1,5 @@
 import { CacheExpiration } from 'workbox-expiration';
-import { workerExport } from '@services/worker';
+import { exportWorker } from '@services/worker';
 
 declare var self: ServiceWorkerGlobalScope;
 
@@ -312,14 +312,14 @@ async function fetchMultipreview(files: any[]) {
 
 /** Will be configured after the worker starts */
 let config: { multiUrl: string };
-export async function configure(_config: typeof config) {
+function configure(_config: typeof config) {
   config = _config;
 }
 
 /** Get BLOB url for image */
-export async function fetchImageSrc(url: string) {
+async function fetchImageSrc(url: string) {
   return URL.createObjectURL(await fetchImage(url));
 }
 
 // Exports to main thread
-workerExport({ fetchImageSrc, configure });
+export default exportWorker({ fetchImageSrc, configure });

src/services/worker.ts

@@ -1,51 +1,105 @@
-/** Set the receiver function for a worker */
-export function workerExport(handlers: Record<string, (...data: any) => Promise<any>>): void {
-  /** Promise API for web worker */
-  self.onmessage = async ({
-    data,
-  }: {
-    data: {
-      id: number;
+/**
+ * Data sent from main thread to worker.
+ */
+type CommRequest = {
+  reqid: number;
   name: string;
   args: any[];
-    };
-  }) => {
+};
+
+/**
+ * Data sent from worker to main thread.
+ */
+type CommResult = {
+  reqid: number;
+  resolve?: any;
+  reject?: string;
+};
+
+/**
+ * Export methods from a worker to the main thread.
+ *
+ * @param handlers Object with methods to export
+ *
+ * @example
+ * ```ts
+ * // my-worker.ts
+ * function foo() { return 'bar'; }
+ *
+ * async function asyncFoo() { return 'bar'; }
+ *
+ * export default exportWorker({
+ *  foo,
+ *  asyncFoo,
+ *  inline: () => 'bar',
+ * });
+ */
+export function exportWorker<T extends { [name: string]: Function }>(handlers: T): T {
+  self.onmessage = async ({ data }: { data: CommRequest }) => {
     try {
+      // Get handler from registrations
       const handler = handlers[data.name];
-      if (!handler) throw new Error(`No handler for type ${data.name}`);
-      const res = await handler.apply(self, data.args);
-      self.postMessage({
-        id: data.id,
-        resolve: res,
-      });
+      if (!handler) throw new Error(`[BUG] No handler for type ${data.name}`);
+
+      // Run handler
+      let result = handler.apply(self, data.args);
+      if (result instanceof Promise) {
+        result = await result;
+      }
+
+      // Success - post back to main thread
+      self.postMessage({ reqid: data.reqid, resolve: result } as CommResult);
     } catch (e) {
-      self.postMessage({
-        id: data.id,
-        reject: e.message,
-      });
+      // Error - post back rejection
+      self.postMessage({ reqid: data.reqid, reject: e.message } as CommResult);
     }
   };
+
+  return null as unknown as T;
 }
 
-/** Get the CALL function for a worker. Call this only once. */
-export function workerImporter(worker: Worker) {
-  const promises = new Map<number, { resolve: any; reject: any }>();
+/**
+ * Import a worker exported with `exportWorker`.
+ *
+ * @param worker Worker to import
+ *
+ * @example
+ * ```ts
+ * // main.ts
+ * import type MyWorker from './my-worker.ts';
+ *
+ * const worker = importWorker<typeof MyWorker>(new Worker(new URL('./XImgWorkerStub.ts', import.meta.url)));
+ *
+ * async (() => {
+ *   // all methods are async
+ *   console.assert(await worker.foo() === 'bar');
+ *   console.assert(await worker.asyncFoo() === 'bar');
+ *   console.assert(await worker.inline() === 'bar');
+ * });
+ */
+export function importWorker<T>(worker: Worker) {
+  const promises = new Map<number, { resolve: Function; reject: Function }>();
 
-  worker.onmessage = ({ data }: { data: any }) => {
-    const { id, resolve, reject } = data;
-    if (resolve) promises.get(id)?.resolve(resolve);
-    if (reject) promises.get(id)?.reject(reject);
-    promises.delete(id);
+  // Handle messages from worker
+  worker.onmessage = ({ data }: { data: CommResult }) => {
+    const { reqid, resolve, reject } = data;
+    if (resolve) promises.get(reqid)?.resolve(resolve);
+    if (reject) promises.get(reqid)?.reject(reject);
+    promises.delete(reqid);
   };
 
-  type PromiseFun = (...args: any) => Promise<any>;
-  return function importer<F extends PromiseFun>(name: string) {
-    return async function fun(...args: Parameters<F>) {
-      return await new Promise<ReturnType<Awaited<F>>>((resolve, reject) => {
-        const id = Math.random();
-        promises.set(id, { resolve, reject });
-        worker.postMessage({ id, name, args });
+  // Create proxy to call worker methods
+  const proxy = new Proxy(worker, {
+    get(target: Worker, name: string) {
+      return async function wrapper(...args: any[]) {
+        return await new Promise((resolve, reject) => {
+          const reqid = Math.random();
+          promises.set(reqid, { resolve, reject });
+          target.postMessage({ reqid, name, args } as CommRequest);
         });
       };
-  };
+    },
+  });
+
+  return proxy as T;
 }