Extract timeout and apply to it to CLI downloads

Author: koesie10

extensions/ql-vscode/src/codeql-cli/distribution.ts

@@ -1,3 +1,4 @@
+import type { WriteStream } from "fs";
 import { createWriteStream, mkdtemp, pathExists, remove } from "fs-extra";
 import { tmpdir } from "os";
 import { delimiter, dirname, join } from "path";
@@ -26,6 +27,8 @@ import { unzipToDirectoryConcurrently } from "../common/unzip-concurrently";
 import { reportUnzipProgress } from "../common/vscode/unzip-progress";
 import type { Release } from "./distribution/release";
 import { ReleasesApiConsumer } from "./distribution/releases-api-consumer";
+import { createTimeoutSignal } from "../common/fetch-stream";
+import { AbortError } from "node-fetch";
 
 /**
  * distribution.ts
@@ -384,15 +387,21 @@ class ExtensionSpecificDistributionManager {
       );
     }
 
-    const assetStream =
-      await this.createReleasesApiConsumer().streamBinaryContentOfAsset(
-        assets[0],
-      );
+    const { signal, onData, dispose: disposeTimeout } = createTimeoutSignal(30);
+
     const tmpDirectory = await mkdtemp(join(tmpdir(), "vscode-codeql"));
 
+    let archiveFile: WriteStream | undefined = undefined;
+
     try {
+      const assetStream =
+        await this.createReleasesApiConsumer().streamBinaryContentOfAsset(
+          assets[0],
+          signal,
+        );
+
       const archivePath = join(tmpDirectory, "distributionDownload.zip");
-      const archiveFile = createWriteStream(archivePath);
+      archiveFile = createWriteStream(archivePath);
 
       const contentLength = assetStream.headers.get("content-length");
       const totalNumBytes = contentLength
@@ -405,12 +414,23 @@ class ExtensionSpecificDistributionManager {
         progressCallback,
       );
 
-      await new Promise((resolve, reject) =>
+      assetStream.body.on("data", onData);
+
+      await new Promise((resolve, reject) => {
+        if (!archiveFile) {
+          throw new Error("Invariant violation: archiveFile not set");
+        }
+
         assetStream.body
           .pipe(archiveFile)
           .on("finish", resolve)
-          .on("error", reject),
-      );
+          .on("error", reject);
+
+        // If an error occurs on the body, we also want to reject the promise (e.g. during a timeout error).
+        assetStream.body.on("error", reject);
+      });
+
+      disposeTimeout();
 
       await this.bumpDistributionFolderIndex();
 
@@ -427,7 +447,19 @@ class ExtensionSpecificDistributionManager {
             )
           : undefined,
       );
+    } catch (e) {
+      if (e instanceof AbortError) {
+        const thrownError = new AbortError("The download timed out.");
+        thrownError.stack = e.stack;
+        throw thrownError;
+      }
+
+      throw e;
     } finally {
+      disposeTimeout();
+
+      archiveFile?.close();
+
       await remove(tmpDirectory);
     }
   }

extensions/ql-vscode/src/codeql-cli/distribution/releases-api-consumer.ts

@@ -90,21 +90,28 @@ export class ReleasesApiConsumer {
 
   public async streamBinaryContentOfAsset(
     asset: ReleaseAsset,
+    signal?: AbortSignal,
   ): Promise<Response> {
     const apiPath = `/repos/${this.repositoryNwo}/releases/assets/${asset.id}`;
 
-    return await this.makeApiCall(apiPath, {
-      accept: "application/octet-stream",
-    });
+    return await this.makeApiCall(
+      apiPath,
+      {
+        accept: "application/octet-stream",
+      },
+      signal,
+    );
   }
 
   protected async makeApiCall(
     apiPath: string,
     additionalHeaders: { [key: string]: string } = {},
+    signal?: AbortSignal,
   ): Promise<Response> {
     const response = await this.makeRawRequest(
       ReleasesApiConsumer.apiBase + apiPath,
       Object.assign({}, this.defaultHeaders, additionalHeaders),
+      signal,
     );
 
     if (!response.ok) {
@@ -129,11 +136,13 @@ export class ReleasesApiConsumer {
   private async makeRawRequest(
     requestUrl: string,
     headers: { [key: string]: string },
+    signal?: AbortSignal,
     redirectCount = 0,
   ): Promise<Response> {
     const response = await fetch(requestUrl, {
       headers,
       redirect: "manual",
+      signal,
     });
 
     const redirectUrl = response.headers.get("location");
@@ -153,7 +162,12 @@ export class ReleasesApiConsumer {
         // mechanism is provided.
         delete headers["authorization"];
       }
-      return await this.makeRawRequest(redirectUrl, headers, redirectCount + 1);
+      return await this.makeRawRequest(
+        redirectUrl,
+        headers,
+        signal,
+        redirectCount + 1,
+      );
     }
 
     return response;

extensions/ql-vscode/src/common/fetch-stream.ts

@@ -0,0 +1,36 @@
+import { clearTimeout } from "node:timers";
+
+export function createTimeoutSignal(timeoutSeconds: number): {
+  signal: AbortSignal;
+  onData: () => void;
+  dispose: () => void;
+} {
+  const timeout = timeoutSeconds * 1000;
+
+  const abortController = new AbortController();
+
+  let timeoutId: NodeJS.Timeout;
+
+  // If we don't get any data within the timeout, abort the download
+  timeoutId = setTimeout(() => {
+    abortController.abort();
+  }, timeout);
+
+  // If we receive any data within the timeout, reset the timeout
+  const onData = () => {
+    clearTimeout(timeoutId);
+    timeoutId = setTimeout(() => {
+      abortController.abort();
+    }, timeout);
+  };
+
+  const dispose = () => {
+    clearTimeout(timeoutId);
+  };
+
+  return {
+    signal: abortController.signal,
+    onData,
+    dispose,
+  };
+}

extensions/ql-vscode/src/databases/database-fetcher.ts

@@ -37,7 +37,7 @@ import { showAndLogInformationMessage } from "../common/logging";
 import { AppOctokit } from "../common/octokit";
 import { getLanguageDisplayName } from "../common/query-language";
 import type { DatabaseOrigin } from "./local-databases/database-origin";
-import { clearTimeout } from "node:timers";
+import { createTimeoutSignal } from "../common/fetch-stream";
 
 /**
  * Prompts a user to fetch a database from a remote location. Database is assumed to be an archive file.
@@ -483,28 +483,23 @@ async function fetchAndUnzip(
     step: 1,
   });
 
-  const abortController = new AbortController();
-
-  const timeout = downloadTimeout() * 1000;
-
-  let timeoutId: NodeJS.Timeout;
-
-  // If we don't get any data within the timeout, abort the download
-  timeoutId = setTimeout(() => {
-    abortController.abort();
-  }, timeout);
+  const {
+    signal,
+    onData,
+    dispose: disposeTimeout,
+  } = createTimeoutSignal(downloadTimeout());
 
   let response: Response;
   try {
     response = await checkForFailingResponse(
       await fetch(databaseUrl, {
         headers: requestHeaders,
-        signal: abortController.signal,
+        signal,
       }),
       "Error downloading database",
     );
   } catch (e) {
-    clearTimeout(timeoutId);
+    disposeTimeout();
 
     if (e instanceof AbortError) {
       const thrownError = new AbortError("The request timed out.");
@@ -526,13 +521,7 @@ async function fetchAndUnzip(
     progress,
   );
 
-  // If we receive any data within the timeout, reset the timeout
-  response.body.on("data", () => {
-    clearTimeout(timeoutId);
-    timeoutId = setTimeout(() => {
-      abortController.abort();
-    }, timeout);
-  });
+  response.body.on("data", onData);
 
   try {
     await new Promise((resolve, reject) => {
@@ -558,7 +547,7 @@ async function fetchAndUnzip(
 
     throw e;
   } finally {
-    clearTimeout(timeoutId);
+    disposeTimeout();
   }
 
   await readAndUnzip(