@cofhe/sdk 0.7.1 → 0.8.0

package/web/test/workerInitRace.web.test.ts

@@ -0,0 +1,79 @@
+import { describe, it, expect, beforeAll, afterAll } from 'vitest';
+import { THREADPOOL_TESTS, fetchNetworkKeys } from './threadPoolBenchHelper';
+
+// What happens when the zkProve worker receives two proof requests before it
+// has finished initializing tfhe — e.g. an app doing
+// `Promise.all([encrypt(a), encrypt(b)])` as its first encryption?
+//
+// `initThreadPool` may only run once per wasm instance. It spawns its N rayon
+// workers first and only then builds rayon's global pool, so a second call
+// spawns N more workers, fails at the build step, and leaves those N stranded
+// for the lifetime of the zkProve worker.
+//
+// The rayon workers are created inside the zkProve worker, out of the page's
+// sight, so this test runs the real worker inside `countingZkProveWorker`,
+// which counts them from the inside.
+
+const THREADS = 2;
+
+type WorkerReply = { id: string; type: string; error?: string; count?: number };
+
+describe.skipIf(!THREADPOOL_TESTS)('@cofhe/sdk/web - zkProve worker init race', () => {
+  let worker: Worker;
+  let keys: { fheKey: string; crs: string };
+  let replies: WorkerReply[];
+  let rayonWorkersStarted: number;
+
+  const request = (message: Record<string, unknown>) =>
+    new Promise<WorkerReply>((resolve) => {
+      worker.addEventListener('message', function onReply(event: MessageEvent) {
+        if (event.data?.id !== message.id) return;
+        worker.removeEventListener('message', onReply);
+        resolve(event.data);
+      });
+      worker.postMessage(message);
+    });
+
+  const prove = (id: string, value: number) =>
+    request({
+      id,
+      type: 'zkProve',
+      fheKeyHex: keys.fheKey,
+      crsHex: keys.crs,
+      items: [{ utype: 'uint32', data: String(value) }],
+      metadata: Array.from(new Uint8Array(32).fill(7)),
+      tfheThreads: THREADS,
+    });
+
+  beforeAll(async () => {
+    keys = await fetchNetworkKeys();
+
+    worker = new Worker(new URL('./countingZkProveWorker.ts', import.meta.url), { type: 'module' });
+    await new Promise<void>((resolve, reject) => {
+      worker.addEventListener('message', function onReady(event: MessageEvent) {
+        if (event.data?.type !== 'ready') return;
+        worker.removeEventListener('message', onReady);
+        resolve();
+      });
+      worker.addEventListener('error', (event) => reject(new Error(`worker failed to load: ${event.message}`)));
+    });
+
+    // Both requests are posted before the worker can finish its first init.
+    replies = await Promise.all([prove('race-a', 1), prove('race-b', 2)]);
+
+    rayonWorkersStarted = (await request({ id: 'count', type: 'countRayonWorkers' })).count ?? -1;
+  }, 240000);
+
+  afterAll(() => worker?.terminate());
+
+  it('generates both proofs', () => {
+    expect(globalThis.crossOriginIsolated).toBe(true);
+    expect(replies.map((reply) => reply.type)).toEqual(['success', 'success']);
+  });
+
+  // The worker memoises its in-flight init, so the second request waits for the
+  // first one's pool instead of starting another.
+  it('starts exactly one rayon pool', () => {
+    expect(rayonWorkersStarted).toBe(THREADS);
+  });
+});

package/web/tfheThreadPool.ts

