@aztec/foundation 0.23.0 → 0.24.0

package/src/worker/browser/web_worker.ts

@@ -0,0 +1,24 @@
+import { DispatchMsg, TransportClient, WorkerConnector, createDispatchProxy } from '../../transport/index.js';
+import { WasmModule } from '../../wasm/index.js';
+import { WasmWorker } from '../wasm_worker.js';
+
+/**
+ * Instantiate a web worker.
+ * @param url - The URL.
+ * @param initialMem - Initial memory pages.
+ * @param maxMem - Maximum memory pages.
+ * @returns The worker.
+ */
+export async function createWebWorker(url: string, initialMem?: number, maxMem?: number): Promise<WasmWorker> {
+  const worker = new Worker(url);
+  const transportConnect = new WorkerConnector(worker);
+  const transportClient = new TransportClient<DispatchMsg>(transportConnect);
+  await transportClient.open();
+  const remoteModule = createDispatchProxy(WasmModule, transportClient) as WasmWorker;
+  remoteModule.destroyWorker = async () => {
+    await transportClient.request({ fn: '__destroyWorker__', args: [] });
+    transportClient.close();
+  };
+  await remoteModule.init(initialMem, maxMem);
+  return remoteModule;
+}

package/src/worker/data_store.ts

@@ -0,0 +1,19 @@
+/**
+ * Simple read/write interface for wasm modules.
+ */
+export interface DataStore {
+  /**
+   * Get a value from our DB.
+   * @param key - The key to look up.
+   * @returns The value.
+   */
+  get(key: string): Promise<Buffer | undefined>;
+
+  /**
+   * Set a value in our DB.
+   * @param key - The key to update.
+   * @param value - The value to set.
+   * @returns Nothing.
+   */
+  set(key: string, value: Buffer): Promise<void>;
+}

package/src/worker/index.ts

@@ -0,0 +1,2 @@
+export { WorkerPool } from './worker_pool.js';
+export { WasmWorker } from './wasm_worker.js';

package/src/worker/node/index.ts

@@ -0,0 +1,2 @@
+export * from './node_worker.js';
+export * from './node_data_store.js';

package/src/worker/node/node_data_store.ts

@@ -0,0 +1,27 @@
+import leveldown from 'leveldown';
+import levelup, { LevelUp } from 'levelup';
+import memdown from 'memdown';
+
+import { DataStore } from '../data_store.js';
+
+/**
+ * Cache for data used by wasm module.
+ */
+export class NodeDataStore implements DataStore {
+  private db: LevelUp;
+
+  // eslint-disable-next-line
+  constructor(path?: string) {
+    // Hack: Cast as any to work around packages "broken" with node16 resolution
+    // See https://github.com/microsoft/TypeScript/issues/49160
+    this.db = levelup(path ? (leveldown as any)(path) : (memdown as any)());
+  }
+
+  async get(key: string): Promise<Buffer | undefined> {
+    return await this.db.get(key).catch(() => {});
+  }
+
+  async set(key: string, value: Buffer): Promise<void> {
+    await this.db.put(key, value);
+  }
+}

package/src/worker/node/node_worker.ts

@@ -0,0 +1,22 @@
+import { Worker } from 'worker_threads';
+
+import { DispatchMsg, NodeConnector, TransportClient, createDispatchProxy } from '../../transport/index.js';
+import { WasmModule } from '../../wasm/wasm_module.js';
+import { WasmWorker } from '../wasm_worker.js';
+
+/**
+ * Creates a node worker.
+ */
+export async function createNodeWorker(filepath: string, initialMem?: number, maxMem?: number): Promise<WasmWorker> {
+  const worker = new Worker(filepath);
+  const transportConnect = new NodeConnector(worker);
+  const transportClient = new TransportClient<DispatchMsg>(transportConnect);
+  await transportClient.open();
+  const remoteModule = createDispatchProxy(WasmModule, transportClient) as WasmWorker;
+  remoteModule.destroyWorker = async () => {
+    await transportClient.request({ fn: '__destroyWorker__', args: [] });
+    transportClient.close();
+  };
+  await remoteModule.init(initialMem, maxMem);
+  return remoteModule;
+}

package/src/worker/node/start_node_module.ts

