Enhance BlockAllocator and MasterSearch with improved failure handling and worker management

- Updated BlockAllocator's fail method to return a ReservedBlock or null for better error handling.
- Introduced failMissing method to release reserved slices for workers no longer on the roster.
- Modified MasterSearch's rebuildRoster method to return the number of reclaimed workers.
- Added requeueBlock and idleWorker methods for better management of worker states and job requeuing.
- Enhanced worker disconnection handling in MasterSearch to ensure proper resource management and logging.

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<string>): 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<string, TrackedWorker>();
     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<TrackedWorker, "host" | "device"> | 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<string, TrackedWorker>): 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<string>();
       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?.();
       }