From 838898f1dfd2d2456d3e9a832d24a7ef6f53be70 Mon Sep 17 00:00:00 2001 From: Shinigami Date: Sun, 14 Feb 2021 18:18:34 +0100 Subject: [PATCH] De-duplicate code for workers (#154) --- src/utility-types.ts | 10 ++++++++-- src/worker/abstract-worker.ts | 36 ++++++++++++++++++++++++++++++++--- src/worker/cluster-worker.ts | 22 +-------------------- src/worker/thread-worker.ts | 35 ++-------------------------------- 4 files changed, 44 insertions(+), 59 deletions(-) diff --git a/src/utility-types.ts b/src/utility-types.ts index aa4d058c..eb3f9727 100644 --- a/src/utility-types.ts +++ b/src/utility-types.ts @@ -1,3 +1,6 @@ +import type { Worker } from 'cluster' +import type { MessagePort } from 'worker_threads' + /** * Make all properties in T non-readonly */ @@ -24,7 +27,10 @@ export type JSONArray = Array /** * Message object that is passed between worker and main worker. */ -export interface MessageValue { +export interface MessageValue< + Data = unknown, + MainWorker extends Worker | MessagePort | unknown = unknown +> { /** * Input data that will be passed to the worker. */ @@ -46,5 +52,5 @@ export interface MessageValue { * * _Only for internal use_ */ - readonly parent?: MessagePort + readonly parent?: MainWorker } diff --git a/src/worker/abstract-worker.ts b/src/worker/abstract-worker.ts index 779088c5..34ab964c 100644 --- a/src/worker/abstract-worker.ts +++ b/src/worker/abstract-worker.ts @@ -1,4 +1,6 @@ import { AsyncResource } from 'async_hooks' +import type { Worker } from 'cluster' +import type { MessagePort } from 'worker_threads' import type { MessageValue } from '../utility-types' import type { WorkerOptions } from './worker-options' @@ -10,7 +12,7 @@ import type { WorkerOptions } from './worker-options' * @template Response Type of response the worker sends back to the main worker. */ export abstract class AbstractWorker< - MainWorker, + MainWorker extends Worker | MessagePort, Data = unknown, Response = unknown > extends AsyncResource { @@ -37,12 +39,14 @@ export abstract class AbstractWorker< * @param type The type of async event. * @param isMain Whether this is the main worker or not. * @param fn Function processed by the worker when the pool's `execution` function is invoked. + * @param mainWorker Reference to main worker. * @param opts Options for the worker. */ public constructor ( type: string, isMain: boolean, fn: (data: Data) => Response, + protected mainWorker?: MainWorker | null, public readonly opts: WorkerOptions = {} ) { super(type) @@ -51,7 +55,7 @@ export abstract class AbstractWorker< this.async = !!this.opts.async this.lastTask = Date.now() if (!fn) throw new Error('fn parameter is mandatory') - // keep the worker active + // Keep the worker active if (!isMain) { this.interval = setInterval( this.checkAlive.bind(this), @@ -59,12 +63,38 @@ export abstract class AbstractWorker< ) this.checkAlive.bind(this)() } + + this.mainWorker?.on('message', (value: MessageValue) => { + if (value?.data && value.id) { + // Here you will receive messages + if (this.async) { + this.runInAsyncScope(this.runAsync.bind(this), this, fn, value) + } else { + this.runInAsyncScope(this.run.bind(this), this, fn, value) + } + } else if (value.parent) { + // Save a reference of the main worker to communicate with it + // This will be received once + this.mainWorker = value.parent + } else if (value.kill) { + // Here is time to kill this worker, just clearing the interval + if (this.interval) clearInterval(this.interval) + this.emitDestroy() + } + }) } /** * Returns the main worker. + * + * @returns Reference to the main worker. */ - protected abstract getMainWorker (): MainWorker + protected getMainWorker (): MainWorker { + if (!this.mainWorker) { + throw new Error('Main worker was not set') + } + return this.mainWorker + } /** * Send a message to the main worker. diff --git a/src/worker/cluster-worker.ts b/src/worker/cluster-worker.ts index 9402ec10..4b1efaca 100644 --- a/src/worker/cluster-worker.ts +++ b/src/worker/cluster-worker.ts @@ -30,27 +30,7 @@ export class ClusterWorker< * @param opts Options for the worker. */ public constructor (fn: (data: Data) => Response, opts: WorkerOptions = {}) { - super('worker-cluster-pool:pioardi', isMaster, fn, opts) - - worker.on('message', (value: MessageValue) => { - if (value?.data && value.id) { - // here you will receive messages - // console.log('This is the main worker ' + isMaster) - if (this.async) { - this.runInAsyncScope(this.runAsync.bind(this), this, fn, value) - } else { - this.runInAsyncScope(this.run.bind(this), this, fn, value) - } - } else if (value.kill) { - // here is time to kill this worker, just clearing the interval - if (this.interval) clearInterval(this.interval) - this.emitDestroy() - } - }) - } - - protected getMainWorker (): Worker { - return worker + super('worker-cluster-pool:pioardi', isMaster, fn, worker, opts) } protected sendToMainWorker (message: MessageValue): void { diff --git a/src/worker/thread-worker.ts b/src/worker/thread-worker.ts index b6fbdae2..615e1d8b 100644 --- a/src/worker/thread-worker.ts +++ b/src/worker/thread-worker.ts @@ -1,3 +1,4 @@ +import type { MessagePort } from 'worker_threads' import { isMainThread, parentPort } from 'worker_threads' import type { JSONValue, MessageValue } from '../utility-types' import { AbstractWorker } from './abstract-worker' @@ -22,11 +23,6 @@ export class ThreadWorker< Data extends JSONValue = JSONValue, Response extends JSONValue = JSONValue > extends AbstractWorker { - /** - * Reference to main thread. - */ - protected parent?: MessagePort - /** * Constructs a new poolifier thread worker. * @@ -34,34 +30,7 @@ export class ThreadWorker< * @param opts Options for the worker. */ public constructor (fn: (data: Data) => Response, opts: WorkerOptions = {}) { - super('worker-thread-pool:pioardi', isMainThread, fn, opts) - - parentPort?.on('message', (value: MessageValue) => { - if (value?.data && value.id) { - // here you will receive messages - // console.log('This is the main worker ' + isMainThread) - if (this.async) { - this.runInAsyncScope(this.runAsync.bind(this), this, fn, value) - } else { - this.runInAsyncScope(this.run.bind(this), this, fn, value) - } - } else if (value.parent) { - // save the port to communicate with the main thread - // this will be received once - this.parent = value.parent - } else if (value.kill) { - // here is time to kill this worker, just clearing the interval - if (this.interval) clearInterval(this.interval) - this.emitDestroy() - } - }) - } - - protected getMainWorker (): MessagePort { - if (!this.parent) { - throw new Error('Parent was not set') - } - return this.parent + super('worker-thread-pool:pioardi', isMainThread, fn, parentPort, opts) } protected sendToMainWorker (message: MessageValue): void { -- 2.34.1