Merge pull request #1150 from actions/kotewar/segment-download-bug-try-2

Implemented custom promise as abortController isn't able to stop the workflow in real scenarios

packages/cache/src/internal/downloadUtils.ts

@@ -249,12 +249,8 @@ export async function downloadCacheStorageSDK(
 
     try {
       downloadProgress.startDisplayTimer()
-      const abortSignal = AbortController.timeout(
-        options.segmentTimeoutInMs || 3600000
-      )
-      abortSignal.addEventListener('abort', () => {
-        core.warning('Aborting cache download as it exceeded the timeout.')
-      })
+      const controller = new AbortController()
+      const abortSignal = controller.signal
       while (!downloadProgress.isDone()) {
         const segmentStart =
           downloadProgress.segmentOffset + downloadProgress.segmentSize
@@ -265,21 +261,41 @@ export async function downloadCacheStorageSDK(
         )
 
         downloadProgress.nextSegment(segmentSize)
-
-        const result = await client.downloadToBuffer(
-          segmentStart,
-          segmentSize,
-          {
+        const result = await promiseWithTimeout(
+          options.segmentTimeoutInMs || 3600000,
+          client.downloadToBuffer(segmentStart, segmentSize, {
             abortSignal,
             concurrency: options.downloadConcurrency,
             onProgress: downloadProgress.onProgress()
-          }
+          })
         )
+        if (result === 'timeout') {
+          controller.abort()
+          throw new Error(
+            'Aborting cache download as the download time exceeded the timeout.'
+          )
+        } else if (Buffer.isBuffer(result)) {
           fs.writeFileSync(fd, result)
         }
+      }
     } finally {
       downloadProgress.stopDisplayTimer()
       fs.closeSync(fd)
     }
   }
 }
+
+const promiseWithTimeout = async (
+  timeoutMs: number,
+  promise: Promise<Buffer>
+): Promise<unknown> => {
+  let timeoutHandle: NodeJS.Timeout
+  const timeoutPromise = new Promise(resolve => {
+    timeoutHandle = setTimeout(() => resolve('timeout'), timeoutMs)
+  })
+
+  return Promise.race([promise, timeoutPromise]).then(result => {
+    clearTimeout(timeoutHandle)
+    return result
+  })
+}