From 8613989147fb54887d73493ce0743c5a4a897cf3 Mon Sep 17 00:00:00 2001
From: Deepak Dahiya <59823596+t-dedah@users.noreply.github.com>
Date: Tue, 29 Mar 2022 22:17:04 +0000
Subject: [PATCH] Added cacheSize in ReserveCache API request

---
 src/cache.ts                    | 42 ++++++++++++++++-----------------
 src/internal/cacheHttpClient.ts | 13 ++++++++--
 src/internal/cacheUtils.ts      |  7 ++++++
 src/internal/contracts.d.ts     |  2 ++
 4 files changed, 40 insertions(+), 24 deletions(-)

diff --git a/src/cache.ts b/src/cache.ts
index 9356e83..6847cf0 100644
--- a/src/cache.ts
+++ b/src/cache.ts
@@ -43,16 +43,6 @@ function checkKey(key: string): void {
   }
 }
 
-/**
- * isFeatureAvailable to check the presence of Actions cache service
- *
- * @returns boolean return true if Actions cache service feature is available, otherwise false
- */
-
-export function isFeatureAvailable(): boolean {
-  return !!process.env['ACTIONS_CACHE_URL']
-}
-
 /**
  * Restores cache from keys
  *
@@ -152,17 +142,7 @@ export async function saveCache(
   checkKey(key)
 
   const compressionMethod = await utils.getCompressionMethod()
-
-  core.debug('Reserving Cache')
-  const cacheId = await cacheHttpClient.reserveCache(key, paths, {
-    compressionMethod
-  })
-  if (cacheId === -1) {
-    throw new ReserveCacheError(
-      `Unable to reserve cache with key ${key}, another job may be creating this cache.`
-    )
-  }
-  core.debug(`Cache ID: ${cacheId}`)
+  let cacheId = null
 
   const cachePaths = await utils.resolvePaths(paths)
   core.debug('Cache Paths:')
@@ -185,7 +165,7 @@ export async function saveCache(
     const fileSizeLimit = 10 * 1024 * 1024 * 1024 // 10GB per repo limit
     const archiveFileSize = utils.getArchiveFileSizeInBytes(archivePath)
     core.debug(`File Size: ${archiveFileSize}`)
-    if (archiveFileSize > fileSizeLimit) {
+    if (archiveFileSize > fileSizeLimit && !utils.isGhes()) {
       throw new Error(
         `Cache size of ~${Math.round(
           archiveFileSize / (1024 * 1024)
@@ -193,6 +173,24 @@ export async function saveCache(
       )
     }
 
+    const cacheSize = utils.isGhes() ? archiveFileSize : undefined
+    core.debug('Reserving Cache')
+    cacheId = await cacheHttpClient.reserveCache(key, paths, {
+      compressionMethod,
+      cacheSize
+    })
+    if (cacheId === -1) {
+      throw new ReserveCacheError(
+        `Unable to reserve cache with key ${key}, another job may be creating this cache.`
+      )
+    }
+    if (cacheId === -2) {
+      throw new ReserveCacheError(
+        `Cache size of ~${Math.round(archiveFileSize / (1024 * 1024))} MB (${archiveFileSize} B) is over the data cap limit, not saving cache.`
+      )
+    }
+    core.debug(`Cache ID: ${cacheId}`)
+
     core.debug(`Saving Cache (ID: ${cacheId})`)
     await cacheHttpClient.saveCache(cacheId, archivePath, options)
   } finally {
diff --git a/src/internal/cacheHttpClient.ts b/src/internal/cacheHttpClient.ts
index 65d7bd0..e0f33bb 100644
--- a/src/internal/cacheHttpClient.ts
+++ b/src/internal/cacheHttpClient.ts
@@ -31,7 +31,12 @@ import {
 const versionSalt = '1.0'
 
 function getCacheApiUrl(resource: string): string {
-  const baseUrl: string = process.env['ACTIONS_CACHE_URL'] || ''
+  // Ideally we just use ACTIONS_CACHE_URL
+  const baseUrl: string = (
+    process.env['ACTIONS_CACHE_URL'] ||
+    process.env['ACTIONS_RUNTIME_URL'] ||
+    ''
+  ).replace('pipelines', 'artifactcache')
   if (!baseUrl) {
     throw new Error('Cache Service Url not found, unable to restore cache.')
   }
@@ -149,7 +154,8 @@ export async function reserveCache(
 
   const reserveCacheRequest: ReserveCacheRequest = {
     key,
-    version
+    version,
+    cacheSize: options?.cacheSize
   }
   const response = await retryTypedResponse('reserveCache', async () =>
     httpClient.postJson<ReserveCacheResponse>(
@@ -157,6 +163,9 @@ export async function reserveCache(
       reserveCacheRequest
     )
   )
+  if(response?.statusCode === 400){
+    return -2
+  }
   return response?.result?.cacheId ?? -1
 }
 
diff --git a/src/internal/cacheUtils.ts b/src/internal/cacheUtils.ts
index 4109610..aefb87b 100644
--- a/src/internal/cacheUtils.ts
+++ b/src/internal/cacheUtils.ts
@@ -123,3 +123,10 @@ export function assertDefined<T>(name: string, value?: T): T {
 
   return value
 }
+
+export function isGhes(): boolean {
+  const ghUrl = new URL(
+      process.env["GITHUB_SERVER_URL"] || "https://github.com"
+  );
+  return ghUrl.hostname.toUpperCase() !== "GITHUB.COM";
+}
diff --git a/src/internal/contracts.d.ts b/src/internal/contracts.d.ts
index 8048476..ca6ea49 100644
--- a/src/internal/contracts.d.ts
+++ b/src/internal/contracts.d.ts
@@ -14,6 +14,7 @@ export interface CommitCacheRequest {
 export interface ReserveCacheRequest {
   key: string
   version?: string
+  cacheSize?: number
 }
 
 export interface ReserveCacheResponse {
@@ -22,4 +23,5 @@ export interface ReserveCacheResponse {
 
 export interface InternalCacheOptions {
   compressionMethod?: CompressionMethod
+  cacheSize?: number
 }