Decrypt keystores in a thread pool

packages/cli/src/cmds/validator/keymanager/decryptKeystoreDefinitions/index.ts

@@ -0,0 +1,94 @@
+import {spawn, ModuleThread, Pool, QueuedTask, Worker} from "@chainsafe/threads";
+import {Signer, SignerLocal, SignerType} from "@lodestar/validator";
+import bls from "@chainsafe/bls";
+import {LocalKeystoreDefinition} from "../interface.js";
+import {clearKeystoreCache, loadKeystoreCache, writeKeystoreCache} from "../keystoreCache.js";
+import {lockFilepath, unlockFilepath} from "../../../../util/lockfile.js";
+import {defaultPoolSize} from "./poolSize.js";
+import {DecryptKeystoreWorkerAPI, KeystoreDecryptOptions} from "./types.js";
+
+/**
+ * Decrypt keystore definitions using a threadpool
+ */
+export async function decryptKeystoreDefinitions(
+  keystoreDefinitions: LocalKeystoreDefinition[],
+  opts: KeystoreDecryptOptions
+): Promise<Signer[]> {
+  if (opts.cacheFilePath) {
+    try {
+      const signers = await loadKeystoreCache(opts.cacheFilePath, keystoreDefinitions);
+      if (opts?.onDecrypt) {
+        opts?.onDecrypt(signers.length - 1);
+      }
+      return signers;
+    } catch {
+      // Some error loading the cache, ignore and invalidate cache
+      await clearKeystoreCache(opts.cacheFilePath);
+    }
+  }
+
+  const signers: SignerLocal[] = [];
+  const passwords: string[] = [];
+  const tasks: QueuedTask<ModuleThread<DecryptKeystoreWorkerAPI>, Uint8Array>[] = [];
+  const errors: Error[] = [];
+
+  const pool = Pool(() => spawn<DecryptKeystoreWorkerAPI>(new Worker("./worker.js")), defaultPoolSize);
+  for (const [index, definition] of keystoreDefinitions.entries()) {
+    try {
+      lockFilepath(definition.keystorePath);
+    } catch (e) {
+      if (opts.force) {
+        // Ignore error, maybe log?
+      } else {
+        throw e;
+      }
+    }
+
+    const task = pool.queue((thread) => thread.decryptKeystoreDefinition(definition, Boolean(opts.force)));
+    tasks.push(task);
+    task
+      .then((secretKeyBytes: Uint8Array) => {
+        const signer: SignerLocal = {
+          type: SignerType.Local,
+          secretKey: bls.SecretKey.fromBytes(secretKeyBytes),
+        };
+
+        signers.push(signer);
+        passwords.push(definition.password);
+
+        if (opts?.onDecrypt) {
+          opts?.onDecrypt(index);
+        }
+      })
+      .catch((e: Error) => {
+        // In-progress tasks can't be canceled, so there's a chance that multiple errors may be caught
+        // add to the list of errors
+        errors.push(e);
+        // cancel all pending tasks, no need to continue decrypting after we hit one error
+        for (const task of tasks) {
+          task.cancel();
+        }
+      });
+  }
+
+  try {
+    // only resolves if there are no errored tasks
+    await pool.completed();
+  } catch (e) {
+    // If an error occurs, the program isn't going to be running,
+    // so we should unlock all lockfiles we created
+    for (const {keystorePath} of keystoreDefinitions) {
+      unlockFilepath(keystorePath);
+    }
+
+    throw new AggregateError(errors);
+  } finally {
+    await pool.terminate();
+  }
+
+  if (opts.cacheFilePath) {
+    await writeKeystoreCache(opts.cacheFilePath, signers, passwords);
+  }
+
+  return signers;
+}

packages/cli/src/cmds/validator/keymanager/decryptKeystoreDefinitions/poolSize.ts

