|
| 1 | +/** |
| 2 | + * Encapsulates embed-worker lifecycle management. |
| 3 | + * |
| 4 | + * @param {() => import("node:events").EventEmitter} workerFactory Creates a worker-like object. |
| 5 | + * @param {object} [opts] |
| 6 | + * @param {number} [opts.embedTimeout=60000] Per-embed timeout in ms. |
| 7 | + * @param {number} [opts.restartDelay=2000] Delay before restarting a crashed worker. |
| 8 | + * @param {number} [opts.workerReadyTimeout=30000] Max time to wait for a restarting worker. |
| 9 | + */ |
| 10 | +export function createEmbedPool(workerFactory, opts = {}) { |
| 11 | + const EMBED_TIMEOUT_MS = opts.embedTimeout ?? 60_000; |
| 12 | + const RESTART_DELAY_MS = opts.restartDelay ?? 2000; |
| 13 | + const WORKER_READY_TIMEOUT_MS = opts.workerReadyTimeout ?? 30_000; |
| 14 | + |
| 15 | + let worker = null; |
| 16 | + let workerAlive = false; |
| 17 | + let shuttingDown = false; |
| 18 | + let embedIdCounter = 0; |
| 19 | + const pendingEmbeds = new Map(); |
| 20 | + |
| 21 | + let workerReadyResolve = null; |
| 22 | + let workerReadyPromise = null; |
| 23 | + |
| 24 | + function rejectAllPending(reason) { |
| 25 | + for (const [, { reject, timer }] of pendingEmbeds) { |
| 26 | + clearTimeout(timer); |
| 27 | + reject(new Error(reason)); |
| 28 | + } |
| 29 | + pendingEmbeds.clear(); |
| 30 | + } |
| 31 | + |
| 32 | + function scheduleRestart(code) { |
| 33 | + if (shuttingDown) return; |
| 34 | + process.stderr.write(`[vector-memory] Worker exited (code ${code}) — restarting in ${RESTART_DELAY_MS}ms\n`); |
| 35 | + workerReadyPromise = new Promise(resolve => { workerReadyResolve = resolve; }); |
| 36 | + setTimeout(() => { |
| 37 | + initWorker(); |
| 38 | + if (workerReadyResolve) { |
| 39 | + workerReadyResolve(); |
| 40 | + workerReadyResolve = null; |
| 41 | + } |
| 42 | + }, RESTART_DELAY_MS); |
| 43 | + } |
| 44 | + |
| 45 | + function initWorker() { |
| 46 | + worker = workerFactory(); |
| 47 | + workerAlive = true; |
| 48 | + workerReadyPromise = null; |
| 49 | + |
| 50 | + worker.on("message", (msg) => { |
| 51 | + if (msg.type === "ready") return; |
| 52 | + if (msg.type === "error") { |
| 53 | + process.stderr.write(`[vector-memory] Embedding model error: ${msg.message}\n`); |
| 54 | + return; |
| 55 | + } |
| 56 | + const pending = pendingEmbeds.get(msg.id); |
| 57 | + if (pending) { |
| 58 | + clearTimeout(pending.timer); |
| 59 | + pendingEmbeds.delete(msg.id); |
| 60 | + pending.resolve(msg.embedding); |
| 61 | + } |
| 62 | + }); |
| 63 | + |
| 64 | + worker.on("error", (err) => { |
| 65 | + process.stderr.write(`[vector-memory] Worker crashed: ${err.message}\n`); |
| 66 | + workerAlive = false; |
| 67 | + rejectAllPending("Worker crashed: " + err.message); |
| 68 | + }); |
| 69 | + |
| 70 | + worker.on("exit", (code) => { |
| 71 | + workerAlive = false; |
| 72 | + if (code !== 0) { |
| 73 | + rejectAllPending("Worker exited with code " + code); |
| 74 | + } |
| 75 | + scheduleRestart(code); |
| 76 | + }); |
| 77 | + } |
| 78 | + |
| 79 | + async function embed(text) { |
| 80 | + // If the worker is down but a restart is pending, wait for it |
| 81 | + if (!workerAlive && workerReadyPromise) { |
| 82 | + let timeoutId; |
| 83 | + const timeout = new Promise((_, reject) => { |
| 84 | + timeoutId = setTimeout( |
| 85 | + () => reject(new Error("Embed worker restart timed out")), |
| 86 | + WORKER_READY_TIMEOUT_MS |
| 87 | + ); |
| 88 | + }); |
| 89 | + try { |
| 90 | + await Promise.race([workerReadyPromise, timeout]); |
| 91 | + } finally { |
| 92 | + clearTimeout(timeoutId); |
| 93 | + } |
| 94 | + } |
| 95 | + |
| 96 | + if (!workerAlive) { |
| 97 | + throw new Error("Embed worker is not running"); |
| 98 | + } |
| 99 | + |
| 100 | + return new Promise((resolve, reject) => { |
| 101 | + const id = embedIdCounter++; |
| 102 | + const timer = setTimeout(() => { |
| 103 | + pendingEmbeds.delete(id); |
| 104 | + reject(new Error("Embedding timed out after " + EMBED_TIMEOUT_MS + "ms")); |
| 105 | + }, EMBED_TIMEOUT_MS); |
| 106 | + pendingEmbeds.set(id, { resolve, reject, timer }); |
| 107 | + try { |
| 108 | + worker.postMessage({ id, text }); |
| 109 | + } catch (err) { |
| 110 | + clearTimeout(timer); |
| 111 | + pendingEmbeds.delete(id); |
| 112 | + reject(err); |
| 113 | + } |
| 114 | + }); |
| 115 | + } |
| 116 | + |
| 117 | + function shutdown() { |
| 118 | + shuttingDown = true; |
| 119 | + rejectAllPending("Pool shutting down"); |
| 120 | + if (worker) { |
| 121 | + worker.terminate(); |
| 122 | + worker = null; |
| 123 | + workerAlive = false; |
| 124 | + } |
| 125 | + if (workerReadyResolve) { |
| 126 | + workerReadyResolve(); |
| 127 | + workerReadyResolve = null; |
| 128 | + } |
| 129 | + } |
| 130 | + |
| 131 | + return { embed, initWorker, shutdown, isAlive: () => workerAlive }; |
| 132 | +} |
0 commit comments