|
| 1 | +import * as accounts from "./accounts"; |
| 2 | +import * as downloadQueue from "./download-queue"; |
| 3 | +import { getJob, listJobs, updateJob, type DownloadErrorKind, type DownloadJob, type DownloadShard } from "./download-jobs-store"; |
| 4 | +import { logger } from "./logger"; |
| 5 | +import { getDownloadManager, type JobEvent, type JsDownloadJob } from "./native-downloader"; |
| 6 | + |
| 7 | +const MAX_AUTO_RETRIES = 3; |
| 8 | + |
| 9 | +// How often raw JobEvents get coalesced into a `downloads:update` broadcast |
| 10 | +// — matches the ~100ms cadence the Rust side's own progress ticker already |
| 11 | +// uses, so this never becomes the bottleneck. A job's terminal state always |
| 12 | +// broadcasts immediately regardless (see `driveJob`'s `finally`), so the UI |
| 13 | +// never has to wait out this throttle to see a job actually finish. |
| 14 | +const BROADCAST_THROTTLE_MS = 100; |
| 15 | +let lastBroadcastAt = 0; |
| 16 | + |
| 17 | +function throttledBroadcast(): void { |
| 18 | + const now = Date.now(); |
| 19 | + if (now - lastBroadcastAt < BROADCAST_THROTTLE_MS) return; |
| 20 | + lastBroadcastAt = now; |
| 21 | + downloadQueue.broadcast(); |
| 22 | +} |
| 23 | + |
| 24 | +// Job ids currently being driven — guards against `wake()` picking up the |
| 25 | +// same "queued" job twice; the real concurrency gate is the Rust manager's |
| 26 | +// own semaphore, this is purely a JS-side duplicate-start guard. |
| 27 | +const active = new Set<string>(); |
| 28 | +const retryTimers = new Map<string, NodeJS.Timeout>(); |
| 29 | + |
| 30 | +export function clearRetryTimer(jobId: string): void { |
| 31 | + const timer = retryTimers.get(jobId); |
| 32 | + if (timer) clearTimeout(timer); |
| 33 | + retryTimers.delete(jobId); |
| 34 | +} |
| 35 | + |
| 36 | +function toJsDownloadJob(job: DownloadJob): JsDownloadJob { |
| 37 | + return { |
| 38 | + id: job.id, |
| 39 | + modelId: job.modelId, |
| 40 | + token: accounts.getAccountToken("huggingface") ?? undefined, |
| 41 | + destinationDir: job.destinationDir, |
| 42 | + shards: job.shards.map((s) => ({ |
| 43 | + filename: s.filename, |
| 44 | + path: s.path, |
| 45 | + expectedBytes: s.expectedBytes, |
| 46 | + receivedBytes: s.receivedBytes, |
| 47 | + sha256: s.sha256, |
| 48 | + })), |
| 49 | + }; |
| 50 | +} |
| 51 | + |
| 52 | +function withShardPatch(job: DownloadJob, filename: string, patch: Partial<DownloadShard>): DownloadShard[] { |
| 53 | + return job.shards.map((s) => (s.filename === filename ? { ...s, ...patch } : s)); |
| 54 | +} |
| 55 | + |
| 56 | +function handleEvent(event: JobEvent): void { |
| 57 | + const job = getJob(event.jobId); |
| 58 | + if (!job) return; // job was deleted while its download was in flight |
| 59 | + |
| 60 | + switch (event.kind) { |
| 61 | + case "shard_state": |
| 62 | + if (event.shardFilename && event.shardState) { |
| 63 | + updateJob(job.id, { shards: withShardPatch(job, event.shardFilename, { state: event.shardState as DownloadShard["state"] }) }); |
| 64 | + } |
| 65 | + break; |
| 66 | + case "shard_progress": |
| 67 | + if (event.shardFilename && event.receivedBytes !== undefined) { |
| 68 | + updateJob(job.id, { |
| 69 | + shards: withShardPatch(job, event.shardFilename, { receivedBytes: event.receivedBytes }), |
| 70 | + jobReceivedBytes: event.jobReceivedBytes, |
| 71 | + totalBytes: event.totalBytes, |
| 72 | + bytesPerSecond: event.bytesPerSec, |
| 73 | + etaSeconds: event.etaSeconds, |
| 74 | + }); |
| 75 | + } |
| 76 | + break; |
| 77 | + case "job_state": |
| 78 | + if (event.jobState) { |
| 79 | + updateJob(job.id, { state: event.jobState as DownloadJob["state"] }); |
| 80 | + } |
| 81 | + break; |
| 82 | + case "job_error": |
| 83 | + updateJob(job.id, { |
| 84 | + state: "failed", |
| 85 | + error: { |
| 86 | + message: event.errorMessage ?? "Download failed.", |
| 87 | + kind: (event.errorKind ?? "unknown") as DownloadErrorKind, |
| 88 | + retryable: event.retryable ?? false, |
| 89 | + }, |
| 90 | + }); |
| 91 | + break; |
| 92 | + } |
| 93 | +} |
| 94 | + |
| 95 | +// A `"failed"` job whose error is marked retryable gets a small number of |
| 96 | +// automatic requeue attempts with exponential backoff before it's left for |
| 97 | +// a manual retry — distinct from the Rust side's own per-chunk retry |
| 98 | +// (MAX_ATTEMPTS_PER_CHUNK), which handles much shorter-lived transient |
| 99 | +// failures within a single shard transfer. |
| 100 | +function maybeAutoRetry(jobId: string): void { |
| 101 | + const job = getJob(jobId); |
| 102 | + if (!job || job.state !== "failed" || !job.error?.retryable || job.retryCount >= MAX_AUTO_RETRIES) return; |
| 103 | + |
| 104 | + const backoffMs = 1000 * 2 ** job.retryCount; |
| 105 | + if (retryTimers.has(jobId)) return; |
| 106 | + const expectedRetryCount = job.retryCount; |
| 107 | + const timer = setTimeout(() => { |
| 108 | + retryTimers.delete(jobId); |
| 109 | + const current = getJob(jobId); |
| 110 | + if (!current || current.state !== "failed" || !current.error?.retryable || current.retryCount !== expectedRetryCount) return; |
| 111 | + updateJob(jobId, { state: "queued", retryCount: expectedRetryCount + 1 }); |
| 112 | + wake(); |
| 113 | + }, backoffMs); |
| 114 | + timer.unref(); |
| 115 | + retryTimers.set(jobId, timer); |
| 116 | +} |
| 117 | + |
| 118 | +async function driveJob(job: DownloadJob): Promise<void> { |
| 119 | + active.add(job.id); |
| 120 | + try { |
| 121 | + await getDownloadManager().startJob(toJsDownloadJob(job), (err, event) => { |
| 122 | + if (err) { |
| 123 | + logger.error(`Download job ${job.id}: event callback error: ${err.message}`); |
| 124 | + return; |
| 125 | + } |
| 126 | + handleEvent(event); |
| 127 | + throttledBroadcast(); |
| 128 | + }); |
| 129 | + } catch (err) { |
| 130 | + // The rejection carries the same message the "job_error" event |
| 131 | + // already recorded via handleEvent — this catch exists only so an |
| 132 | + // unexpected throw here doesn't become an unhandled rejection. |
| 133 | + logger.error(`Download job ${job.id}: ${(err as Error).message}`); |
| 134 | + } finally { |
| 135 | + active.delete(job.id); |
| 136 | + downloadQueue.broadcast(); |
| 137 | + maybeAutoRetry(job.id); |
| 138 | + } |
| 139 | +} |
| 140 | + |
| 141 | +// Picks up every currently-"queued" Hugging Face job not already being |
| 142 | +// driven. Safe to call repeatedly/redundantly — `active` prevents double |
| 143 | +// starts. Called once at app-ready and again whenever something might have |
| 144 | +// made a "queued" job appear (a fresh `downloads:create`, a manual |
| 145 | +// retry/resume, or `download-queue.ts`'s startup requeue). |
| 146 | +export function wake(): void { |
| 147 | + for (const job of listJobs()) { |
| 148 | + if (job.state === "queued" && job.kind === "huggingface" && !active.has(job.id)) { |
| 149 | + void driveJob(job); |
| 150 | + } |
| 151 | + } |
| 152 | +} |
| 153 | + |
| 154 | +export function start(): void { |
| 155 | + wake(); |
| 156 | +} |
| 157 | + |
| 158 | +export function pause(jobId: string): void { |
| 159 | + clearRetryTimer(jobId); |
| 160 | + getDownloadManager().pauseJob(jobId); |
| 161 | +} |
| 162 | + |
| 163 | +export function cancel(jobId: string): void { |
| 164 | + clearRetryTimer(jobId); |
| 165 | + getDownloadManager().cancelJob(jobId); |
| 166 | +} |
0 commit comments