Add option for concurrent cache downloads with timeout

packages/cache/RELEASES.md

@@ -160,3 +160,7 @@
 ### 3.2.1
 
 - Updated @azure/storage-blob to `v12.13.0`
+
+### 3.2.2
+
+- Add new default cache download method to improve performance and reduce hangs [#1484](https://github.com/actions/toolkit/pull/1484)

packages/cache/__tests__/cacheHttpClient.test.ts

@@ -84,18 +84,24 @@
     'downloadCacheStorageSDK'
   )
 
+  const downloadCacheHttpClientConcurrentMock = jest.spyOn(
+    downloadUtils,
+    'downloadCacheHttpClientConcurrent'
+  )
+
   const archiveLocation = 'http://foo.blob.core.windows.net/bar/baz'
   const archivePath = '/foo/bar'
 
   await downloadCache(archiveLocation, archivePath)
 
-  expect(downloadCacheStorageSDKMock).toHaveBeenCalledTimes(1)
-  expect(downloadCacheStorageSDKMock).toHaveBeenCalledWith(
+  expect(downloadCacheHttpClientConcurrentMock).toHaveBeenCalledTimes(1)
+  expect(downloadCacheHttpClientConcurrentMock).toHaveBeenCalledWith(
     archiveLocation,
     archivePath,
     getDownloadOptions()
   )
 
+  expect(downloadCacheStorageSDKMock).toHaveBeenCalledTimes(0)
   expect(downloadCacheHttpClientMock).toHaveBeenCalledTimes(0)
 })
 
@@ -109,20 +115,26 @@
     'downloadCacheStorageSDK'
   )
 
+  const downloadCacheHttpClientConcurrentMock = jest.spyOn(
+    downloadUtils,
+    'downloadCacheHttpClientConcurrent'
+  )
+
   const archiveLocation = 'http://foo.blob.core.windows.net/bar/baz'
   const archivePath = '/foo/bar'
   const options: DownloadOptions = {downloadConcurrency: 4}
 
   await downloadCache(archiveLocation, archivePath, options)
 
-  expect(downloadCacheStorageSDKMock).toHaveBeenCalledTimes(1)
-  expect(downloadCacheStorageSDKMock).toHaveBeenCalled()
-  expect(downloadCacheStorageSDKMock).toHaveBeenCalledWith(
+  expect(downloadCacheHttpClientConcurrentMock).toHaveBeenCalledTimes(1)
+  expect(downloadCacheHttpClientConcurrentMock).toHaveBeenCalled()
+  expect(downloadCacheHttpClientConcurrentMock).toHaveBeenCalledWith(
     archiveLocation,
     archivePath,
     getDownloadOptions(options)
   )
 
+  expect(downloadCacheStorageSDKMock).toHaveBeenCalledTimes(0)
   expect(downloadCacheHttpClientMock).toHaveBeenCalledTimes(0)
 })
 
@@ -138,7 +150,7 @@
 
   const archiveLocation = 'http://foo.blob.core.windows.net/bar/baz'
   const archivePath = '/foo/bar'
-  const options: DownloadOptions = {useAzureSdk: false}
+  const options: DownloadOptions = {useAzureSdk: false, concurrentBlobDownloads: false}
 
   await downloadCache(archiveLocation, archivePath, options)
 

packages/cache/__tests__/options.test.ts

@@ -5,7 +5,8 @@ import {
   getUploadOptions
 } from '../src/options'
 
-const useAzureSdk = true
+const useAzureSdk = false
+const concurrentBlobDownloads = true
 const downloadConcurrency = 8
 const timeoutInMs = 30000
 const segmentTimeoutInMs = 600000
