diff --git a/src/search/orchestrator.ts b/src/search/orchestrator.ts index 45367e7..4c53282 100644 --- a/src/search/orchestrator.ts +++ b/src/search/orchestrator.ts @@ -132,13 +132,29 @@ class BlockAllocator { this.completed += block.count; } - fail(jobId: string): void { + fail(jobId: string): ReservedBlock | null { const block = this.inflight.get(jobId); if (!block) { - return; + return null; } this.inflight.delete(jobId); this.retries.push({ start: block.start, count: block.count }); + return block; + } + + /** Release reserved slices whose worker is no longer on the roster. */ + failMissing(liveWorkerIds: Iterable): ReservedBlock[] { + const live = liveWorkerIds instanceof Set ? liveWorkerIds : new Set(liveWorkerIds); + const released: ReservedBlock[] = []; + for (const jobId of [...this.inflight.keys()]) { + const block = this.inflight.get(jobId); + if (block && !live.has(block.workerId)) { + this.inflight.delete(jobId); + this.retries.push({ start: block.start, count: block.count }); + released.push(block); + } + } + return released; } has(jobId: string): boolean { @@ -323,11 +339,12 @@ export class MasterSearch { return [...workers].sort((a, b) => this.priorityRate(b) - this.priorityRate(a)); } - private rebuildRoster(): void { + private rebuildRoster(): number { + const previous = this.roster; const next = new Map(); for (const device of this.params.devices) { const id = localWorkerId(device.key); - const prev = this.roster.get(id); + const prev = previous.get(id); next.set(id, { id, host: "local", @@ -350,7 +367,7 @@ export class MasterSearch { if (this.server) { for (const remote of this.server.listWorkers()) { const id = remoteWorkerId(remote.slaveId, remote.device.key); - const prev = this.roster.get(id); + const prev = previous.get(id); next.set(id, { id, host: remote.hostname, @@ -373,6 +390,68 @@ export class MasterSearch { } } this.roster = next; + return this.reclaimMissingWorkers(previous); + } + + private workerLabel(worker: Pick | undefined, fallback: string): string { + return worker ? `${worker.host} / ${worker.device}` : fallback; + } + + private requeueBlock(block: ReservedBlock | null, reason: string, worker?: TrackedWorker): void { + if (!block) { + return; + } + this.events.onLog?.( + `Requeued ${formatNumber(block.count)} keys from ${this.workerLabel(worker, block.workerId)} (${reason})`, + ); + this.syncStageCompleted(); + } + + private idleWorker(worker: TrackedWorker, keys: bigint): void { + endWork(worker, keys); + worker.jobId = null; + worker.busy = false; + worker.rate = 0; + worker.done = 0n; + worker.count = 0n; + } + + private reclaimMissingWorkers(previous: Map): number { + for (const worker of previous.values()) { + if (this.roster.has(worker.id)) { + continue; + } + if (this.calibrating && worker.jobId) { + this.finishRemoteCalibration(worker.id, worker.jobId, 0); + } + } + if (this.calibrating || !this.stageAlloc) { + return 0; + } + let reclaimed = 0; + const released = this.stageAlloc.failMissing(this.roster.keys()); + for (const block of released) { + this.requeueBlock(block, "worker disconnected", previous.get(block.workerId)); + reclaimed += 1; + } + if (!this.server) { + return reclaimed; + } + for (const worker of this.roster.values()) { + if (!worker.slaveId || !worker.jobId || !worker.busy || !this.stageAlloc.has(worker.jobId)) { + continue; + } + const remote = this.server.listWorkers().find( + (row) => row.slaveId === worker.slaveId && row.device.key === worker.deviceKey, + ); + if (remote && !remote.idle && remote.jobId === worker.jobId) { + continue; + } + this.requeueBlock(this.stageAlloc.fail(worker.jobId), "worker disconnected", worker); + this.idleWorker(worker, worker.done); + reclaimed += 1; + } + return reclaimed; } private bindServer(data: Buffer, firstCt: Buffer): void { @@ -430,16 +509,28 @@ export class MasterSearch { this.emitProgress(); return; } + const slaveGone = !this.server!.list().some((s) => s.id === slave.id); + const reason = slaveGone ? "worker disconnected" : "block failed"; + const jobIds = new Set(); if (jobId) { - this.stageAlloc?.fail(jobId); - this.syncStageCompleted(); - for (const tracked of this.roster.values()) { - if (tracked.jobId === jobId) { - endWork(tracked, tracked.done); - tracked.jobId = null; - tracked.busy = false; - tracked.rate = 0; - } + jobIds.add(jobId); + } + for (const tracked of this.roster.values()) { + if (tracked.slaveId !== slave.id) { + continue; + } + if (jobId && tracked.jobId !== jobId) { + continue; + } + if (tracked.jobId) { + jobIds.add(tracked.jobId); + } + } + for (const droppedId of jobIds) { + const tracked = [...this.roster.values()].find((w) => w.jobId === droppedId); + this.requeueBlock(this.stageAlloc?.fail(droppedId) ?? null, reason, tracked); + if (tracked) { + this.idleWorker(tracked, tracked.done); } } this.rebuildRoster(); @@ -449,9 +540,9 @@ export class MasterSearch { this.server.events.onSlaves = (_slaves) => { prev.onSlaves?.(_slaves); const before = [...this.roster.keys()].sort().join("|"); - this.rebuildRoster(); + const reclaimed = this.rebuildRoster(); const after = [...this.roster.keys()].sort().join("|"); - if (before !== after) { + if (before !== after || reclaimed > 0) { this.emitProgress(); this.dispatch?.(); }