feat: add optional init function

targos · targos · commit 6738b469c05e · 2019-06-22T21:33:41.000+02:00
diff --git a/src/Ework.ts b/src/Ework.ts
@@ -25,6 +25,11 @@ export interface IEworkOptions {
    * Default: 1.
    */
   minFreeThreads?: number;
+  /**
+   * Initialization function that will be executed after spawning the worker and
+   * before sending jobs to it.
+   */
+  init?: () => any;
 }
 
 type WorkerResolveFn<Output> = (result: Output) => void;
@@ -43,13 +48,28 @@ interface IWorkerJob<Input, Output> {
   reject: WorkerRejectFn;
 }
 
-interface IWorkerMessage<Output> {
-  id: number;
-  type: 'result';
-  status: 'success' | 'error';
-  result?: Output;
-  error?: any;
-}
+type IWorkerMessage<Output> =
+  | {
+    type: 'result';
+    id: number;
+    status: 'success';
+    result: Output;
+  }
+  | {
+    type: 'result';
+    id: number;
+    status: 'error';
+    error: any;
+  }
+  | {
+    type: 'init';
+    status: 'success';
+  }
+  | {
+    type: 'init';
+    status: 'error';
+    error: any;
+  };
 
 export class Ework<Input, Output> {
   private totalWorkers: number;
@@ -69,7 +89,11 @@ export class Ework<Input, Output> {
       throw new TypeError('options must be an object');
     }
 
-    const { maxWorkers = numCpus, minFreeThreads = 1 } = options;
+    const {
+      maxWorkers = numCpus,
+      minFreeThreads = 1,
+      init = () => null,
+    } = options;
 
     if (!Number.isInteger(maxWorkers) || maxWorkers < 1) {
       throw new RangeError('options.maxWorkers must be a positive integer');
@@ -80,28 +104,41 @@ export class Ework<Input, Output> {
       );
     }
 
+    if (typeof init !== 'function') {
+      throw new TypeError('options.init must be a function');
+    }
+
+    const initString = init.toString();
     const workerString = worker.toString();
-    const workerCode = makeWorkerCode(workerString);
+    const workerCode = makeWorkerCode(initString, workerString);
 
     this.totalWorkers = Math.max(
       Math.min(maxWorkers, numCpus - minFreeThreads),
       1,
     );
-    this.freeWorkers = this.totalWorkers;
+    this.freeWorkers = 0;
 
     this.workers = [];
     for (let i = 0; i < this.totalWorkers; i++) {
       const worker = spawnWorker(workerCode);
       const workerObj: IWorker<Input, Output> = {
         worker,
-        isWorking: false,
+        isWorking: true,
         job: null,
       };
       addWorkerListener(
         worker,
         'message',
         (message: IWorkerMessage<Output>) => {
-          if (message.type === 'result') {
+          if (message.type === 'init') {
+            if (message.status === 'success') {
+              workerObj.isWorking = false;
+              this.freeWorkers++;
+              this.run();
+            } else {
+              // TODO handle init error
+            }
+          } else if (message.type === 'result') {
             const job = workerObj.job;
             if (job === null) {
               throw new Error('UNREACHABLE');
@@ -110,9 +147,6 @@ export class Ework<Input, Output> {
               throw new Error('UNREACHABLE');
             }
             if (message.status === 'success') {
-              if (message.result === undefined) {
-                throw new Error('UNREACHABLE');
-              }
               job.resolve(message.result);
             } else {
               job.reject(new Error(message.error));
@@ -121,9 +155,12 @@ export class Ework<Input, Output> {
             workerObj.isWorking = false;
             this.freeWorkers++;
             this.run();
+          } else {
+            throw new Error('UNREACHABLE');
           }
         },
       );
+      worker.postMessage({ type: 'init' });
       this.workers.push(workerObj);
     }
 
diff --git a/src/__tests__/Ework.test.ts b/src/__tests__/Ework.test.ts
@@ -36,6 +36,18 @@ describe('constructor errors', () => {
       ).toThrow(/options\.minFreeThreads must be a positive integer or 0/);
     },
   );
+
+  it.each([null, 42, 'test'])(
+    'should throw if init is not a function (%s)',
+    (value: any) => {
+      expect(
+        () =>
+          new Ework(() => null, {
+            init: value,
+          }),
+      ).toThrow(/init must be a function/);
+    },
+  );
 });
 
 describe('execute', () => {
@@ -115,3 +127,28 @@ describe('map', () => {
     await worker.terminate();
   });
 });
+
+describe('init function', () => {
+  it('should execute and wait for the init function', async () => {
+    const worker = new Ework(
+      () => {
+        // @ts-ignore
+        return global.initValue;
+      },
+      {
+        init: async () => {
+          await new Promise((resolve) => {
+            setTimeout(() => {
+              // @ts-ignore
+              global.initValue = 42;
+              resolve();
+            }, 500);
+          });
+        },
+      },
+    );
+    const result = await worker.execute(null);
+    expect(result).toBe(42);
+    await worker.terminate();
+  });
+});
diff --git a/src/worker.browser.ts b/src/worker.browser.ts
@@ -28,48 +28,68 @@ export function addWorkerListener(
   worker.addEventListener(type, (event) => listener(event.data));
 }
 
-export function makeWorkerCode(workerString: string): string {
+export function makeWorkerCode(
+  initString: string,
+  workerString: string,
+): string {
   return `'use strict';
 
+const initFunction = ${initString};
 const workerFunction = ${workerString};
+const messagePort = self;
 
-self.onmessage = async (event) => {
+function postError(obj, error) {
+  let errorToTransfer = error;
+  if (error) {
+    if (error.stack) {
+      errorToTransfer = error.stack;
+    } else if (error.message) {
+      errorToTransfer = error.message;
+    }
+  }
+  try {
+    messagePort.postMessage(Object.assign({}, obj, {
+      status: 'error',
+      error: errorToTransfer
+    }));
+  } catch (e) {
+    messagePort.postMessage({
+      status: 'error',
+      error: 'Work failed but error could not be transferred: ' + e.message
+    });
+  }
+}
+
+messagePort.onmessage = async (event) => {
   const { data: message } = event;
   if (message.type === 'work') {
     try {
       const result = await workerFunction(message.value);
-      self.postMessage({
+      messagePort.postMessage({
         id: message.id,
         type: 'result',
         status: 'success',
         result
       });
     } catch (error) {
-      let errorToTransfer = error;
-      if (error) {
-        if (error.stack) {
-          errorToTransfer = error.stack;
-        } else if (error.message) {
-          errorToTransfer = error.message;
-        }
-      }
-      try {
-        self.postMessage({
-          id: message.id,
-          type: 'result',
-          status: 'error',
-          error: errorToTransfer
-        });
-      } catch (e) {
-        self.postMessage({
-          id: message.id,
-          type: 'result',
-          status: 'error',
-          error: 'Work failed but error could not be transferred: ' + e.message
-        });
-      }
+      postError({
+        id: message.id,
+        type: 'result'
+      }, error);
+    }
+  } else if (message.type === 'init') {
+    try {
+      await initFunction();
+      messagePort.postMessage({
+        type: 'init',
+        status: 'success'
+      });
+    } catch (error) {
+      postError({
+        type: 'init'
+      }, error);
     }
   }
 };
-  `;
+`;
 }
diff --git a/src/worker.ts b/src/worker.ts
@@ -31,49 +31,69 @@ export function addWorkerListener(
   worker.addListener(type, listener);
 }
 
-export function makeWorkerCode(workerString: string): string {
+export function makeWorkerCode(
+  initString: string,
+  workerString: string,
+): string {
   return `'use strict';
 
 const { parentPort } = require('worker_threads');
 
+const initFunction = ${initString};
 const workerFunction = ${workerString};
+const messagePort = parentPort;
 
-parentPort.on('message', async (message) => {
+function postError(obj, error) {
+  let errorToTransfer = error;
+  if (error) {
+    if (error.stack) {
+      errorToTransfer = error.stack;
+    } else if (error.message) {
+      errorToTransfer = error.message;
+    }
+  }
+  try {
+    messagePort.postMessage(Object.assign({}, obj, {
+      status: 'error',
+      error: errorToTransfer
+    }));
+  } catch (e) {
+    messagePort.postMessage({
+      status: 'error',
+      error: 'Work failed but error could not be transferred: ' + e.message
+    });
+  }
+}
+
+messagePort.on('message', async (message) => {
   if (message.type === 'work') {
     try {
       const result = await workerFunction(message.value);
-      parentPort.postMessage({
+      messagePort.postMessage({
         id: message.id,
         type: 'result',
         status: 'success',
         result
       });
     } catch (error) {
-      let errorToTransfer = error;
-      if (error) {
-        if (error.stack) {
-          errorToTransfer = error.stack;
-        } else if (error.message) {
-          errorToTransfer = error.message;
-        }
-      }
-      try {
-        parentPort.postMessage({
-          id: message.id,
-          type: 'result',
-          status: 'error',
-          error: errorToTransfer
-        });
-      } catch (e) {
-        parentPort.postMessage({
-          id: message.id,
-          type: 'result',
-          status: 'error',
-          error: 'Work failed but error could not be transferred: ' + e.message
-        });
-      }
+      postError({
+        id: message.id,
+        type: 'result'
+      }, error);
+    }
+  } else if (message.type === 'init') {
+    try {
+      await initFunction();
+      messagePort.postMessage({
+        type: 'init',
+        status: 'success'
+      });
+    } catch (error) {
+      postError({
+        type: 'init'
+      }, error);
     }
   }
 });
-  `;
+`;
 }
