Decrypt keystores in a thread pool

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

@@ -0,0 +1,104 @@
+import {spawn, ModuleThread, Pool, QueuedTask, Worker} from "@chainsafe/threads";
+import {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<SignerLocal[]> {
+  if (opts.cacheFilePath) {
+    try {
+      const signers = await loadKeystoreCache(opts.cacheFilePath, keystoreDefinitions);
+      if (opts?.onDecrypt) {
+        opts?.onDecrypt(signers.length - 1);
+      }
+      opts.logger.debug("Loaded keystores via keystore cache");
+      return signers;
+    } catch {
+      // Some error loading the cache, ignore and invalidate cache
+      await clearKeystoreCache(opts.cacheFilePath);
+    }
+  }
+
+  const signers = new Array(keystoreDefinitions.length) as SignerLocal[];
+  const passwords = new Array(keystoreDefinitions.length) as string[];
+  const tasks: QueuedTask<ModuleThread<DecryptKeystoreWorkerAPI>, Uint8Array>[] = [];
+  const errors: Error[] = [];
+  const pool = Pool(
+    () =>
+      spawn<DecryptKeystoreWorkerAPI>(new Worker("./worker.js"), {
+        // The number below is big enough to almost disable the timeout which helps during tests run on unpredictablely slow hosts
+        timeout: 5 * 60 * 1000,
+      }),
+    defaultPoolSize
+  );
+  for (const [index, definition] of keystoreDefinitions.entries()) {
+    try {
+      lockFilepath(definition.keystorePath);
+    } catch (e) {
+      if (opts.ignoreLockFile) {
+        opts.logger.warn("Keystore forcefully loaded even though lockfile exists", {
+          path: definition.keystorePath,
+        });
+      } else {
+        throw e;
+      }
+    }
+
+    const task = pool.queue((thread) => thread.decryptKeystoreDefinition(definition));
+    tasks.push(task);
+    task
+      .then((secretKeyBytes: Uint8Array) => {
+        const signer: SignerLocal = {
+          type: SignerType.Local,
+          secretKey: bls.SecretKey.fromBytes(secretKeyBytes),
+        };
+
+        signers[index] = signer;
+        passwords[index] = 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(true);
+  } 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);
+    opts.logger.debug("Written keystores to keystore cache");
+  }
+
+  return signers;
+}

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

@@ -0,0 +1,17 @@
+let defaultPoolSize: number;
+
+try {
+  if (typeof navigator !== "undefined") {
+    defaultPoolSize = navigator.hardwareConcurrency ?? 4;
+  } else {
+    // TODO change this line to use os.availableParallelism() once we upgrade to node v20
+    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,14 @@
+import {Logger} from "@lodestar/utils";
+import {LocalKeystoreDefinition} from "../interface.js";
+
+export type DecryptKeystoreWorkerAPI = {
+  decryptKeystoreDefinition({keystorePath, password}: LocalKeystoreDefinition): Promise<Uint8Array>;
+};
+
+export type KeystoreDecryptOptions = {
+  ignoreLockFile?: boolean;
+  onDecrypt?: (index: number) => void;
+  // Try to use the cache file if it exists
+  cacheFilePath?: string;
+  logger: Pick<Logger, "info" | "warn" | "debug">;
+};

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

@@ -0,0 +1,26 @@
+import fs from "node:fs";
+import {expose} from "@chainsafe/threads/worker";
+import {Transfer, TransferDescriptor} from "@chainsafe/threads";
+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 a memory (and cpu) -intensive process, since decrypting the keystore involves running a key derivation function (either pbkdf2 or scrypt)
+ */
+export async function decryptKeystoreDefinition({
+  keystorePath,
+  password,
+}: LocalKeystoreDefinition): Promise<TransferDescriptor<Uint8Array>> {
+  const keystore = Keystore.parse(fs.readFileSync(keystorePath, "utf8"));
+
+  // Memory-hogging function
+  const secret = await keystore.decrypt(password);
+  // Transfer the underlying ArrayBuffer back to the main thread: https://threads.js.org/usage-advanced#transferable-objects
+  // This small performance gain may help in cases where this is run for many keystores
+  return Transfer(secret, [secret.buffer]);
+}
+
+expose({decryptKeystoreDefinition} as unknown 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";
 
@@ -43,7 +44,7 @@ const KEYSTORE_IMPORT_PROGRESS_MS = 10000;
 export async function getSignersFromArgs(
   args: IValidatorCliArgs & GlobalArgs,
   network: string,
-  {logger, signal}: {logger: Pick<Logger, "info">; signal: AbortSignal}
+  {logger, signal}: {logger: Pick<Logger, "info" | "warn" | "debug">; signal: AbortSignal}
 ): Promise<Signer[]> {
   const accountPaths = getAccountPaths(args, network);
 
@@ -95,9 +96,10 @@ export async function getSignersFromArgs(
       },
     });
     return decryptKeystoreDefinitions(keystoreDefinitions, {
-      ...args,
+      ignoreLockFile: args.force,
       onDecrypt: needle,
       cacheFilePath: path.join(accountPaths.cacheDir, "imported_keystores.cache"),
+      logger,
     });
   }
 
@@ -127,9 +129,10 @@ export async function getSignersFromArgs(
     });
 
     const keystoreSigners = await decryptKeystoreDefinitions(keystoreDefinitions, {
-      ...args,
+      ignoreLockFile: args.force,
       onDecrypt: needle,
       cacheFilePath: path.join(accountPaths.cacheDir, "local_keystores.cache"),
+      logger,
     });
 
     // Read local remote keys, imported via keymanager api

packages/cli/src/util/progress.ts

@@ -35,7 +35,7 @@ export function showProgress({
       current,
       total,
       ratePerSec: processTime === 0 ? 0 : ((current - last) / processTime) * 1000,
-      percentage: current && total ? (current / total) * 100 : 100,
+      percentage: total ? (current / total) * 100 : 100,
     });
 
     last = current;

packages/cli/test/unit/util/progress.test.ts

@@ -28,14 +28,16 @@ describe("progress", () => {
       const frequencyMs = 50;
       const total = 8;
       const needle = showProgress({total, signal: new AbortController().signal, frequencyMs, progress});
+      sandbox.clock.tick(frequencyMs);
       needle(1);
       sandbox.clock.tick(frequencyMs);
       needle(3);
       sandbox.clock.tick(frequencyMs);
 
-      expect(progress).to.be.calledTwice;
-      expect(progress.firstCall.args[0]).to.eql({total, current: 2, ratePerSec: 40, percentage: 25});
-      expect(progress.secondCall.args[0]).to.eql({total, current: 4, ratePerSec: 40, percentage: 50});
+      expect(progress).to.be.calledThrice;
+      expect(progress.firstCall.args[0]).to.eql({total, current: 0, ratePerSec: 0, percentage: 0});
+      expect(progress.secondCall.args[0]).to.eql({total, current: 2, ratePerSec: 40, percentage: 25});
+      expect(progress.thirdCall.args[0]).to.eql({total, current: 4, ratePerSec: 40, percentage: 50});
     });
 
     it("should call progress with correct values when reach total", () => {