Merge pull request #1293 from github/edoardo/fix-with-timeout

Fix `withTimeout` helper function

src/trap-caching.ts

@@ -113,15 +113,15 @@ export async function downloadTrapCaches(
     logger.info(
       `Looking in Actions cache for TRAP cache with key ${preferredKey}`
     );
-    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`
-        );
+    const found = await cache.restoreCache(
+      [cacheDir],
+      preferredKey,
+      [
+        // Fall back to any cache with the right key prefix
+        await cachePrefix(codeql, language),
+      ],
+      {
+        segmentTimeoutInMs: MAX_CACHE_OPERATION_MS,
       }
     );
     if (found === undefined) {

src/util.test.ts

@@ -632,3 +632,18 @@ test("withTimeout on short task", async (t) => {
   t.deepEqual(shortTaskTimedOut, false);
   t.deepEqual(result, 99);
 });
+
+test("withTimeout doesn't call callback if promise resolves", async (t) => {
+  let shortTaskTimedOut = false;
+  const shortTask = new Promise((resolve) => {
+    setTimeout(() => {
+      resolve(99);
+    }, shortTime);
+  });
+  const result = await util.withTimeout(100, shortTask, () => {
+    shortTaskTimedOut = true;
+  });
+  await new Promise((r) => setTimeout(r, 200));
+  t.deepEqual(shortTaskTimedOut, false);
+  t.deepEqual(result, 99);
+});

src/util.ts

@@ -899,6 +899,9 @@ export async function tryGetFolderBytes(
  * 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.
  *
+ * Note that this does NOT cancel the original promise, so that promise will
+ * continue in the background even after the timeout has expired.
+ *
  * @param timeoutMs The timeout in milliseconds.
  * @param promise The promise to run.
  * @param onTimeout A callback to call if the promise times out.
@@ -909,12 +912,18 @@ export async function withTimeout<T>(
   promise: Promise<T>,
   onTimeout: () => void
 ): Promise<T | undefined> {
+  let finished = false;
+  const mainTask = async () => {
+    const result = await promise;
+    finished = true;
+    return result;
+  };
   const timeout: Promise<undefined> = new Promise((resolve) => {
     setTimeout(() => {
-      onTimeout();
+      if (!finished) onTimeout();
       resolve(undefined);
     }, timeoutMs);
   });
 
-  return await Promise.race([promise, timeout]);
+  return await Promise.race([mainTask(), timeout]);
 }