@@ -0,0 +1,16 @@
+let defaultPoolSize: number;
+
+try {
+  if (typeof navigator !== "undefined") {
+    defaultPoolSize = navigator.hardwareConcurrency ?? 4;
+  } else {
+    defaultPoolSize = (await import("node:os")).cpus().length;
+  }
+} catch (e) {
+  defaultPoolSize = 8;
+}
+
+/**
+ * Cross-platform aprox number of logical cores
+ */
+export {defaultPoolSize};

packages/cli/src/cmds/validator/keymanager/decryptKeystoreDefinitions/types.ts

@@ -0,0 +1,12 @@
+import {LocalKeystoreDefinition} from "../interface.js";
+
+export type DecryptKeystoreWorkerAPI = {
+  decryptKeystoreDefinition({keystorePath, password}: LocalKeystoreDefinition, force: boolean): Promise<Uint8Array>;
+};
+
+export type KeystoreDecryptOptions = {
+  force?: boolean;
+  onDecrypt?: (index: number) => void;
+  // Try to use the cache file if it exists
+  cacheFilePath?: string;
+};

packages/cli/src/cmds/validator/keymanager/decryptKeystoreDefinitions/worker.ts

@@ -0,0 +1,22 @@
+import fs from "node:fs";
+import {expose} from "@chainsafe/threads/worker";
+import {Keystore} from "@chainsafe/bls-keystore";
+import {LocalKeystoreDefinition} from "../interface.js";
+import {DecryptKeystoreWorkerAPI} from "./types.js";
+
+/**
+ * Decrypt a single keystore definition, returning the secret key as a Uint8Array
+ *
+ * NOTE: This is memory-intensive process, since decrypting the keystore involves running a key derivation function (either pbkdf2 or scrypt)
+ */
+export async function decryptKeystoreDefinition({
+  keystorePath,
+  password,
+}: LocalKeystoreDefinition): Promise<Uint8Array> {
+  const keystore = Keystore.parse(fs.readFileSync(keystorePath, "utf8"));
+
+  // Memory-hogging function
+  return keystore.decrypt(password);
+}
+
+expose({decryptKeystoreDefinition} as DecryptKeystoreWorkerAPI);

packages/cli/src/cmds/validator/keymanager/persistedKeys.ts

@@ -1,8 +1,7 @@
 import fs from "node:fs";
 import path from "node:path";
-import bls from "@chainsafe/bls";
 import {Keystore} from "@chainsafe/bls-keystore";
