From c7340e91af6697c78a17111d7d63141b793a0924 Mon Sep 17 00:00:00 2001
From: Sankalp Kotewar <98868223+kotewar@users.noreply.github.com>
Date: Thu, 11 Aug 2022 04:43:43 +0000
Subject: [PATCH] Reverted to custom promise implementation

---
 packages/cache/src/internal/downloadUtils.ts | 42 ++++++++++++++------
 1 file changed, 29 insertions(+), 13 deletions(-)

diff --git a/packages/cache/src/internal/downloadUtils.ts b/packages/cache/src/internal/downloadUtils.ts
index c92cc999..265c34f7 100644
--- a/packages/cache/src/internal/downloadUtils.ts
+++ b/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,17 +261,22 @@ 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()
-          }
+          })
         )
-        fs.writeFileSync(fd, result)
+        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()
@@ -283,3 +284,18 @@ export async function downloadCacheStorageSDK(
     }
   }
 }
+
+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
+  })
+}