| 1 | import { createWriteStream } from "node:fs"; |
| 2 | import { open, stat } from "node:fs/promises"; |
| 3 | import { totalmem } from "node:os"; |
| 4 | import { Readable } from "node:stream"; |
| 5 | import { pipeline } from "node:stream/promises"; |
| 6 | import type { ReadableStream as WebReadableStream } from "node:stream/web"; |
| 7 | import type { JobSpec } from "@diffinery/shared"; |
| 8 | |
| 9 | /** |
| 10 | * Client for the Actions runtime services (ACTIONS_RESULTS_URL): the Twirp |
| 11 | * cache and artifact services and their signed blob URLs. The built-in cache |
| 12 | * and artifact actions use it, so they share storage (and archive formats) |
| 13 | * with the official actions. |
| 14 | */ |
| 15 | |
| 16 | export class RuntimeError extends Error { |
| 17 | constructor( |
| 18 | readonly status: number, |
| 19 | readonly code: string, |
| 20 | message: string, |
| 21 | ) { |
| 22 | super(message); |
| 23 | } |
| 24 | } |
| 25 | |
| 26 | const sleep = (ms: number) => new Promise((r) => setTimeout(r, ms)); |
| 27 | const retryable = (status: number) => status === 408 || status === 429 || status >= 500; |
| 28 | |
| 29 | /** Uploads at most this much in one request; larger files go up as blocks. */ |
| 30 | const SINGLE_SHOT = 32 * 1024 * 1024; |
| 31 | /** |
| 32 | * Larger files go up in blocks of this size, this many at a time (measured on hosted runners: 4 of |
| 33 | * 32 MiB ran at 23-33 MiB/s, 16 of 16 MiB at 58-62 MiB/s). Blocks in flight are held in memory, so |
| 34 | * the 1 GiB runner sends fewer and smaller ones (32 MiB at most). |
| 35 | */ |
| 36 | const SMALL_MACHINE = totalmem() < 2 * 1024 ** 3; |
| 37 | const BLOCK_SIZE = (SMALL_MACHINE ? 8 : 16) * 1024 * 1024; |
| 38 | const UPLOAD_CONCURRENCY = SMALL_MACHINE ? 4 : 8; |
| 39 | /** Downloads come in ranges of this size, this many at a time (measured on hosted runners: 16 of 8 MiB ran 2-5x faster than one stream). */ |
| 40 | const DOWNLOAD_CHUNK = 8 * 1024 * 1024; |
| 41 | // Fewer on the smallest runner (1 GiB), so ranges waiting their turn hold at most 32 MiB there. |
| 42 | const DOWNLOAD_CONCURRENCY = SMALL_MACHINE ? 4 : 16; |
| 43 | |
| 44 | export class RuntimeClient { |
| 45 | constructor( |
| 46 | readonly token: string, |
| 47 | readonly resultsUrl: string, |
| 48 | private readonly opts: { retries?: number; minBackoffMs?: number } = {}, |
| 49 | ) {} |
| 50 | |
| 51 | static fromSpec(spec: JobSpec): RuntimeClient | null { |
| 52 | return spec.runtime ? new RuntimeClient(spec.runtime.token, spec.runtime.resultsUrl) : null; |
| 53 | } |
| 54 | |
| 55 | async withRetries<T>(what: string, fn: () => Promise<T>): Promise<T> { |
| 56 | const retries = this.opts.retries ?? 5; |
| 57 | let last: unknown; |
| 58 | for (let attempt = 0; attempt <= retries; attempt++) { |
| 59 | if (attempt) await sleep((this.opts.minBackoffMs ?? 1000) * 2 ** (attempt - 1)); |
| 60 | try { |
| 61 | return await fn(); |
| 62 | } catch (e) { |
| 63 | if (e instanceof RuntimeError && !retryable(e.status)) throw e; |
| 64 | last = e; |
| 65 | } |
| 66 | } |
| 67 | throw last instanceof Error ? new Error(`${what}: ${last.message}`) : new Error(`${what} failed`); |
| 68 | } |
| 69 | |
| 70 | /** Calls a Twirp method (JSON). Errors become RuntimeError with the Twirp code. */ |
| 71 | async twirp<T>(service: "CacheService" | "ArtifactService", method: string, body: Record<string, unknown>): Promise<T> { |
| 72 | const url = `${new URL(this.resultsUrl).origin}/twirp/github.actions.results.api.v1.${service}/${method}`; |
| 73 | return this.withRetries(method, async () => { |
| 74 | const res = await fetch(url, { |
| 75 | method: "POST", |
| 76 | headers: { authorization: `Bearer ${this.token}`, "content-type": "application/json", "user-agent": "diffinery-runner" }, |
| 77 | body: JSON.stringify(body), |
| 78 | }); |
| 79 | const text = await res.text(); |
| 80 | if (!res.ok) { |
| 81 | let code = "unknown"; |
| 82 | let msg = text.slice(0, 300); |
| 83 | try { |
| 84 | const j = JSON.parse(text) as { code?: string; msg?: string }; |
| 85 | code = j.code ?? code; |
| 86 | msg = j.msg ?? msg; |
| 87 | } catch { |
| 88 | /* not Twirp JSON */ |
| 89 | } |
| 90 | throw new RuntimeError(res.status, code, `${method} failed (${res.status} ${code}): ${msg}`); |
| 91 | } |
| 92 | return JSON.parse(text) as T; |
| 93 | }); |
| 94 | } |
| 95 | |
| 96 | /** Uploads a file to a signed URL (Put Blob, or Put Block + Put Block List). */ |
| 97 | async upload(url: string, file: string, contentType = "application/octet-stream"): Promise<number> { |
| 98 | const size = (await stat(file)).size; |
| 99 | const fh = await open(file, "r"); |
| 100 | try { |
| 101 | const read = async (offset: number, length: number) => { |
| 102 | const buf = Buffer.alloc(length); |
| 103 | let got = 0; |
| 104 | while (got < length) { |
| 105 | const { bytesRead } = await fh.read(buf, got, length - got, offset + got); |
| 106 | if (bytesRead === 0) break; |
| 107 | got += bytesRead; |
| 108 | } |
| 109 | return buf.subarray(0, got); |
| 110 | }; |
| 111 | const put = (target: string, body: Buffer, headers: Record<string, string>) => |
| 112 | this.withRetries("upload", async () => { |
| 113 | const res = await fetch(target, { method: "PUT", headers: { "x-ms-version": "2025-01-05", ...headers }, body }); |
| 114 | if (res.status !== 201) { |
| 115 | const text = await res.text().catch(() => ""); |
| 116 | throw new RuntimeError(res.status, "upload", `upload failed (${res.status}): ${text.slice(0, 300)}`); |
| 117 | } |
| 118 | }); |
| 119 | if (size <= SINGLE_SHOT) { |
| 120 | await put(url, await read(0, size), { "x-ms-blob-type": "BlockBlob", "x-ms-blob-content-type": contentType, "content-type": "application/octet-stream" }); |
| 121 | return size; |
| 122 | } |
| 123 | const blocks: string[] = []; |
| 124 | for (let off = 0, i = 0; off < size; off += BLOCK_SIZE, i++) blocks.push(Buffer.from(`block-${String(i).padStart(8, "0")}`).toString("base64")); |
| 125 | let next = 0; |
| 126 | const worker = async () => { |
| 127 | for (let i = next++; i < blocks.length; i = next++) { |
| 128 | const body = await read(i * BLOCK_SIZE, Math.min(BLOCK_SIZE, size - i * BLOCK_SIZE)); |
| 129 | await put(`${url}&comp=block&blockid=${encodeURIComponent(blocks[i]!)}`, body, { "content-type": "application/octet-stream" }); |
| 130 | } |
| 131 | }; |
| 132 | await Promise.all(Array.from({ length: Math.min(UPLOAD_CONCURRENCY, blocks.length) }, worker)); |
| 133 | const xml = `<?xml version="1.0" encoding="utf-8"?><BlockList>${blocks.map((b) => `<Latest>${b}</Latest>`).join("")}</BlockList>`; |
| 134 | await put(`${url}&comp=blocklist`, Buffer.from(xml), { "x-ms-blob-content-type": contentType, "content-type": "application/xml" }); |
| 135 | return size; |
| 136 | } finally { |
| 137 | await fh.close(); |
| 138 | } |
| 139 | } |
| 140 | |
| 141 | /** |
| 142 | * Downloads a signed URL to a file. Large blobs come down in ranges, several at once, which is |
| 143 | * what makes a cache or artifact of a few hundred MB restore in seconds (one stream from R2 through |
| 144 | * the Worker often runs at 20-30 MB/s). Returns the response's content type and file name. |
| 145 | */ |
| 146 | async download(url: string, file: string): Promise<{ contentType: string; fileName: string | null }> { |
| 147 | const out = createWriteStream(file); |
| 148 | const dl = this.downloadStream(url); |
| 149 | await pipeline(dl, out); |
| 150 | return dl.meta!; |
| 151 | } |
| 152 | |
| 153 | /** |
| 154 | * The blob at a signed URL as a stream, in order. It asks for the first range; a server that |
| 155 | * answers with ranges (as Diffinery's does) gets the rest DOWNLOAD_CONCURRENCY at a time, each |
| 156 | * range retried on its own, and at most that many ranges are held in memory. A server that |
| 157 | * ignores ranges sends the whole body in one response. Every range must come from the same |
| 158 | * object (size and ETag), or the stream fails rather than mix two versions. |
| 159 | */ |
| 160 | downloadStream(url: string, o: { onSize?: (size: number) => void } = {}): Readable & { meta?: { contentType: string; fileName: string | null }; size?: number } { |
| 161 | const abort = new AbortController(); |
| 162 | const client = this; |
| 163 | const get = (range: { start: number; end: number } | null) => |
| 164 | fetch(url, { headers: { "user-agent": "diffinery-runner", ...(range ? { range: `bytes=${range.start}-${range.end}` } : {}) }, signal: abort.signal }); |
| 165 | const fail = async (res: Response) => { |
| 166 | const text = await res.text().catch(() => ""); |
| 167 | return new RuntimeError(res.status, "download", `download failed (${res.status}): ${text.slice(0, 200)}`); |
| 168 | }; |
| 169 | async function* chunks(this: void): AsyncGenerator<Buffer> { |
| 170 | // The first range, or the whole body from a server that doesn't do ranges (or an empty blob: 416). |
| 171 | const first = await client.withRetries("download", async () => { |
| 172 | let res = await get({ start: 0, end: DOWNLOAD_CHUNK - 1 }); |
| 173 | if (res.status === 416) { |
| 174 | await res.body?.cancel(); |
| 175 | res = await get(null); |
| 176 | } |
| 177 | if (!res.ok || !res.body) throw await fail(res); |
| 178 | return res; |
| 179 | }); |
| 180 | stream.meta = { contentType: (first.headers.get("content-type") ?? "").split(";")[0]!.trim().toLowerCase(), fileName: fileNameOf(first.headers.get("content-disposition") ?? "") }; |
| 181 | const total = first.status === 206 ? Number(/\/(\d+)$/.exec(first.headers.get("content-range") ?? "")?.[1]) : NaN; |
| 182 | if (first.status !== 206 || !Number.isSafeInteger(total)) { |
| 183 | const declared = first.headers.get("content-length"); |
| 184 | // fetch decodes a compressed body, so its length only counts when the body isn't encoded. |
| 185 | const length = declared === null || first.headers.get("content-encoding") ? NaN : Number(declared); |
| 186 | if (Number.isSafeInteger(length)) o.onSize?.(length); |
| 187 | let n = 0; |
| 188 | for await (const c of Readable.fromWeb(first.body as unknown as WebReadableStream)) { |
| 189 | n += (c as Buffer).length; |
| 190 | yield c as Buffer; |
| 191 | } |
| 192 | if (Number.isSafeInteger(length) && n !== length) throw new Error(`download: got ${n} of ${length} bytes`); |
| 193 | stream.size = n; |
| 194 | return; |
| 195 | } |
| 196 | stream.size = total; |
| 197 | o.onSize?.(total); |
| 198 | const etag = first.headers.get("etag"); |
| 199 | const ranges: Array<{ start: number; end: number }> = []; |
| 200 | for (let start = 0; start < total; start += DOWNLOAD_CHUNK) ranges.push({ start, end: Math.min(start + DOWNLOAD_CHUNK, total) - 1 }); |
| 201 | const body = async (res: Response, r: { start: number; end: number }) => { |
| 202 | const buf = Buffer.from(await res.arrayBuffer()); |
| 203 | if (buf.length !== r.end - r.start + 1) throw new Error(`download: got ${buf.length} bytes for bytes ${r.start}-${r.end}`); |
| 204 | return buf; |
| 205 | }; |
| 206 | const fetchRange = (r: { start: number; end: number }) => |
| 207 | client.withRetries("download", async () => { |
| 208 | const res = await get(r); |
| 209 | if (res.status !== 206) { |
| 210 | await res.body?.cancel(); |
| 211 | throw res.ok ? new RuntimeError(412, "download", `download failed: the server stopped answering with ranges (${res.status})`) : await fail(res); |
| 212 | } |
| 213 | const m = /^bytes (\d+)-(\d+)\/(\d+)$/.exec(res.headers.get("content-range") ?? ""); |
| 214 | if (!m || Number(m[1]) !== r.start || Number(m[2]) !== r.end || Number(m[3]) !== total || (etag && res.headers.get("etag") !== etag)) { |
| 215 | await res.body?.cancel(); |
| 216 | throw new RuntimeError(412, "download", "download failed: the blob changed while it was being downloaded"); |
| 217 | } |
| 218 | return body(res, r); |
| 219 | }); |
| 220 | const pending: Array<Promise<Buffer>> = [body(first, ranges[0]!).catch(() => fetchRange(ranges[0]!))]; |
| 221 | let next = 1; |
| 222 | const fill = () => { |
| 223 | while (next < ranges.length && pending.length < DOWNLOAD_CONCURRENCY) { |
| 224 | const p = fetchRange(ranges[next++]!); |
| 225 | p.catch(() => {}); // awaited in order below; this keeps an early failure from going unhandled meanwhile |
| 226 | pending.push(p); |
| 227 | } |
| 228 | }; |
| 229 | fill(); |
| 230 | while (pending.length) { |
| 231 | const buf = await pending.shift()!; |
| 232 | fill(); |
| 233 | yield buf; |
| 234 | } |
| 235 | } |
| 236 | const stream: Readable & { meta?: { contentType: string; fileName: string | null }; size?: number } = Readable.from(chunks(), { objectMode: false }); |
| 237 | stream.once("close", () => abort.abort()); |
| 238 | return stream; |
| 239 | } |
| 240 | } |
| 241 | |
| 242 | function fileNameOf(cd: string): string | null { |
| 243 | const star = /filename\*\s*=\s*UTF-8''([^;\r\n]*)/i.exec(cd)?.[1]; |
| 244 | const plain = /(?<!\*)filename\s*=\s*"?([^;\r\n"]*)"?/i.exec(cd)?.[1]; |
| 245 | try { |
| 246 | return star ? decodeURIComponent(star) : (plain ?? null); |
| 247 | } catch { |
| 248 | return plain ?? null; |
| 249 | } |
| 250 | } |