@@ -0,0 +1,140 @@
+/**
+ * Rayon thread pool setup for the browser build of tfhe.
+ *
+ * `tfhe` is built with wasm-bindgen-rayon, so the ZK proof (the
+ * `build_with_proof_packed` call, by far the most expensive step of
+ * `encryptInputs`) can run across several threads. Those threads are real Web
+ * Workers that share the wasm linear memory, and `initThreadPool` hands them
+ * that memory over `postMessage`.
+ *
+ * Sharing memory across threads is only allowed in a cross-origin-isolated
+ * context, so the host page must be served with:
+ *
+ *   Cross-Origin-Opener-Policy:   same-origin
+ *   Cross-Origin-Embedder-Policy: require-corp   (or `credentialless`)
+ *
+ * Plenty of apps can't or don't set those headers, and calling
+ * `initThreadPool` without them throws ("SharedArrayBuffer transfer requires
+ * self.crossOriginIsolated"). So this is strictly an optimisation: we probe for
+ * isolation, and when it isn't available we leave tfhe single-threaded rather
+ * than failing the encryption. Single-threaded tfhe works in every context.
+ */
+
+import type { TfheThreadsSetting } from '@/core';
+
+/**
+ * Cap for the auto-derived thread count.
+ *
+ * Proving keeps getting faster all the way up to core count, but with sharply
+ * diminishing returns — measured on a 12-core machine (median of 5 proofs of
+ * uint128 + uint64 + uint32):
+ *
+ *   threads |  1   |  2   |  4   |  6   |  8   |  12
+ *   median  | 3199 | 1687 |  964 |  728 |  657 |  624   (ms)
+ *   speedup | 1.0x | 1.9x | 3.3x | 4.4x | 4.9x | 5.1x
+ *
+ * Past 8 threads the curve is flat (8 -> 12 is ~5% for 50% more workers), and
+ * every rayon thread is a real Web Worker competing with the host app's main
+ * thread. 8 captures ~95% of the achievable speedup, so we stop there.
+ */
+export const MAX_AUTO_THREADS = 8;
+
+/**
+ * Resolve a `config.tfheThreads` setting into a concrete thread count.
+ * Returns 1 to mean "don't start a pool".
+ *
+ * For `'auto'`, Zama's guidance is `initThreadPool(navigator.hardwareConcurrency)`;
+ * we follow that but clamp to {@link MAX_AUTO_THREADS}. `hardwareConcurrency`
+ * counts logical cores and some browsers deliberately misreport it (Firefox with
+ * `privacy.resistFingerprinting` returns 2, Safari clamps it), and it can be
+ * absent outside a browser — hence the `?? 1` fallback and the floor of 1.
+ */
+export function resolveThreadCount(setting: TfheThreadsSetting = 'auto'): number {
+  // Explicit opt-out.
+  if (setting === false) return 1;
+
+  // Explicit count: trust the caller over `hardwareConcurrency` (which browsers
+  // clamp for privacy) but keep it a sane integer.
+  if (typeof setting === 'number') {
+    return Number.isFinite(setting) && setting >= 1 ? Math.floor(setting) : 1;
+  }
+
+  const cores = (globalThis as { navigator?: { hardwareConcurrency?: number } }).navigator?.hardwareConcurrency ?? 1;
+
+  if (!Number.isFinite(cores) || cores < 1) return 1;
+
+  return Math.min(Math.floor(cores), MAX_AUTO_THREADS);
+}
+
+export type TfheThreadPoolResult = {
+  /** Whether a multi-threaded rayon pool is now running */
+  enabled: boolean;
+  /** Threads requested from `initThreadPool` (1 when single-threaded) */
+  threads: number;
+  /** Why the pool was not started, when `enabled` is false */
+  reason?: string;
+};
+
+/** Minimal shape of the parts of the tfhe module this file touches. */
+type TfheThreadPoolModule = {
+  initThreadPool?: (numThreads: number) => Promise<unknown>;
+};
+
+/** True when the current context may share wasm memory between threads. */
+export function isCrossOriginIsolated(): boolean {
+  return (globalThis as { crossOriginIsolated?: boolean }).crossOriginIsolated === true;
+}
+
+/**
+ * Start tfhe's rayon thread pool, if the environment allows it.
+ *
+ * Safe to call when unsupported — it reports why it bailed instead of throwing,
+ * leaving tfhe single-threaded. `initThreadPool` must only run once per wasm
+ * instance, so callers are responsible for calling this once (both call sites
+ * in this package memoise their whole tfhe init).
+ */
+export async function initTfheThreadPool(
+  mod: TfheThreadPoolModule,
+  setting: TfheThreadsSetting = 'auto'
+): Promise<TfheThreadPoolResult> {
+  if (setting === false) {
+    return { enabled: false, threads: 1, reason: 'disabled via config.tfheThreads' };
+  }
+
+  if (typeof mod.initThreadPool !== 'function') {
+    return { enabled: false, threads: 1, reason: 'tfhe build has no initThreadPool (single-threaded wasm)' };
+  }
+
+  if (typeof Worker === 'undefined') {
+    return { enabled: false, threads: 1, reason: 'Worker is unavailable' };
+  }
+
+  // The real gate. Without isolation `initThreadPool` throws on postMessage.
+  if (!isCrossOriginIsolated()) {
+    return {
+      enabled: false,
+      threads: 1,
+      reason:
+        'not cross-origin isolated — serve the page with Cross-Origin-Opener-Policy: same-origin and ' +
+        'Cross-Origin-Embedder-Policy: require-corp (or credentialless) to enable multi-threaded proving',
+    };
+  }
+
+  const threads = resolveThreadCount(setting);
+  if (threads <= 1) {
+    return { enabled: false, threads: 1, reason: 'resolved to a single thread' };
+  }
+
+  try {
+    await mod.initThreadPool(threads);
+    return { enabled: true, threads };
+  } catch (error) {
+    // Nested-worker spawning and shared-memory transfer are the flaky parts
+    // across browsers; a failure here must not take the encryption down.
+    return {
+      enabled: false,
+      threads: 1,
+      reason: `initThreadPool failed: ${error instanceof Error ? error.message : String(error)}`,
+    };
+  }
+}