-import {Signer, SignerType, ProposerConfig, SignerLocal} from "@lodestar/validator";
+import {ProposerConfig} from "@lodestar/validator";
 import {DeletionStatus, ImportStatus, PubkeyHex, SignerDefinition} from "@lodestar/api/keymanager";
 import {
   getPubkeyHexFromKeystore,
@@ -13,8 +12,7 @@ import {
   readProposerConfigDir,
 } from "../../../util/index.js";
 import {lockFilepath} from "../../../util/lockfile.js";
-import {IPersistedKeysBackend} from "./interface.js";
-import {clearKeystoreCache, loadKeystoreCache, writeKeystoreCache} from "./keystoreCache.js";
+import {IPersistedKeysBackend, LocalKeystoreDefinition} from "./interface.js";
 
 export {ImportStatus, DeletionStatus};
 
@@ -25,11 +23,6 @@ type PathArgs = {
   proposerDir: string;
 };
 
-export type LocalKeystoreDefinition = {
-  keystorePath: string;
-  password: string;
-};
-
 /**
  * Class to unify read+write of keystores and remoteKeys from disk.
  * Consumers of this class include:
@@ -245,82 +238,6 @@ export class PersistedKeysBackend implements IPersistedKeysBackend {
   }
 }
 
-type KeystoreDecryptOptions = {
-  force?: boolean;
-  onDecrypt?: (index: number, signer: Signer) => void;
-  // Try to use the cache file if it exists
-  cacheFilePath?: string;
-};
-
-export async function decryptKeystoreDefinitions(
-  keystoreDefinitions: LocalKeystoreDefinition[],
-  opts: KeystoreDecryptOptions
-): Promise<Signer[]> {
-  if (opts.cacheFilePath) {
-    try {
-      const signers = await loadKeystoreCache(opts.cacheFilePath, keystoreDefinitions);
-      if (opts?.onDecrypt) {
-        opts?.onDecrypt(signers.length - 1, signers[signers.length - 1]);
-      }
-      return signers;
-    } catch {
-      // Some error loading the cache, ignore and invalidate cache
-      await clearKeystoreCache(opts.cacheFilePath);
-    }
-  }
-
-  const signers: SignerLocal[] = [];
-  const passwords: string[] = [];
-
-  for (const [index, {keystorePath, password}] of keystoreDefinitions.entries()) {
-    try {
-      lockFilepath(keystorePath);
-    } catch (e) {
-      if (opts.force) {
-        // Ignore error, maybe log?
-      } else {
-        throw e;
-      }
-    }
-
-    const keystore = Keystore.parse(fs.readFileSync(keystorePath, "utf8"));
-
-    // PPS: OOM error issue while decripting validators in parallel
-    // https://github.com/ChainSafe/lodestar/issues/4166
-    //
-    // Below call has been serialized as a hotfix for now as even for 10 vals
-    // it causes 2.5GB memory hog, which doesn't go down even when the promise
-    // resolves and all validators have been decrypted.
-    //
-    // return await Promise.all(validators.map(async (validator) =>
-    //  validator.votingKeypair(this.secretsDir)));
-    //
-    // The new serialized decryption takes full 5 minutes to decrypt 100 validators
-    // on a 100% single core engagement! This needs to be invesigated deeply and
-    // fixed most prefered to the above `Promise.all(...)` flow
-    //
-    const secretKeyBytes = await keystore.decrypt(password);
-
-    const signer: SignerLocal = {
-      type: SignerType.Local,
-      secretKey: bls.SecretKey.fromBytes(secretKeyBytes),
-    };
-
-    signers.push(signer);
-    passwords.push(password);
-
-    if (opts?.onDecrypt) {
-      opts?.onDecrypt(index, signer);
-    }
-  }
-
-  if (opts.cacheFilePath) {
-    await writeKeystoreCache(opts.cacheFilePath, signers, passwords);
-  }
-
-  return signers;
-}
-
 /**
  * Validate SignerDefinition from un-trusted disk file.
  * Performs type validation and re-maps only expected properties.

packages/cli/src/cmds/validator/signers/importExternalKeystores.ts

@@ -1,6 +1,6 @@
 import inquirer from "inquirer";
 import {readPassphraseFile, recursiveLookup} from "../../../util/index.js";
-import {LocalKeystoreDefinition} from "../keymanager/persistedKeys.js";
+import {LocalKeystoreDefinition} from "../keymanager/interface.js";
 
 /**
  * Imports keystores from un-controlled directories provided by the user.

packages/cli/src/cmds/validator/signers/index.ts

@@ -9,7 +9,8 @@ import {defaultNetwork, GlobalArgs} from "../../../options/index.js";
 import {assertValidPubkeysHex, isValidHttpUrl, parseRange, YargsError} from "../../../util/index.js";
 import {getAccountPaths} from "../paths.js";
 import {IValidatorCliArgs} from "../options.js";
-import {decryptKeystoreDefinitions, PersistedKeysBackend} from "../keymanager/persistedKeys.js";
+import {PersistedKeysBackend} from "../keymanager/persistedKeys.js";
+import {decryptKeystoreDefinitions} from "../keymanager/decryptKeystoreDefinitions/index.js";
 import {showProgress} from "../../../util/progress.js";
 import {importKeystoreDefinitionsFromExternalDir, readPassphraseOrPrompt} from "./importExternalKeystores.js";