RunsOn implem

src/custom/downloadUtils.ts

@@ -0,0 +1,328 @@
+// Just a copy of the original file from the toolkit/actions/cache repository, with a change for byte range used in the downloadCacheHttpClientConcurrent function.
+import * as core from "@actions/core";
+import { HttpClient } from "@actions/http-client";
+import { TransferProgressEvent } from "@azure/ms-rest-js";
+import * as fs from "fs";
+import { DownloadOptions } from "@actions/cache/lib/options";
+import { retryHttpClientResponse } from "@actions/cache/lib/internal/requestUtils";
+
+/**
+ * Class for tracking the download state and displaying stats.
+ */
+export class DownloadProgress {
+    contentLength: number;
+    segmentIndex: number;
+    segmentSize: number;
+    segmentOffset: number;
+    receivedBytes: number;
+    startTime: number;
+    displayedComplete: boolean;
+    timeoutHandle?: ReturnType<typeof setTimeout>;
+
+    constructor(contentLength: number) {
+        this.contentLength = contentLength;
+        this.segmentIndex = 0;
+        this.segmentSize = 0;
+        this.segmentOffset = 0;
+        this.receivedBytes = 0;
+        this.displayedComplete = false;
+        this.startTime = Date.now();
+    }
+
+    /**
+     * Progress to the next segment. Only call this method when the previous segment
+     * is complete.
+     *
+     * @param segmentSize the length of the next segment
+     */
+    nextSegment(segmentSize: number): void {
+        this.segmentOffset = this.segmentOffset + this.segmentSize;
+        this.segmentIndex = this.segmentIndex + 1;
+        this.segmentSize = segmentSize;
+        this.receivedBytes = 0;
+
+        core.debug(
+            `Downloading segment at offset ${this.segmentOffset} with length ${this.segmentSize}...`
+        );
+    }
+
+    /**
+     * Sets the number of bytes received for the current segment.
+     *
+     * @param receivedBytes the number of bytes received
+     */
+    setReceivedBytes(receivedBytes: number): void {
+        this.receivedBytes = receivedBytes;
+    }
+
+    /**
+     * Returns the total number of bytes transferred.
+     */
+    getTransferredBytes(): number {
+        return this.segmentOffset + this.receivedBytes;
+    }
+
+    /**
+     * Returns true if the download is complete.
+     */
+    isDone(): boolean {
+        return this.getTransferredBytes() === this.contentLength;
+    }
+
+    /**
+     * Prints the current download stats. Once the download completes, this will print one
+     * last line and then stop.
+     */
+    display(): void {
+        if (this.displayedComplete) {
+            return;
+        }
+
+        const transferredBytes = this.segmentOffset + this.receivedBytes;
+        const percentage = (
+            100 *
+            (transferredBytes / this.contentLength)
+        ).toFixed(1);
+        const elapsedTime = Date.now() - this.startTime;
+        const downloadSpeed = (
+            transferredBytes /
+            (1024 * 1024) /
+            (elapsedTime / 1000)
+        ).toFixed(1);
+
+        core.info(
+            `Received ${transferredBytes} of ${this.contentLength} (${percentage}%), ${downloadSpeed} MBs/sec`
+        );
+
+        if (this.isDone()) {
+            this.displayedComplete = true;
+        }
+    }
+
+    /**
+     * Returns a function used to handle TransferProgressEvents.
+     */
+    onProgress(): (progress: TransferProgressEvent) => void {
+        return (progress: TransferProgressEvent) => {
+            this.setReceivedBytes(progress.loadedBytes);
+        };
+    }
+
+    /**
+     * Starts the timer that displays the stats.
+     *
+     * @param delayInMs the delay between each write
+     */
+    startDisplayTimer(delayInMs = 1000): void {
+        const displayCallback = (): void => {
+            this.display();
+
+            if (!this.isDone()) {
+                this.timeoutHandle = setTimeout(displayCallback, delayInMs);
+            }
+        };
+
+        this.timeoutHandle = setTimeout(displayCallback, delayInMs);
+    }
+
+    /**
+     * Stops the timer that displays the stats. As this typically indicates the download
+     * is complete, this will display one last line, unless the last line has already
+     * been written.
+     */
+    stopDisplayTimer(): void {
+        if (this.timeoutHandle) {
+            clearTimeout(this.timeoutHandle);
+            this.timeoutHandle = undefined;
+        }
+
+        this.display();
+    }
+}
+
+/**
+ * Download the cache using the Actions toolkit http-client concurrently
+ *
+ * @param archiveLocation the URL for the cache
+ * @param archivePath the local path where the cache is saved
+ */
+export async function downloadCacheHttpClientConcurrent(
+    archiveLocation: string,
+    archivePath: fs.PathLike,
+    options: DownloadOptions
+): Promise<void> {
+    const archiveDescriptor = await fs.promises.open(archivePath, "w");
+    const httpClient = new HttpClient("actions/cache", undefined, {
+        socketTimeout: options.timeoutInMs,
+        keepAlive: true
+    });
+    try {
+        const res = await retryHttpClientResponse(
+            "downloadCacheMetadata",
+            async () =>
+                await httpClient.request("GET", archiveLocation, null, {
+                    Range: "bytes=0-1"
+                })
+        );
+
+        const contentRange = res.message.headers["content-range"];
+        if (!contentRange) {
+            throw new Error("Range request not supported by server");
+        }
+        const match = contentRange?.match(/bytes \d+-\d+\/(\d+)/);
+        if (!match) {
+            throw new Error(
+                "Content-Range header in server response not in correct format"
+            );
+        }
+        const length = parseInt(match[1]);
+        if (Number.isNaN(length)) {
+            throw new Error(`Could not interpret Content-Length: ${length}`);
+        }
+
+        const downloads: {
+            offset: number;
+            promiseGetter: () => Promise<DownloadSegment>;
+        }[] = [];
+
+        const blockSize = 32 * 1024 * 1024;
+
+        for (let offset = 0; offset < length; offset += blockSize) {
+            const count = Math.min(blockSize, length - offset);
+            downloads.push({
+                offset,
+                promiseGetter: async () => {
+                    return await downloadSegmentRetry(
+                        httpClient,
+                        archiveLocation,
+                        offset,
+                        count
+                    );
+                }
+            });
+        }
+
+        // reverse to use .pop instead of .shift
+        downloads.reverse();
+        let actives = 0;
+        let bytesDownloaded = 0;
+        const progress = new DownloadProgress(length);
+        progress.startDisplayTimer();
+        const progressFn = progress.onProgress();
+
+        const activeDownloads: { [offset: number]: Promise<DownloadSegment> } =
+            [];
+        let nextDownload:
+            | { offset: number; promiseGetter: () => Promise<DownloadSegment> }
+            | undefined;
+
+        const waitAndWrite: () => Promise<void> = async () => {
+            const segment = await Promise.race(Object.values(activeDownloads));
+            await archiveDescriptor.write(
+                segment.buffer,
+                0,
+                segment.count,
+                segment.offset
+            );
+            actives--;
+            delete activeDownloads[segment.offset];
+            bytesDownloaded += segment.count;
+            progressFn({ loadedBytes: bytesDownloaded });
+        };
+
+        while ((nextDownload = downloads.pop())) {
+            activeDownloads[nextDownload.offset] = nextDownload.promiseGetter();
+            actives++;
+
+            if (actives >= (options.downloadConcurrency ?? 10)) {
+                await waitAndWrite();
+            }
+        }
+
+        while (actives > 0) {
+            await waitAndWrite();
+        }
+    } finally {
+        httpClient.dispose();
+        await archiveDescriptor.close();
+    }
+}
+
+async function downloadSegmentRetry(
+    httpClient: HttpClient,
+    archiveLocation: string,
+    offset: number,
+    count: number
+): Promise<DownloadSegment> {
+    const retries = 5;
+    let failures = 0;
+
+    while (true) {
+        try {
+            const timeout = 30000;
+            const result = await promiseWithTimeout(
+                timeout,
+                downloadSegment(httpClient, archiveLocation, offset, count)
+            );
+            if (typeof result === "string") {
+                throw new Error("downloadSegmentRetry failed due to timeout");
+            }
+
+            return result;
+        } catch (err) {
+            if (failures >= retries) {
+                throw err;
+            }
+
+            failures++;
+        }
+    }
+}
+
+async function downloadSegment(
+    httpClient: HttpClient,
+    archiveLocation: string,
+    offset: number,
+    count: number
+): Promise<DownloadSegment> {
+    const partRes = await retryHttpClientResponse(
+        "downloadCachePart",
+        async () =>
+            await httpClient.get(archiveLocation, {
+                Range: `bytes=${offset}-${offset + count - 1}`
+            })
+    );
+
+    if (!partRes.readBodyBuffer) {
+        throw new Error(
+            "Expected HttpClientResponse to implement readBodyBuffer"
+        );
+    }
+
+    return {
+        offset,
+        count,
+        buffer: await partRes.readBodyBuffer()
+    };
+}
+
+declare class DownloadSegment {
+    offset: number;
+    count: number;
+    buffer: Buffer;
+}
+
+const promiseWithTimeout = async <T>(
+    timeoutMs: number,
+    promise: Promise<T>
+): Promise<T | string> => {
+    let timeoutHandle: NodeJS.Timeout;
+    const timeoutPromise = new Promise<string>(resolve => {
+        timeoutHandle = setTimeout(() => resolve("timeout"), timeoutMs);
+    });
+
+    return Promise.race([promise, timeoutPromise]).then(result => {
+        clearTimeout(timeoutHandle);
+        return result;
+    });
+};