package/web/workerManager.ts

@@ -3,7 +3,7 @@
  * Manages worker lifecycle and request/response handling
  */
 
-import type { ZkProveWorkerRequest, ZkProveWorkerResponse } from '@/core';
+import type { TfheThreadsSetting, ZkProveWorkerRequest, ZkProveWorkerResponse } from '@/core';
 
 // Declare Worker type for environments where it's not available
 declare const Worker: any;
@@ -124,7 +124,8 @@ class ZkProveWorkerManager {
     fheKeyHex: string,
     crsHex: string,
     items: Array<{ utype: string; data: any }>,
-    metadata: Uint8Array
+    metadata: Uint8Array,
+    tfheThreads?: TfheThreadsSetting
   ): Promise<Uint8Array> {
     // Initialize worker if needed
     await this.initializeWorker();
@@ -150,6 +151,7 @@ class ZkProveWorkerManager {
         crsHex,
         items,
         metadata: Array.from(metadata),
+        tfheThreads,
       };
 
       (this.worker as any).postMessage(message);

package/web/zkProve.worker.ts

@@ -8,28 +8,53 @@
 
 import { TFHE_RS_SAFE_SERIALIZATION_SIZE_LIMIT } from '../core/consts';
 import type { ZkProveWorkerRequest, ZkProveWorkerResponse } from '../core/encrypt/zkPackProveVerify.js';
+import { initTfheThreadPool, type TfheThreadPoolResult } from './tfheThreadPool.js';
+import type { TfheThreadsSetting } from '../core/types.js';
 
 // TFHE module (will be initialized on first use)
 let tfheModule: any = null;
-let initialized = false;
+let initPromise: Promise<void> | null = null;
+let threadPool: TfheThreadPoolResult | null = null;
 
 /**
  * Initialize TFHE in worker context
+ *
+ * This is where the heavy `build_with_proof_packed` call runs, so it's also
+ * where tfhe's rayon thread pool matters most — the main-thread pool only helps
+ * the non-worker fallback path. The pool spawns nested Workers that share this
+ * worker's wasm memory, which requires the page to be cross-origin isolated;
+ * `initTfheThreadPool` degrades to single-threaded instead of throwing when it
+ * isn't.
+ *
+ * Memoised on the in-flight promise rather than on completion: further proof
+ * requests can arrive while the first one is still initializing, and
+ * `initThreadPool` may only run once per wasm instance — a second call spawns
+ * its workers, fails to build the pool, and leaves those workers stranded.
  */
-async function initTfhe() {
-  if (initialized) return;
-
-  try {
-    // Dynamic import of tfhe module
-    tfheModule = await import('tfhe');
-    await tfheModule.default();
-    await tfheModule.init_panic_hook();
-    initialized = true;
-    console.log('[Worker] TFHE initialized');
-  } catch (error) {
-    console.error('[Worker] Failed to initialize TFHE:', error);
-    throw error;
-  }
+function initTfhe(tfheThreads: TfheThreadsSetting = 'auto'): Promise<void> {
+  initPromise ??= (async () => {
+    try {
+      // Dynamic import of tfhe module
+      tfheModule = await import('tfhe');
+      await tfheModule.default();
+      await tfheModule.init_panic_hook();
+
+      threadPool = await initTfheThreadPool(tfheModule, tfheThreads);
+
+      console.log(
+        threadPool.enabled
+          ? `[Worker] TFHE initialized (rayon thread pool: ${threadPool.threads} threads)`
+          : `[Worker] TFHE initialized (single-threaded: ${threadPool.reason})`
+      );
+    } catch (error) {
+      console.error('[Worker] Failed to initialize TFHE:', error);
+      // Let the next request retry from scratch.
+      initPromise = null;
+      throw error;
+    }
+  })();
+
+  return initPromise;
 }
 
 /**
@@ -51,7 +76,7 @@ if (typeof self !== 'undefined') {
    * Main message handler
    */
   self.onmessage = async (event: MessageEvent) => {
-    const { id, type, fheKeyHex, crsHex, items, metadata } = event.data as ZkProveWorkerRequest;
+    const { id, type, fheKeyHex, crsHex, items, metadata, tfheThreads } = event.data as ZkProveWorkerRequest;
 
     if (type !== 'zkProve') {
       self.postMessage({
@@ -64,7 +89,7 @@ if (typeof self !== 'undefined') {
 
     try {
       // Initialize TFHE if needed
-      await initTfhe();
+      await initTfhe(tfheThreads);
 
       if (!tfheModule) {
         throw new Error('TFHE module not initialized');