Merge pull request #121 from lidofinance/fix/key-validation-threads

fix: remove initialization of thread pool on each validation

packages/key-validation/package.json

@@ -34,7 +34,7 @@
     "@lido-nestjs/contracts": "workspace:*",
     "@lido-nestjs/di": "workspace:*",
     "@lido-nestjs/utils": "workspace:*",
-    "piscina": "3.2.0"
+    "piscina": "4.6.1"
   },
   "peerDependencies": {
     "@nestjs/common": "^8.2.5",

packages/key-validation/src/interfaces/key-validator.interface.ts

@@ -1,5 +1,6 @@
 import { Key } from './common';
 import { createInterface } from '@lido-nestjs/di';
+import { KeyValidatorExecutorInterface } from './key-validator.executor.interface';
 
 export const KeyValidatorInterface = createInterface<KeyValidatorInterface>(
   'KeyValidatorInterface',
@@ -24,4 +25,11 @@ export interface KeyValidatorInterface {
    * the same data will be returned with the result
    */
   validateKeys<T>(keys: (Key & T)[]): Promise<[Key & T, boolean][]>;
+
+  /**
+   * Executor of the validation process.
+   *
+   * Can be single or multithreaded
+   */
+  readonly executor: KeyValidatorExecutorInterface;
 }

packages/key-validation/src/services/key-validator.ts

@@ -9,9 +9,7 @@ import { ImplementsAtRuntime } from '@lido-nestjs/di';
 @Injectable()
 @ImplementsAtRuntime(KeyValidatorInterface)
 export class KeyValidator implements KeyValidatorInterface {
-  public constructor(
-    private readonly executor: KeyValidatorExecutorInterface,
-  ) {}
+  public constructor(public readonly executor: KeyValidatorExecutorInterface) {}
 
   public async validateKey<T>(key: Key & T): Promise<boolean> {
     return this.executor.validateKey(key);

packages/key-validation/src/services/multi-threaded-key-validator-executor.ts

@@ -8,6 +8,12 @@ import assert from 'assert';
 export class MultiThreadedKeyValidatorExecutor
   implements KeyValidatorExecutorInterface
 {
+  threadPool: Piscina<unknown, [index: number, valid: boolean][]> | null = null;
+
+  public constructor() {
+    this.enableGracefulShutdown();
+  }
+
   public async validateKey<T = never>(key: Key & T): Promise<boolean> {
     const serialized = serialize(key);
 
@@ -19,9 +25,12 @@ export class MultiThreadedKeyValidatorExecutor
   public async validateKeys<T = never>(
     keys: (Key & T)[],
   ): Promise<[Key & T, boolean][]> {
-    const threadPool = new Piscina({
-      filename: worker.filename,
-    });
+    const threadPool =
+      this.threadPool ??
+      (this.threadPool = new Piscina({
+        filename: worker.filename,
+        minThreads: 2,
+      }));
 
     type Runner = (
       task: Parameters<typeof worker>[0],
@@ -56,8 +65,25 @@ export class MultiThreadedKeyValidatorExecutor
       }),
     );
 
-    await threadPool.destroy();
-
     return results.flat();
   }
+
+  protected enableGracefulShutdown() {
+    // handling process termination
+    /* istanbul ignore next */
+    process.on('SIGTERM', async () => {
+      await this.destroy();
+    });
+
+    // handling Ctrl+C
+    /* istanbul ignore next */
+    process.on('SIGINT', async () => {
+      await this.destroy();
+    });
+  }
+
+  public async destroy() {
+    await this.threadPool?.destroy();
+    this.threadPool = null;
+  }
 }

packages/key-validation/test/key-validator.spec.ts

@@ -11,6 +11,7 @@ import {
   KeyValidator,
   KeyValidatorInterface,
   KeyValidatorModule,
+  MultiThreadedKeyValidatorExecutor,
 } from '../src';
 import { Test } from '@nestjs/testing';
 import { ModuleMetadata } from '@nestjs/common';
@@ -157,6 +158,34 @@ describe('KeyValidator', () => {
     expect(time).toBeLessThan(0.01);
   });
 
+  test('[multi-thread] should correctly handle destroy after validation', async () => {
+    const keyValidator = await getKeyValidator(true);
+
+    expect(keyValidator.executor).toBeInstanceOf(
+      MultiThreadedKeyValidatorExecutor,
+    );
+    if (keyValidator.executor instanceof MultiThreadedKeyValidatorExecutor) {
+      expect(keyValidator.executor.threadPool).toBeNull();
+      const res = await keyValidator.validateKeys([validKey]);
+      expect(res.length).toBe(1);
+      expect(keyValidator.executor.threadPool).not.toBeNull();
+      await keyValidator.executor.destroy();
+      expect(keyValidator.executor.threadPool).toBeNull();
+    }
+  });
+
+  test('[multi-thread] should correctly handle destroy before validation', async () => {
+    const keyValidator = await getKeyValidator(true);
+
+    expect(keyValidator.executor).toBeInstanceOf(
+      MultiThreadedKeyValidatorExecutor,
+    );
+    if (keyValidator.executor instanceof MultiThreadedKeyValidatorExecutor) {
+      await keyValidator.executor.destroy();
+      expect(keyValidator.executor.threadPool).toBeNull();
+    }
+  });
+
   test('[single-thread] should successfully validate a valid key with custom properties', async () => {
     const keyValidator = await getKeyValidator(false);
 
@@ -235,5 +264,9 @@ describe('KeyValidator', () => {
 
     expect(res.length).toBe(1000);
     expect(time).toBeLessThan(15); // 30 seconds
+
+    if (keyValidator.executor instanceof MultiThreadedKeyValidatorExecutor) {
+      await keyValidator.executor.destroy();
+    }
   });
 });

yarn.lock

@@ -14,13 +14,6 @@ __metadata:
   languageName: node
   linkType: hard
 
-"@assemblyscript/loader@npm:^0.10.1":
-  version: 0.10.1
-  resolution: "@assemblyscript/loader@npm:0.10.1"
-  checksum: fd1f57bdf2c55252a48c2d93fbec3c5a9ef4ca40e581e8709dd8ee437613eb47af74c8cdba011a324077eda9605d6f24983f429c2ce18b0b582ddcc5acf75c26
-  languageName: node
-  linkType: hard
-
 "@babel/code-frame@npm:^7.0.0, @babel/code-frame@npm:^7.12.13, @babel/code-frame@npm:^7.16.7":
   version: 7.16.7
   resolution: "@babel/code-frame@npm:7.16.7"
@@ -2768,7 +2761,7 @@ __metadata:
     "@nestjs/common": ^8.2.5
     "@nestjs/core": ^8.2.5
     "@nestjs/testing": ^8.2.5
-    piscina: 3.2.0
+    piscina: 4.6.1
     reflect-metadata: ^0.1.13
     rxjs: ^7.5.2
   peerDependencies:
@@ -5314,13 +5307,6 @@ __metadata:
   languageName: node
   linkType: hard
 
-"base64-js@npm:^1.2.0":
-  version: 1.5.1
-  resolution: "base64-js@npm:1.5.1"
-  checksum: 669632eb3745404c2f822a18fc3a0122d2f9a7a13f7fb8b5823ee19d1d2ff9ee5b52c53367176ea4ad093c332fd5ab4bd0ebae5a8e27917a4105a4cfc86b1005
-  languageName: node
-  linkType: hard
-
 "base@npm:^0.11.1":
   version: 0.11.2
   resolution: "base@npm:0.11.2"
@@ -7642,13 +7628,6 @@ __metadata:
   languageName: node
   linkType: hard
 
-"eventemitter-asyncresource@npm:^1.0.0":
-  version: 1.0.0
-  resolution: "eventemitter-asyncresource@npm:1.0.0"
-  checksum: 3cfbbc3490bd429a165bff6336289ff810f7df214796f25000d2097a5a0883eae51542a78674916ff99bbd4c66811911b310df1cb4fc96dfc9546ba9dfc89f8f
-  languageName: node
-  linkType: hard
-
 "eventemitter3@npm:^3.1.0":
   version: 3.1.2
   resolution: "eventemitter3@npm:3.1.2"
@@ -8944,24 +8923,6 @@ __metadata:
   languageName: node
   linkType: hard
 
-"hdr-histogram-js@npm:^2.0.1":
-  version: 2.0.3
-  resolution: "hdr-histogram-js@npm:2.0.3"
-  dependencies:
-    "@assemblyscript/loader": ^0.10.1
-    base64-js: ^1.2.0
-    pako: ^1.0.3
-  checksum: 7bb252ba3596bed72b90427ffc6f6fa332a460c4810788faa9b9a743f7ac6f1cb42dccd7ae7555740f0a8c0602884944d00d1ccfb746af4976a816772361a6d6
-  languageName: node
-  linkType: hard
-
-"hdr-histogram-percentiles-obj@npm:^3.0.0":
-  version: 3.0.0
-  resolution: "hdr-histogram-percentiles-obj@npm:3.0.0"
-  checksum: ab238edcb38d9b60d23ca53da0ecd9a6b1c8ee9a49e30a6146bd3f8f70f26244652f28b79974157c00504e7ddf3129e0ddb217baf71d32330e3fae0105bf30ed
-  languageName: node
-  linkType: hard
-
 "hmac-drbg@npm:^1.0.1":
   version: 1.0.1
   resolution: "hmac-drbg@npm:1.0.1"
@@ -13751,13 +13712,6 @@ __metadata:
   languageName: node
   linkType: hard
 
-"pako@npm:^1.0.3":
-  version: 1.0.11
-  resolution: "pako@npm:1.0.11"
-  checksum: 1be2bfa1f807608c7538afa15d6f25baa523c30ec870a3228a89579e474a4d992f4293859524e46d5d87fd30fa17c5edf34dbef0671251d9749820b488660b16
-  languageName: node
-  linkType: hard
-
 "parallel-transform@npm:^1.1.0":
   version: 1.2.0
   resolution: "parallel-transform@npm:1.2.0"
@@ -14047,18 +14001,15 @@ __metadata:
   languageName: node
   linkType: hard
 
-"piscina@npm:3.2.0":
-  version: 3.2.0
-  resolution: "piscina@npm:3.2.0"
+"piscina@npm:4.6.1":
+  version: 4.6.1
+  resolution: "piscina@npm:4.6.1"
   dependencies:
-    eventemitter-asyncresource: ^1.0.0
-    hdr-histogram-js: ^2.0.1
-    hdr-histogram-percentiles-obj: ^3.0.0
     nice-napi: ^1.0.2
   dependenciesMeta:
     nice-napi:
       optional: true
-  checksum: c1980c7d45d85f53265652dd2fc62a2b9e9d2321f5bbb9fc1796edb9c1324bb77c153e823a0d6454c3c35098820efedff584737cc282207480afe478a3b8a166
+  checksum: a0d79d1798965a0fc4aef65045b2deacbc3af7bd6effa8d08b4e1d59dfacc082e26d496701bb87b8ff2dc0798e1868ec612a6ad75ffe837f7c13cfb98d4f4e01
   languageName: node
   linkType: hard