@@ -18,6 +19,7 @@ test('getDownloadOptions sets defaults', async () => {
 
   expect(actualOptions).toEqual({
     useAzureSdk,
+    concurrentBlobDownloads,
     downloadConcurrency,
     timeoutInMs,
     segmentTimeoutInMs,
@@ -27,7 +29,8 @@ test('getDownloadOptions sets defaults', async () => {
 
 test('getDownloadOptions overrides all settings', async () => {
   const expectedOptions: DownloadOptions = {
-    useAzureSdk: false,
+    useAzureSdk: true,
+    concurrentBlobDownloads: false,
     downloadConcurrency: 14,
     timeoutInMs: 20000,
     segmentTimeoutInMs: 3600000,

packages/cache/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@actions/cache",
-  "version": "3.2.1",
+  "version": "3.2.2",
   "preview": true,
   "description": "Actions cache lib",
   "keywords": [
@@ -40,7 +40,7 @@
     "@actions/core": "^1.10.0",
     "@actions/exec": "^1.0.1",
     "@actions/glob": "^0.1.0",
-    "@actions/http-client": "^2.0.1",
+    "@actions/http-client": "^2.1.1",
     "@actions/io": "^1.0.1",
     "@azure/abort-controller": "^1.1.0",
     "@azure/ms-rest-js": "^2.6.0",

packages/cache/src/internal/cacheHttpClient.ts

@@ -20,7 +20,7 @@
   ITypedResponseWithError,
   ArtifactCacheList
 } from './contracts'
-import {downloadCacheHttpClient, downloadCacheStorageSDK} from './downloadUtils'
+import {downloadCacheHttpClient, downloadCacheHttpClientConcurrent, downloadCacheStorageSDK} from './downloadUtils'
 import {
   DownloadOptions,
   UploadOptions,
@@ -171,15 +171,19 @@
   const archiveUrl = new URL(archiveLocation)
   const downloadOptions = getDownloadOptions(options)
 
-  if (
-    downloadOptions.useAzureSdk &&
-    archiveUrl.hostname.endsWith('.blob.core.windows.net')
-  ) {
-    // Use Azure storage SDK to download caches hosted on Azure to improve speed and reliability.
-    await downloadCacheStorageSDK(archiveLocation, archivePath, downloadOptions)
+  if (archiveUrl.hostname.endsWith('.blob.core.windows.net')) {
+    if (downloadOptions.useAzureSdk) {
+      // Use Azure storage SDK to download caches hosted on Azure to improve speed and reliability.
+      await downloadCacheStorageSDK(archiveLocation, archivePath, downloadOptions)
+    } else if (downloadOptions.concurrentBlobDownloads) {
+      // Use concurrent implementation with HttpClient to work around blob SDK issue
+      await downloadCacheHttpClientConcurrent(archiveLocation, archivePath, downloadOptions)
+    } else {
+      // Otherwise, download using the Actions http-client.
+      await downloadCacheHttpClient(archiveLocation, archivePath)
+    }
   } else {
-    // Otherwise, download using the Actions http-client.
-    await downloadCacheHttpClient(archiveLocation, archivePath)
+      await downloadCacheHttpClient(archiveLocation, archivePath)
   }
 }
 

packages/cache/src/internal/downloadUtils.ts

@@ -203,6 +203,153 @@
   }
 }
 
+/**
+ * Download the cache using the Actions toolkit http-client concurrently
+ *
+ * @param archiveLocation the URL for the cache
+ * @param archivePath the local path where the cache is saved
+ */
+export async function downloadCacheHttpClientConcurrent(
+  archiveLocation: string,
+  archivePath: fs.PathLike,
+  options: DownloadOptions
+): Promise<void> {
+  const archiveDescriptor = await fs.promises.open(archivePath, 'w')
+  const httpClient = new HttpClient('actions/cache', undefined, {
+    socketTimeout: options.timeoutInMs,
+    keepAlive: true
+  })
+  try {
+    const res = await retryHttpClientResponse(
+      'downloadCacheMetadata',
+      async () => await httpClient.request('HEAD', archiveLocation, null, {})
+    )
+
+    const lengthHeader = res.message.headers['content-length']
+    if (lengthHeader === undefined || lengthHeader === null) {
+      throw new Error('Content-Length not found on blob response')
+    }
+
+    const length = parseInt(lengthHeader)
+    if (Number.isNaN(length)) {
+      throw new Error(`Could not interpret Content-Length: ${length}`)
+    }
+
+    const downloads: {
+      offset: number
+      promiseGetter: () => Promise<DownloadSegment>
+    }[] = []
+    const blockSize = 4 * 1024 * 1024
+
+    for (let offset = 0; offset < length; offset += blockSize) {
+      const count = Math.min(blockSize, length - offset)
+      downloads.push({
+        offset,
+        promiseGetter: async () => {
+          return await downloadSegmentRetry(
+            httpClient,
+            archiveLocation,
+            offset,
+            count
+          )
+        }
+      })
+    }
+
+    // reverse to use .pop instead of .shift
+    downloads.reverse()
+    let actives = 0
+    let bytesDownloaded = 0
+    const progress = new DownloadProgress(length)
+    progress.startDisplayTimer()
+    const progressFn = progress.onProgress()
+
+    const activeDownloads: {[offset: number]: Promise<DownloadSegment>} = []
+    let nextDownload:
+      | {offset: number; promiseGetter: () => Promise<DownloadSegment>}
+      | undefined
+
+    const waitAndWrite: () => Promise<void> = async () => {
+      const segment = await Promise.race(Object.values(activeDownloads))
+      await archiveDescriptor.write(segment.buffer, 0, segment.count, segment.offset);
+      actives--
+      delete activeDownloads[segment.offset]
+      bytesDownloaded += segment.count
+      progressFn({ loadedBytes: bytesDownloaded })
+    }
+
+    while ((nextDownload = downloads.pop())) {
+      activeDownloads[nextDownload.offset] = nextDownload.promiseGetter()
+      actives++
+
+      if (actives >= (options.downloadConcurrency ?? 10)) {
+        await waitAndWrite()
+      }
+    }
+
+    while (actives > 0) {
+      await waitAndWrite()
+    }
+  } finally {
+    httpClient.dispose()
+    await archiveDescriptor.close()
+  }
+}
+
+async function downloadSegmentRetry(
+  httpClient: HttpClient,
+  archiveLocation: string,
+  offset: number,
+  count: number
+): Promise<DownloadSegment> {
+  const retries = 5
+  let failures = 0
+
+  while (true) {
+    try {
+      const timeout = 30000
+      const result = await promiseWithTimeout(timeout, downloadSegment(httpClient, archiveLocation, offset, count))
+      if (typeof result === 'string') {
+        throw new Error('downloadSegmentRetry failed due to timeout')
+      }
+
+      return result
+    } catch (err) {
+      if (failures >= retries) {
+        throw err
+      }
+
+      failures++
+    }
+  }
+}
+
+async function downloadSegment(
+  httpClient: HttpClient,
+  archiveLocation: string,
+  offset: number,
+  count: number
+): Promise<DownloadSegment> {
+  const partRes = await retryHttpClientResponse(
+    'downloadCachePart',
+    async () =>
+      await httpClient.get(archiveLocation, {
+        Range: `bytes=${offset}-${offset + count - 1}`
+      })
+  )
+  return {
+    offset,
+    count,
+    buffer: await partRes.readBodyBuffer!()
+  }
+}
+
+declare class DownloadSegment {
+  offset: number
+  count: number
+  buffer: Buffer
+}
+
 /**
  * Download the cache using the Azure Storage SDK.  Only call this method if the
  * URL points to an Azure Storage endpoint.
@@ -287,12 +434,12 @@
   }
 }
 
-const promiseWithTimeout = async (
+const promiseWithTimeout = async<T> (
   timeoutMs: number,
-  promise: Promise<Buffer>
-): Promise<unknown> => {
+  promise: Promise<T>
+): Promise<T | string> => {
   let timeoutHandle: NodeJS.Timeout
-  const timeoutPromise = new Promise(resolve => {
+  const timeoutPromise = new Promise<string>(resolve => {
     timeoutHandle = setTimeout(() => resolve('timeout'), timeoutMs)
   })