feat: add support for async kill handler
[poolifier.git] / src / pools / thread / fixed.ts
1 import {
2 type MessageChannel,
3 type MessagePort,
4 SHARE_ENV,
5 type TransferListItem,
6 Worker,
7 type WorkerOptions,
8 isMainThread
9 } from 'node:worker_threads'
10 import type { MessageValue } from '../../utility-types'
11 import { AbstractPool } from '../abstract-pool'
12 import { type PoolOptions, type PoolType, PoolTypes } from '../pool'
13 import { type WorkerType, WorkerTypes } from '../worker'
14
15 /**
16 * Options for a poolifier thread pool.
17 */
18 export interface ThreadPoolOptions extends PoolOptions<Worker> {
19 /**
20 * Worker options.
21 *
22 * @see https://nodejs.org/api/worker_threads.html#new-workerfilename-options
23 */
24 workerOptions?: WorkerOptions
25 }
26
27 /**
28 * A thread pool with a fixed number of threads.
29 *
30 * @typeParam Data - Type of data sent to the worker. This can only be structured-cloneable data.
31 * @typeParam Response - Type of execution response. This can only be structured-cloneable data.
32 * @author [Alessandro Pio Ardizio](https://github.com/pioardi)
33 * @since 0.0.1
34 */
35 export class FixedThreadPool<
36 Data = unknown,
37 Response = unknown
38 > extends AbstractPool<Worker, Data, Response> {
39 /**
40 * Constructs a new poolifier fixed thread pool.
41 *
42 * @param numberOfThreads - Number of threads for this pool.
43 * @param filePath - Path to an implementation of a `ThreadWorker` file, which can be relative or absolute.
44 * @param opts - Options for this fixed thread pool.
45 */
46 public constructor (
47 numberOfThreads: number,
48 filePath: string,
49 protected readonly opts: ThreadPoolOptions = {}
50 ) {
51 super(numberOfThreads, filePath, opts)
52 }
53
54 /** @inheritDoc */
55 protected isMain (): boolean {
56 return isMainThread
57 }
58
59 /** @inheritDoc */
60 protected async destroyWorkerNode (workerNodeKey: number): Promise<void> {
61 this.flushTasksQueue(workerNodeKey)
62 // FIXME: wait for tasks to be finished
63 const workerNode = this.workerNodes[workerNodeKey]
64 const worker = workerNode.worker
65 const waitWorkerExit = new Promise<void>((resolve) => {
66 worker.on('exit', () => {
67 resolve()
68 })
69 })
70 this.sendToWorker(workerNodeKey, { kill: true, workerId: worker.threadId })
71 workerNode.closeChannel()
72 await worker.terminate()
73 await waitWorkerExit
74 }
75
76 /** @inheritDoc */
77 protected sendToWorker (
78 workerNodeKey: number,
79 message: MessageValue<Data>,
80 transferList?: TransferListItem[]
81 ): void {
82 (
83 this.getWorkerInfo(workerNodeKey).messageChannel as MessageChannel
84 ).port1.postMessage(message, transferList)
85 }
86
87 /** @inheritDoc */
88 protected sendStartupMessageToWorker (workerNodeKey: number): void {
89 const worker = this.workerNodes[workerNodeKey].worker
90 const port2: MessagePort = (
91 this.getWorkerInfo(workerNodeKey).messageChannel as MessageChannel
92 ).port2
93 worker.postMessage(
94 {
95 ready: false,
96 workerId: worker.threadId,
97 port: port2
98 },
99 [port2]
100 )
101 }
102
103 /** @inheritDoc */
104 protected registerWorkerMessageListener<Message extends Data | Response>(
105 workerNodeKey: number,
106 listener: (message: MessageValue<Message>) => void
107 ): void {
108 (
109 this.getWorkerInfo(workerNodeKey).messageChannel as MessageChannel
110 ).port1.on('message', listener)
111 }
112
113 /** @inheritDoc */
114 protected createWorker (): Worker {
115 return new Worker(this.filePath, {
116 env: SHARE_ENV,
117 ...this.opts.workerOptions
118 })
119 }
120
121 /** @inheritDoc */
122 protected get type (): PoolType {
123 return PoolTypes.fixed
124 }
125
126 /** @inheritDoc */
127 protected get worker (): WorkerType {
128 return WorkerTypes.thread
129 }
130
131 /** @inheritDoc */
132 protected get minSize (): number {
133 return this.numberOfWorkers
134 }
135
136 /** @inheritDoc */
137 protected get maxSize (): number {
138 return this.numberOfWorkers
139 }
140
141 /** @inheritDoc */
142 protected get busy (): boolean {
143 return this.internalBusy()
144 }
145 }