From 5afc042a7457ece5073f00cb20c423145d045d1e Mon Sep 17 00:00:00 2001
From: Bassem Dghaidi <568794+Link-@users.noreply.github.com>
Date: Mon, 17 Jun 2024 01:17:10 -0700
Subject: [PATCH] Add download cache v2

---
 packages/cache/src/cache.ts                   | 73 ++++++++++++++++++-
 .../cache/src/internal/v2/download-cache.ts   | 67 +++++++++++++++++
 2 files changed, 139 insertions(+), 1 deletion(-)
 create mode 100644 packages/cache/src/internal/v2/download-cache.ts

diff --git a/packages/cache/src/cache.ts b/packages/cache/src/cache.ts
index 5a582f8d..e93ffd4b 100644
--- a/packages/cache/src/cache.ts
+++ b/packages/cache/src/cache.ts
@@ -6,8 +6,14 @@ import * as cacheHttpClient from './internal/cacheHttpClient'
 import * as cacheTwirpClient from './internal/cacheTwirpClient'
 import {createTar, extractTar, listTar} from './internal/tar'
 import {DownloadOptions, UploadOptions} from './options'
-import {GetCacheBlobUploadURLRequest, GetCacheBlobUploadURLResponse} from './generated/results/api/v1/blobcache'
+import {
+  GetCacheBlobUploadURLRequest,
+  GetCacheBlobUploadURLResponse,
+  GetCachedBlobRequest,
+  GetCachedBlobResponse
+} from './generated/results/api/v1/blobcache'
 import {UploadCacheStream} from './internal/v2/upload-cache'
+import {StreamExtract} from './internal/v2/download-cache'
 import {
   UploadZipSpecification,
   getUploadZipSpecification
@@ -81,6 +87,23 @@ export async function restoreCache(
 ): Promise<string | undefined> {
   checkPaths(paths)
 
+  console.debug(`Cache Service Version: ${CacheServiceVersion}`)
+  switch (CacheServiceVersion) {
+    case "v2":
+      return await restoreCachev2(paths, primaryKey, restoreKeys, options, enableCrossOsArchive)
+    case "v1":
+    default:
+      return await restoreCachev1(paths, primaryKey, restoreKeys, options, enableCrossOsArchive)
+  }
+}
+
+async function restoreCachev1(
+  paths: string[],
+  primaryKey: string,
+  restoreKeys?: string[],
+  options?: DownloadOptions,
+  enableCrossOsArchive = false
+) {
   restoreKeys = restoreKeys || []
   const keys = [primaryKey, ...restoreKeys]
 
@@ -162,6 +185,54 @@ export async function restoreCache(
   return undefined
 }
 
+async function restoreCachev2(
+  paths: string[],
+  primaryKey: string,
+  restoreKeys?: string[],
+  options?: DownloadOptions,
+  enableCrossOsArchive = false
+) {
+  
+  restoreKeys = restoreKeys || []
+  const keys = [primaryKey, ...restoreKeys]
+
+  core.debug('Resolved Keys:')
+  core.debug(JSON.stringify(keys))
+
+  if (keys.length > 10) {
+    throw new ValidationError(
+      `Key Validation Error: Keys are limited to a maximum of 10.`
+    )
+  }
+  for (const key of keys) {
+    checkKey(key)
+  }
+
+  try {
+    const twirpClient = cacheTwirpClient.internalBlobCacheTwirpClient()
+    const getSignedDownloadURLRequest: GetCachedBlobRequest = {
+      owner: "github",
+      keys: keys,
+    }
+    const signedDownloadURL: GetCachedBlobResponse = await twirpClient.GetCachedBlob(getSignedDownloadURLRequest)
+    core.info(`GetCachedBlobResponse: ${JSON.stringify(signedDownloadURL)}`)
+
+    if (signedDownloadURL.blobs.length === 0) {
+      // Cache not found
+      core.warning(`Cache not found for keys: ${keys.join(', ')}`)
+      return undefined
+    }
+
+    core.info(`Starting download of artifact to: ${paths[0]}`)
+    await StreamExtract(signedDownloadURL.blobs[0].signedUrl, paths[0])
+    core.info(`Artifact download completed successfully.`)
+  } catch (error) {
+    throw new Error(`Unable to download and extract cache: ${error.message}`)
+  }
+  
+  return undefined
+}
+
 /**
  * Saves a list of files with the specified key
  *
diff --git a/packages/cache/src/internal/v2/download-cache.ts b/packages/cache/src/internal/v2/download-cache.ts
new file mode 100644
index 00000000..bfba0d70
--- /dev/null
+++ b/packages/cache/src/internal/v2/download-cache.ts
@@ -0,0 +1,67 @@
+import * as core from '@actions/core'
+import * as httpClient from '@actions/http-client'
+import unzip from 'unzip-stream'
+const packageJson = require('../../../package.json')
+
+export async function StreamExtract(url: string, directory: string): Promise<void> {
+    let retryCount = 0
+    while (retryCount < 5) {
+      try {
+        await streamExtractExternal(url, directory)
+        return
+      } catch (error) {
+        retryCount++
+        core.debug(
+          `Failed to download cache after ${retryCount} retries due to ${error.message}. Retrying in 5 seconds...`
+        )
+        // wait 5 seconds before retrying
+        await new Promise(resolve => setTimeout(resolve, 5000))
+      }
+    }
+
+    throw new Error(`Cache download failed after ${retryCount} retries.`)
+}
+
+export async function streamExtractExternal(
+    url: string,
+    directory: string
+  ): Promise<void> {
+    const client = new httpClient.HttpClient(`@actions/cache-${packageJson.version}`)
+    const response = await client.get(url)
+    if (response.message.statusCode !== 200) {
+      throw new Error(
+        `Unexpected HTTP response from blob storage: ${response.message.statusCode} ${response.message.statusMessage}`
+      )
+    }
+  
+    const timeout = 30 * 1000 // 30 seconds
+  
+    return new Promise((resolve, reject) => {
+      const timerFn = (): void => {
+        response.message.destroy(
+          new Error(`Blob storage chunk did not respond in ${timeout}ms`)
+        )
+      }
+      const timer = setTimeout(timerFn, timeout)
+  
+      response.message
+        .on('data', () => {
+          timer.refresh()
+        })
+        .on('error', (error: Error) => {
+          core.debug(
+            `response.message: Cache download failed: ${error.message}`
+          )
+          clearTimeout(timer)
+          reject(error)
+        })
+        .pipe(unzip.Extract({path: directory}))
+        .on('close', () => {
+          clearTimeout(timer)
+          resolve()
+        })
+        .on('error', (error: Error) => {
+          reject(error)
+        })
+    })
+  }
\ No newline at end of file