wip: working methods proxy

Author: tegefaulkes

src/WorkerManager.ts

@@ -1,7 +1,8 @@
-import type { WorkerFactory, WorkerManifest } from './types.js';
+import type { WorkerFactory, WorkerManifest, WorkerResult } from './types.js';
+import type { TransferListItem } from 'node:worker_threads';
 import Logger from '@matrixai/logger';
 import { CreateDestroy, ready } from '@matrixai/async-init/CreateDestroy.js';
-import { type WorkerTaskInput } from './types.js';
+import { type WorkerTaskInformation } from './types.js';
 import * as errors from './errors.js';
 import WorkerPool from './WorkerPool.js';
 
@@ -42,9 +43,8 @@ class WorkerManager<T extends WorkerManifest> {
   }
 
   protected pool: WorkerPool;
-  protected manifest: T;
-  public proxy: T;
   protected logger: Logger;
+  public methods: T;
 
   public constructor({
     workerFactory,
@@ -58,17 +58,18 @@ class WorkerManager<T extends WorkerManifest> {
     logger: Logger;
   }) {
     this.logger = logger;
-    this.manifest = manifest;
     this.pool = new WorkerPool(cores, workerFactory);
-    this.proxy = new Proxy(this.manifest, {
-      apply: () => {
-        throw Error('TMP IMP NEVER');
-      },
-      get: (target: T, prop) => {
+    // TODO: Add support for running functions main thread?
+    this.methods = new Proxy(manifest, {
+      get: (target: T, prop: string | symbol) => {
         if (typeof prop === 'symbol') return;
-        return (data, transferList) => {
-          console.log(`${prop} called with `, data, transferList);
-          return this.call({ type: prop, data, transferList });
+        return async (
+          data: WorkerResult,
+          transferList: Array<TransferListItem>,
+        ) => {
+          const result = await this.call({ type: prop, data, transferList });
+          if (result.transferList == null) return { data: result.data };
+          return result;
         };
       },
     });
@@ -83,12 +84,12 @@ class WorkerManager<T extends WorkerManifest> {
   }
 
   @ready(new errors.ErrorWorkerManagerDestroyed())
-  public async call(task: WorkerTaskInput): Promise<unknown> {
+  public async call(task: WorkerTaskInformation): Promise<WorkerResult> {
     return await this.queue(task);
   }
 
   @ready(new errors.ErrorWorkerManagerDestroyed())
-  public queue(task: WorkerTaskInput): Promise<unknown> {
+  public queue(task: WorkerTaskInformation): Promise<WorkerResult> {
     return new Promise((resolve, reject) => {
       this.pool.runTask(task, (result, error) => {
         if (error != null) return reject(error);

src/WorkerPool.ts

@@ -5,7 +5,6 @@ import type {
   WorkerFactory,
   WorkerResultInternal,
   WorkerTaskInformation,
-  WorkerTaskInput,
 } from './types.js';
 import { AsyncResource } from 'node:async_hooks';
 import { Subject } from 'rxjs';
@@ -90,7 +89,10 @@ class WorkerPool {
       if (result.error != null) {
         worker[taskInfoSymbol].done(undefined, result.error);
       } else {
-        worker[taskInfoSymbol].done(result.data, undefined);
+        worker[taskInfoSymbol].done(
+          { data: result.data, transferList: result.transferList },
+          undefined,
+        );
       }
       worker[taskInfoSymbol] = undefined;
       this.freeWorkers.push(worker);
@@ -136,7 +138,7 @@ class WorkerPool {
     this.$workerCreated.next();
   }
 
-  public runTask(task: WorkerTaskInput, callback: TaskCallback) {
+  public runTask(task: WorkerTaskInformation, callback: TaskCallback) {
     if (this.terminatedError != null) throw this.terminatedError;
     if (this.freeWorkers.length === 0) {
       this.queue.push({ task, callback });
@@ -148,7 +150,10 @@ class WorkerPool {
     const worker = this.freeWorkers.pop()!;
     if (wasIdle) this.$poolStatus.next('working');
     worker[taskInfoSymbol] = new WorkerTask(callback);
-    worker.postMessage({ type: task.type, data: task.data }, task.transferList);
+    worker.postMessage(
+      { type: task.type, data: task.data, transferList: task.transferList },
+      task.transferList,
+    );
   }
 
   public async terminate(force: boolean) {

src/expose.ts

@@ -33,7 +33,10 @@ function expose(workerManifest: WorkerManifest) {
       parentPort!.postMessage({ error });
       return;
     }
-    parentPort!.postMessage({ data: result.data }, result.transferList);
+    parentPort!.postMessage(
+      { data: result.data, transferList: result.transferList },
+      result.transferList,
+    );
   };
   parentPort!.on('message', handleMessage);
   parentPort!.once('close', () => {

src/types.ts

@@ -5,29 +5,32 @@ type WorkerFactory = () => Worker;
 type WorkerTaskInformation<T extends string = string, D = unknown> = {
   type: T;
   data: D;
-};
-
-type WorkerTaskInput = WorkerTaskInformation & {
   transferList?: Array<TransferListItem>;
 };
 
+// TODO: structure this like a normal function as much as possible.
+// TODO: limit types to primitives and transferrables.
 type WorkerFunction<I = unknown, O = unknown> = (
   data: I,
   transferList?: Array<TransferListItem>,
 ) => Promise<WorkerResult<O>>;
-type WorkerResult<O = unknown> = { data: O };
+type WorkerResult<O = unknown> = {
+  data: O;
+  transferList?: Array<TransferListItem>;
+};
 type WorkerResultInternal<O = unknown> = WorkerResult<O> & {
   error?: Error;
-  transferList?: Array<TransferListItem>;
 };
 type WorkerManifest = Record<string, WorkerFunction>;
 
-type TaskCallback<O = unknown> = (result: O, error?: Error) => void;
+type TaskCallback<O = unknown> = (
+  result: WorkerResult<O>,
+  error?: Error,
+) => void;
 
 export type {
   WorkerFactory,
   WorkerTaskInformation,
-  WorkerTaskInput,
   WorkerFunction,
   WorkerResult,
   WorkerResultInternal,

src/worker.ts

@@ -4,7 +4,7 @@ import type { WorkerManifest } from '#types.js';
 import { expose } from './expose.js';
 
 const worker = {
-  test: async () => {
+  test: async (): Promise<{ data: 'hello world!' }> => {
     return { data: 'hello world!' };
   },
   add: async (data: { a: number; b: number }) => {

tests/WorkerManager.test.ts

@@ -1,4 +1,4 @@
-import type { WorkerFactory } from '#types.js';
+import type { WorkerFactory, WorkerResult } from '#types.js';
 import { Worker } from 'node:worker_threads';
 import url from 'url';
 import path from 'node:path';
@@ -19,7 +19,7 @@ describe('WorkerManager', () => {
     return new Worker(path.join(dirname, '../dist/worker.js'));
   };
 
-  let workerManager: WorkerManager;
+  let workerManager: WorkerManager<typeof workerManifest>;
 
   afterEach(async () => {
     await workerManager?.destroy();
@@ -28,12 +28,13 @@ describe('WorkerManager', () => {
   test('async construction and async destroy', async () => {
     workerManager = await WorkerManager.createWorkerManager({
       workerFactory,
+      manifest: workerManifest,
       cores: 1,
       logger,
     });
     expect(workerManager[destroyed]).toBe(false);
-    expect(await workerManager.call({ type: 'test', data: undefined })).toBe(
-      'hello world!',
+    expect(await workerManager.call({ type: 'test', data: undefined })).toEqual(
+      { data: 'hello world!' },
     );
     await workerManager.destroy();
     expect(workerManager[destroyed]).toBe(true);
@@ -45,6 +46,7 @@ describe('WorkerManager', () => {
     await expect(
       WorkerManager.createWorkerManager({
         workerFactory,
+        manifest: workerManifest,
         cores: 0,
         logger,
       }),
@@ -53,11 +55,12 @@ describe('WorkerManager', () => {
   test('start with 1 worker core', async () => {
     workerManager = await WorkerManager.createWorkerManager({
       workerFactory,
+      manifest: workerManifest,
       cores: 1,
       logger,
     });
-    expect(await workerManager.call({ type: 'test', data: undefined })).toBe(
-      'hello world!',
+    expect(await workerManager.call({ type: 'test', data: undefined })).toEqual(
+      { data: 'hello world!' },
     );
     await workerManager.destroy();
   });
@@ -66,6 +69,7 @@ describe('WorkerManager', () => {
     // if you only use 1 core, this test will be much slower
     workerManager = await WorkerManager.createWorkerManager({
       workerFactory,
+      manifest: workerManifest,
       cores: 1,
       logger,
     });
@@ -77,16 +81,17 @@ describe('WorkerManager', () => {
     }
     const rs = await Promise.all(tasks);
     expect(rs.length).toBe(taskCount);
-    expect(rs.every((x) => x === undefined)).toBe(true);
+    expect(rs.every((x: WorkerResult) => x.data === undefined)).toBe(true);
     const r = await task;
-    expect(r).toBeUndefined();
+    expect(r).toEqual({ data: undefined });
     await workerManager.destroy();
   });
   test('queueing up tasks', async () => {
     // Use all possible cores
     // if you only use 1 core, this test will be much slower
     workerManager = await WorkerManager.createWorkerManager({
       workerFactory,
+      manifest: workerManifest,
       cores: 1,
       logger,
     });
@@ -95,10 +100,10 @@ describe('WorkerManager', () => {
     const t3 = workerManager.queue({ type: 'sleep', data: 500 });
     const t4 = workerManager.queue({ type: 'sleep', data: 500 });
     await workerManager.completed();
-    expect(await t1).toBeUndefined();
-    expect(await t2).toBeUndefined();
-    expect(await t3).toBeUndefined();
-    expect(await t4).toBeUndefined();
+    expect(await t1).toEqual({ data: undefined });
+    expect(await t2).toEqual({ data: undefined });
+    expect(await t3).toEqual({ data: undefined });
+    expect(await t4).toEqual({ data: undefined });
     void workerManager.queue({ type: 'sleep', data: 500 });
     void workerManager.queue({ type: 'sleep', data: 500 });
     void workerManager.queue({ type: 'sleep', data: 500 });
@@ -109,6 +114,7 @@ describe('WorkerManager', () => {
   test('zero-copy buffer transfer', async () => {
     workerManager = await WorkerManager.createWorkerManager({
       workerFactory,
+      manifest: workerManifest,
       cores: 1,
       logger,
     });
@@ -128,7 +134,7 @@ describe('WorkerManager', () => {
     // The input ArrayBuffer is detached so the length is now 0
     expect(input.byteLength).toBe(0);
     // The output should be filled with 0xF
-    expect(Buffer.from(output as ArrayBuffer)).toEqual(
+    expect(Buffer.from(output.data as ArrayBuffer)).toEqual(
       Buffer.alloc(inputBuffer.byteLength, 0xf),
     );
 
@@ -144,11 +150,23 @@ describe('WorkerManager', () => {
       manifest: workerManifest,
       logger,
     });
+    expect(await workerManager.methods.test()).toEqual({
+      data: 'hello world!',
+    });
+    expect(await workerManager.methods.add({ a: 1, b: 2 })).toEqual({
+      data: 3,
+    });
+    expect(await workerManager.methods.sub({ a: 1, b: 2 })).toEqual({
+      data: -1,
+    });
+    expect(await workerManager.methods.fac(5)).toEqual({ data: 120 });
+    expect(await workerManager.methods.sleep(10)).toEqual({ data: undefined });
     const arrayBuffer = new ArrayBuffer(100);
-    const result = await workerManager.proxy.transferBuffer(arrayBuffer, [
+    const result = await workerManager.methods.transferBuffer(arrayBuffer, [
       arrayBuffer,
     ]);
-    console.log(result);
+    expect(new Uint8Array(result.data)).toEqual(new Uint8Array(100).fill(0xf));
+
     await workerManager.destroy();
   });
 });

tests/WorkerPool.test.ts

@@ -46,38 +46,47 @@ describe('WorkerPool', () => {
         return resolve(result);
       });
     });
-    await expect(task1).resolves.toBe('hello world!');
+    await expect(task1).resolves.toEqual({
+      data: 'hello world!',
+      transferList: undefined,
+    });
 
     const task2 = new Promise((resolve, reject) => {
       pool.runTask({ type: 'add', data: { a: 2, b: 2 } }, (result, error) => {
         if (error != null) return reject(error);
         return resolve(result);
       });
     });
-    await expect(task2).resolves.toBe(4);
+    await expect(task2).resolves.toEqual({ data: 4, transferList: undefined });
 
     const task3 = new Promise((resolve, reject) => {
       pool.runTask({ type: 'fac', data: 5 }, (result, error) => {
         if (error != null) return reject(error);
         return resolve(result);
       });
     });
-    await expect(task3).resolves.toBe(120);
+    await expect(task3).resolves.toEqual({
+      data: 120,
+      transferList: undefined,
+    });
 
     const task4 = new Promise((resolve, reject) => {
       pool.runTask({ type: 'sleep', data: 10 }, (result, error) => {
         if (error != null) return reject(error);
         return resolve(result);
       });
     });
-    await expect(task4).resolves.toBe(undefined);
+    await expect(task4).resolves.toEqual({
+      data: undefined,
+      transferList: undefined,
+    });
   });
   test('WorkerPool handles failure to create worker', async () => {
     pool = new WorkerPool(3, () => new Worker('./badPath'));
-    const taskP = new Promise<string>((resolve, reject) => {
+    const taskP = new Promise<unknown>((resolve, reject) => {
       pool.runTask({ type: 'test', data: undefined }, (result, error) => {
         if (error != null) return reject(error);
-        return resolve(result as string);
+        return resolve(result);
       });
     });
     await expect(taskP).rejects.toThrow();