|
| 1 | +import { browser } from "$app/environment"; |
| 2 | +import MarkdownWorker from "$lib/workers/markdownWorker?worker"; |
| 3 | +import { fallbackBlocks, processBlocks, type BlockToken } from "$lib/utils/marked"; |
| 4 | + |
| 5 | +type Source = { title?: string; link: string }; |
| 6 | +type ResultCallback = (blocks: BlockToken[], requestId: number) => void; |
| 7 | + |
| 8 | +interface Job { |
| 9 | + clientId: number; |
| 10 | + requestId: number; |
| 11 | + content: string; |
| 12 | + sources: Source[]; |
| 13 | + streaming: boolean; |
| 14 | + onResult: ResultCallback; |
| 15 | +} |
| 16 | + |
| 17 | +interface InFlight { |
| 18 | + requestId: number; |
| 19 | + clientId: number; |
| 20 | + content: string; |
| 21 | + worker: Worker; |
| 22 | + onResult: ResultCallback; |
| 23 | + cancelled: boolean; |
| 24 | +} |
| 25 | + |
| 26 | +// The full markdown pipeline (katex + highlight.js + marked + DOMPurify) is loaded into |
| 27 | +// every worker's JS VM, so the pool is deliberately tiny: a couple of workers shared |
| 28 | +// across every MarkdownRenderer instance, instead of one worker per message. The old |
| 29 | +// per-instance design spawned hundreds of worker VMs in a long thread, which exhausted |
| 30 | +// mobile WebKit's ~2 GB per-tab renderer memory cap and got the tab jetsam-killed. |
| 31 | +const POOL_SIZE = browser ? Math.max(1, Math.min(2, (navigator.hardwareConcurrency || 2) - 1)) : 0; |
| 32 | + |
| 33 | +const hasWorkerSupport = browser && typeof Worker !== "undefined"; |
| 34 | + |
| 35 | +// Flips permanently to true if the Worker constructor throws (strict CSP `worker-src`, |
| 36 | +// sandboxed iframe, Safari Lockdown Mode, quota denial) or if workers keep dying on |
| 37 | +// startup. `typeof Worker !== "undefined"` only proves the global exists, not that |
| 38 | +// construction succeeds, so we also need this runtime guard. Once set, every render |
| 39 | +// degrades to the async main-thread path instead of going silent on the fallback. |
| 40 | +let workersBroken = false; |
| 41 | +const canUseWorker = () => hasWorkerSupport && !workersBroken; |
| 42 | + |
| 43 | +// A worker that dies before ever replying triggers a recreate; cap how many times in a |
| 44 | +// row that can happen before we give up on workers entirely (avoids a recreate storm if |
| 45 | +// the worker module itself fails to initialize). Reset on any successful reply. |
| 46 | +const MAX_CONSECUTIVE_DEATHS = 3; |
| 47 | +let consecutiveDeaths = 0; |
| 48 | + |
| 49 | +let nextRequestId = 1; |
| 50 | +let nextClientId = 1; |
| 51 | + |
| 52 | +const idle: Worker[] = []; |
| 53 | +let createdWorkers = 0; |
| 54 | + |
| 55 | +// Latest queued job per client. Older queued jobs for the same client are coalesced away |
| 56 | +// (only the most recent content matters), which keeps a fast-streaming message from |
| 57 | +// flooding the pool with superseded work. |
| 58 | +const pending = new Map<number, Job>(); |
| 59 | +const queue: number[] = []; |
| 60 | +const queued = new Set<number>(); |
| 61 | + |
| 62 | +// Jobs handed to a worker and awaiting its reply, keyed by requestId. |
| 63 | +const inFlight = new Map<number, InFlight>(); |
| 64 | +// Clients currently occupying a worker. Used to serialize per client so a newer request |
| 65 | +// never races ahead of (or runs concurrently with) the one already in flight. |
| 66 | +const inFlightClients = new Set<number>(); |
| 67 | + |
| 68 | +function runOnMainThread(job: { |
| 69 | + content: string; |
| 70 | + sources: Source[]; |
| 71 | + streaming: boolean; |
| 72 | + requestId: number; |
| 73 | + onResult: ResultCallback; |
| 74 | +}) { |
| 75 | + void processBlocks(job.content, job.sources, job.streaming) |
| 76 | + .then((blocks) => job.onResult(blocks, job.requestId)) |
| 77 | + // Symmetry with the worker path ("never go silent"): on a parse error still |
| 78 | + // resolve the render with the lightweight fallback instead of leaking an |
| 79 | + // unhandled rejection and stranding the renderer on escaped plaintext. |
| 80 | + .catch(() => job.onResult(fallbackBlocks(job.content), job.requestId)); |
| 81 | +} |
| 82 | + |
| 83 | +// Workers became unavailable: flush everything still queued to the main-thread path so no |
| 84 | +// render is left stranded. |
| 85 | +function drainToMainThread() { |
| 86 | + for (const job of pending.values()) runOnMainThread(job); |
| 87 | + pending.clear(); |
| 88 | + queue.length = 0; |
| 89 | + queued.clear(); |
| 90 | +} |
| 91 | + |
| 92 | +function ensureWorker() { |
| 93 | + if (idle.length > 0 || createdWorkers >= POOL_SIZE) return; |
| 94 | + let worker: Worker; |
| 95 | + try { |
| 96 | + worker = new MarkdownWorker(); |
| 97 | + } catch { |
| 98 | + // Construction itself failed (CSP / Lockdown Mode / sandbox). Give up on workers |
| 99 | + // for the session and degrade gracefully. |
| 100 | + workersBroken = true; |
| 101 | + drainToMainThread(); |
| 102 | + return; |
| 103 | + } |
| 104 | + worker.onmessage = (event: MessageEvent) => onWorkerMessage(worker, event); |
| 105 | + worker.onerror = () => handleWorkerDeath(worker); |
| 106 | + worker.onmessageerror = () => handleWorkerDeath(worker); |
| 107 | + createdWorkers++; |
| 108 | + idle.push(worker); |
| 109 | +} |
| 110 | + |
| 111 | +// Detach a worker from the pool, freeing its slot so a replacement can be created. Returns |
| 112 | +// the in-flight entries that were attached to it for the caller to resolve or discard. |
| 113 | +function dropWorker(worker: Worker): InFlight[] { |
| 114 | + const i = idle.indexOf(worker); |
| 115 | + if (i !== -1) idle.splice(i, 1); |
| 116 | + if (createdWorkers > 0) createdWorkers--; |
| 117 | + const orphaned: InFlight[] = []; |
| 118 | + for (const entry of inFlight.values()) { |
| 119 | + if (entry.worker === worker) orphaned.push(entry); |
| 120 | + } |
| 121 | + for (const entry of orphaned) { |
| 122 | + inFlight.delete(entry.requestId); |
| 123 | + inFlightClients.delete(entry.clientId); |
| 124 | + } |
| 125 | + return orphaned; |
| 126 | +} |
| 127 | + |
| 128 | +// A worker crashed without replying (top-level error, structured-clone messageerror, or |
| 129 | +// the OS reclaiming it). Don't strand its render or wedge the pool: deliver the fallback, |
| 130 | +// free the slot, and retry queued work on a fresh worker. |
| 131 | +function handleWorkerDeath(worker: Worker) { |
| 132 | + worker.onmessage = null; |
| 133 | + worker.onerror = null; |
| 134 | + worker.onmessageerror = null; |
| 135 | + const orphaned = dropWorker(worker); |
| 136 | + worker.terminate(); |
| 137 | + for (const entry of orphaned) { |
| 138 | + if (!entry.cancelled) entry.onResult(fallbackBlocks(entry.content), entry.requestId); |
| 139 | + } |
| 140 | + consecutiveDeaths++; |
| 141 | + if (consecutiveDeaths >= MAX_CONSECUTIVE_DEATHS) { |
| 142 | + workersBroken = true; |
| 143 | + drainToMainThread(); |
| 144 | + return; |
| 145 | + } |
| 146 | + pump(); |
| 147 | +} |
| 148 | + |
| 149 | +function onWorkerMessage(worker: Worker, event: MessageEvent) { |
| 150 | + const data = event.data as { type?: string; blocks?: BlockToken[]; requestId?: number }; |
| 151 | + if (data?.type !== "processed" || data.requestId === undefined) return; |
| 152 | + |
| 153 | + consecutiveDeaths = 0; |
| 154 | + idle.push(worker); |
| 155 | + const entry = inFlight.get(data.requestId); |
| 156 | + inFlight.delete(data.requestId); |
| 157 | + if (entry) { |
| 158 | + inFlightClients.delete(entry.clientId); |
| 159 | + if (!entry.cancelled) entry.onResult(data.blocks ?? [], data.requestId); |
| 160 | + } |
| 161 | + pump(); |
| 162 | +} |
| 163 | + |
| 164 | +function pump() { |
| 165 | + while (queue.length > 0) { |
| 166 | + if (idle.length === 0) ensureWorker(); |
| 167 | + if (idle.length === 0) return; // every worker busy and pool at capacity (or broken) |
| 168 | + // First queued client that isn't already running (one job at a time per client). |
| 169 | + const idx = queue.findIndex((clientId) => !inFlightClients.has(clientId)); |
| 170 | + if (idx === -1) return; // all waiting clients are mid-flight; wait for a reply |
| 171 | + const clientId = queue.splice(idx, 1)[0]; |
| 172 | + queued.delete(clientId); |
| 173 | + const job = pending.get(clientId); |
| 174 | + pending.delete(clientId); |
| 175 | + if (!job) continue; |
| 176 | + |
| 177 | + const worker = idle.pop() as Worker; |
| 178 | + inFlightClients.add(clientId); |
| 179 | + inFlight.set(job.requestId, { |
| 180 | + requestId: job.requestId, |
| 181 | + clientId, |
| 182 | + content: job.content, |
| 183 | + worker, |
| 184 | + onResult: job.onResult, |
| 185 | + cancelled: false, |
| 186 | + }); |
| 187 | + worker.postMessage({ |
| 188 | + type: "process", |
| 189 | + content: job.content, |
| 190 | + sources: job.sources, |
| 191 | + requestId: job.requestId, |
| 192 | + streaming: job.streaming, |
| 193 | + }); |
| 194 | + } |
| 195 | +} |
| 196 | + |
| 197 | +/** Stable id for one MarkdownRenderer instance, used to coalesce its successive renders. */ |
| 198 | +export function acquireMarkdownClientId(): number { |
| 199 | + return nextClientId++; |
| 200 | +} |
| 201 | + |
| 202 | +/** |
| 203 | + * Queue a markdown render for `clientId` on the shared worker pool. Returns the requestId; |
| 204 | + * the caller should treat only the highest requestId it has issued as current and ignore |
| 205 | + * late/stale callbacks. Falls back to async main-thread processing when workers aren't |
| 206 | + * available (SSR, no Worker support, or workers that failed at runtime). |
| 207 | + */ |
| 208 | +export function renderMarkdownBlocks( |
| 209 | + clientId: number, |
| 210 | + content: string, |
| 211 | + sources: Source[], |
| 212 | + streaming: boolean, |
| 213 | + onResult: ResultCallback |
| 214 | +): number { |
| 215 | + const requestId = nextRequestId++; |
| 216 | + |
| 217 | + if (!canUseWorker()) { |
| 218 | + runOnMainThread({ content, sources, streaming, requestId, onResult }); |
| 219 | + return requestId; |
| 220 | + } |
| 221 | + |
| 222 | + pending.set(clientId, { clientId, requestId, content, sources, streaming, onResult }); |
| 223 | + if (!queued.has(clientId)) { |
| 224 | + queued.add(clientId); |
| 225 | + queue.push(clientId); |
| 226 | + } |
| 227 | + pump(); |
| 228 | + return requestId; |
| 229 | +} |
| 230 | + |
| 231 | +/** Drop any queued/in-flight work for a client that is going away (component destroyed). */ |
| 232 | +export function cancelMarkdownClient(clientId: number): void { |
| 233 | + pending.delete(clientId); |
| 234 | + if (queued.delete(clientId)) { |
| 235 | + const i = queue.indexOf(clientId); |
| 236 | + if (i !== -1) queue.splice(i, 1); |
| 237 | + } |
| 238 | + |
| 239 | + // If a render is already running for this client, abort it: the component is gone, so |
| 240 | + // terminate its worker (truly stopping the obsolete parse, rather than letting it hold |
| 241 | + // a scarce pool slot) and free the slot. A fresh worker is created lazily on the next |
| 242 | + // pump for whatever view is now active. At most one in-flight job per client (renders |
| 243 | + // are serialized per client), so this terminates at most one worker. |
| 244 | + let abortedWorker: Worker | undefined; |
| 245 | + for (const entry of inFlight.values()) { |
| 246 | + if (entry.clientId === clientId) { |
| 247 | + entry.cancelled = true; |
| 248 | + abortedWorker = entry.worker; |
| 249 | + break; |
| 250 | + } |
| 251 | + } |
| 252 | + if (abortedWorker) { |
| 253 | + abortedWorker.onmessage = null; |
| 254 | + abortedWorker.onerror = null; |
| 255 | + abortedWorker.onmessageerror = null; |
| 256 | + dropWorker(abortedWorker); |
| 257 | + abortedWorker.terminate(); |
| 258 | + pump(); |
| 259 | + } |
| 260 | +} |
| 261 | + |
| 262 | +// Debug hook: read `__mdWorkerPool.workers()` in the console to confirm the live worker |
| 263 | +// count stays at most POOL_SIZE regardless of conversation length. Before pooling, that |
| 264 | +// count equaled the number of mounted MarkdownRenderer instances (hundreds in a long |
| 265 | +// thread), which is what exhausted the mobile renderer's memory. |
| 266 | +if (browser) { |
| 267 | + (globalThis as unknown as { __mdWorkerPool?: unknown }).__mdWorkerPool = { |
| 268 | + poolSize: POOL_SIZE, |
| 269 | + workers: () => createdWorkers, |
| 270 | + idle: () => idle.length, |
| 271 | + queued: () => queue.length, |
| 272 | + inFlight: () => inFlight.size, |
| 273 | + broken: () => workersBroken, |
| 274 | + }; |
| 275 | +} |
0 commit comments