@@ -0,0 +1,29 @@
+import { parentPort } from 'worker_threads';
+
+import { DispatchMsg, NodeListener, TransportServer } from '../../transport/index.js';
+import { WasmModule } from '../../wasm/wasm_module.js';
+
+if (!parentPort) {
+  throw new Error('InvalidWorker');
+}
+
+/**
+ * Start the transport server corresponding to this module.
+ * @param module - The WasmModule to host.
+ */
+export function startNodeModule(module: WasmModule) {
+  const dispatch = async ({ fn, args }: DispatchMsg) => {
+    if (fn === '__destroyWorker__') {
+      transportServer.stop();
+      return;
+    }
+    if (!(module as any)[fn]) {
+      throw new Error(`dispatch error, function not found: ${fn}`);
+    }
+    return await (module as any)[fn](...args);
+  };
+  const transportListener = new NodeListener();
+  const transportServer = new TransportServer<DispatchMsg>(transportListener, dispatch);
+  module.addLogger((...args: any[]) => transportServer.broadcast({ fn: 'emit', args: ['log', ...args] }));
+  transportServer.start();
+}

package/src/worker/wasm_worker.ts

@@ -0,0 +1,7 @@
+import { Proxify } from '../transport/index.js';
+import { WasmModule } from '../wasm/wasm_module.js';
+
+/**
+ * Represents either a WASM web worker, or node.js worker.
+ */
+export type WasmWorker = Proxify<WasmModule> & { destroyWorker(): void };

package/src/worker/worker_pool.ts

@@ -0,0 +1,73 @@
+import { createDebugLogger } from '../log/index.js';
+import { WasmWorker } from './wasm_worker.js';
+
+const debug = createDebugLogger('bb:worker_pool');
+
+/**
+ * Type of a worker factory.
+ * Used to customize WorkerPool worker construction.
+ */
+export type CreateWorker = (name: string, minMem: number, maxMem: number) => WasmWorker;
+/**
+ * Allocates a pool of WasmWorker's.
+ * Worker 0 is allocated MAX_PAGES memory pages. This is because worker 0 will need to hold the proving key
+ * (i.e. Has state), whereas the others are pure compute (they hold a little crs state).
+ */
+export class WorkerPool {
+  // TODO(AD): Revisit what this means in aztec 3 context
+  // --
+  // Introduction of low mem prover work (polynomial cache) may actually increase mem usage when the backing store isn't
+  // enabled. We were seeing intermittent failings related to memory in production for some users when limiting to
+  // 6660 (416MB). It would be nice to understand why this is (the non determinism and/or the increased mem usage).
+  // For now, increasing mem usage to 512MB. This maybe preferable to backing out the low mem work, but
+  // ironically may break the chance of us using it in mobile.
+  // We *could* enable the low memory backing store, but this needs a little bit of work to actually
+  // read/write from indexeddb, performance testing, and actual further memory load testing.
+  // At this point it's hard to know what our memory savings would be relative to just fully reverting the LMP.
+  // public static MAX_PAGES = 6660;
+  /**
+   * The maximum number of memory pages to be used by the webassembly.
+   */
+  public static MAX_PAGES = 8192;
+  /**
+   * The workers in the pool.
+   */
+  private workers: WasmWorker[] = [];
+
+  /**
+   * Create an instance and initialize the workers.
+   * @param createWorker - Worker factory.
+   * @param poolSize - Pool size.
+   * @returns An initialized WorkerPool.
+   */
+  static async new(createWorker: CreateWorker, poolSize: number) {
+    const pool = new WorkerPool();
+    await pool.init(createWorker, poolSize);
+    return pool;
+  }
+
+  /**
+   * Initialize the workers.
+   * @param createWorker - Worker factory().
+   * @param poolSize - Pool size.
+   * @param maxMem - Max memory pages.
+   */
+  public async init(createWorker: CreateWorker, poolSize: number, maxMem = WorkerPool.MAX_PAGES) {
+    debug(`creating ${poolSize} workers...`);
+    const start = new Date().getTime();
+    this.workers = await Promise.all(
+      Array(poolSize)
+        .fill(0)
+        .map((_, i) => createWorker(`${i}`, i === 0 ? Math.min(WorkerPool.MAX_PAGES, maxMem) : 768, maxMem)),
+    );
+
+    debug(`created workers: ${new Date().getTime() - start}ms`);
+  }
+
+  /**
+   * Tell all workers in the pool to stop processing.
+   */
+  public async destroy() {
+    await Promise.all(this.workers.map(w => w.destroyWorker()));
+  }
+}