TRAP Caching: Add timeouts to upload/download operations

src/trap-caching.ts

@@ -8,7 +8,7 @@ import { CodeQL, CODEQL_VERSION_BETTER_RESOLVE_LANGUAGES } from "./codeql";
 import { Config } from "./config-utils";
 import { Language } from "./languages";
 import { Logger } from "./logging";
-import { codeQlVersionAbove, tryGetFolderBytes } from "./util";
+import { codeQlVersionAbove, tryGetFolderBytes, withTimeout } from "./util";
 
 // This constant should be bumped if we make a breaking change
 // to how the CodeQL Action stores or retrieves the TRAP cache,
@@ -24,6 +24,12 @@ const CACHE_SIZE_MB = 1024;
 // cache for us to consider it worth uploading.
 const MINIMUM_CACHE_MB_TO_UPLOAD = 10;
 
+// The maximum number of milliseconds to wait for TRAP cache
+// uploads or downloads to complete before continuing. Note
+// this timeout is per operation, so will be run as many
+// times as there are languages with TRAP caching enabled.
+const MAX_CACHE_OPERATION_MS = 120_000; // Two minutes
+
 export async function getTrapCachingExtractorConfigArgs(
   config: Config
 ): Promise<string[]> {
@@ -107,9 +113,17 @@ export async function downloadTrapCaches(
     logger.info(
       `Looking in Actions cache for TRAP cache with key ${preferredKey}`
     );
-    const found = await cache.restoreCache([cacheDir], preferredKey, [
-      await cachePrefix(codeql, language), // Fall back to any cache with the right key prefix
-    ]);
+    const found = await withTimeout(
+      MAX_CACHE_OPERATION_MS,
+      cache.restoreCache([cacheDir], preferredKey, [
+        await cachePrefix(codeql, language), // Fall back to any cache with the right key prefix
+      ]),
+      () => {
+        logger.info(
+          `Timed out waiting for TRAP cache download for ${language}, will continue without it`
+        );
+      }
+    );
     if (found === undefined) {
       // We didn't find a TRAP cache in the Actions cache, so the directory on disk is
       // still just an empty directory. There's no reason to tell the extractor to use it,
@@ -136,7 +150,6 @@ export async function uploadTrapCaches(
 ): Promise<boolean> {
   if (!(await actionsUtil.isAnalyzingDefaultBranch())) return false; // Only upload caches from the default branch
 
-  const toAwait: Array<Promise<number>> = [];
   for (const language of config.languages) {
     const cacheDir = config.trapCaches[language];
     if (cacheDir === undefined) continue;
@@ -159,9 +172,16 @@ export async function uploadTrapCaches(
       process.env.GITHUB_SHA || "unknown"
     );
     logger.info(`Uploading TRAP cache to Actions cache with key ${key}`);
-    toAwait.push(cache.saveCache([cacheDir], key));
+    await withTimeout(
+      MAX_CACHE_OPERATION_MS,
+      cache.saveCache([cacheDir], key),
+      () => {
+        logger.info(
+          `Timed out waiting for TRAP cache for ${language} to upload, will continue without uploading`
+        );
+      }
+    );
   }
-  await Promise.all(toAwait);
   return true;
 }
 

src/util.test.ts

@@ -601,3 +601,34 @@ function mockVersion(version) {
     },
   } as CodeQL;
 }
+
+const longTime = 999_999;
+const shortTime = 10;
+
+test("withTimeout on long task", async (t) => {
+  let longTaskTimedOut = false;
+  const longTask = new Promise((resolve) => {
+    setTimeout(() => {
+      resolve(42);
+    }, longTime);
+  });
+  const result = await util.withTimeout(shortTime, longTask, () => {
+    longTaskTimedOut = true;
+  });
+  t.deepEqual(longTaskTimedOut, true);
+  t.deepEqual(result, undefined);
+});
+
+test("withTimeout on short task", async (t) => {
+  let shortTaskTimedOut = false;
+  const shortTask = new Promise((resolve) => {
+    setTimeout(() => {
+      resolve(99);
+    }, shortTime);
+  });
+  const result = await util.withTimeout(longTime, shortTask, () => {
+    shortTaskTimedOut = true;
+  });
+  t.deepEqual(shortTaskTimedOut, false);
+  t.deepEqual(result, 99);
+});

src/util.ts

@@ -878,3 +878,27 @@ export async function tryGetFolderBytes(
     return undefined;
   }
 }
+
+/**
+ * Run a promise for a given amount of time, and if it doesn't resolve within
+ * that time, call the provided callback and then return undefined.
+ *
+ * @param timeoutMs The timeout in milliseconds.
+ * @param promise The promise to run.
+ * @param onTimeout A callback to call if the promise times out.
+ * @returns The result of the promise, or undefined if the promise times out.
+ */
+export async function withTimeout<T>(
+  timeoutMs: number,
+  promise: Promise<T>,
+  onTimeout: () => void
+): Promise<T | undefined> {
+  const timeout: Promise<undefined> = new Promise((resolve) => {
+    setTimeout(() => {
+      onTimeout();
+      resolve(undefined);
+    }, timeoutMs);
+  });
+
+  return await Promise.race([promise, timeout]);
+}