Merge pull request #1928 from actions/yacaovsnc/artifact_upload_concurrency_and_timeout

Make both upload concurrency and timeout settings configurable with env variables.

Author: yacaovsnc

packages/artifact/__tests__/config.test.ts

@@ -1,4 +1,10 @@
 import * as config from '../src/internal/shared/config'
+import os from 'os'
+
+// Mock the 'os' module
+jest.mock('os', () => ({
+  cpus: jest.fn()
+}))
 
 beforeEach(() => {
   jest.resetModules()
@@ -30,3 +36,66 @@ describe('isGhes', () => {
     expect(config.isGhes()).toBe(true)
   })
 })
+
+describe('uploadChunkTimeoutEnv', () => {
+  it('should return default 300000 when no env set', () => {
+    expect(config.getUploadChunkTimeout()).toBe(300000)
+  })
+
+  it('should return value set in ACTIONS_ARTIFACT_UPLOAD_TIMEOUT_MS', () => {
+    process.env.ACTIONS_ARTIFACT_UPLOAD_TIMEOUT_MS = '150000'
+    expect(config.getUploadChunkTimeout()).toBe(150000)
+  })
+
+  it('should throw if value set in ACTIONS_ARTIFACT_UPLOAD_TIMEOUT_MS is invalid', () => {
+    process.env.ACTIONS_ARTIFACT_UPLOAD_TIMEOUT_MS = 'abc'
+    expect(() => {
+      config.getUploadChunkTimeout()
+    }).toThrow()
+  })
+})
+
+describe('uploadConcurrencyEnv', () => {
+  it('should return default 32 when cpu num is <= 4', () => {
+    ;(os.cpus as jest.Mock).mockReturnValue(new Array(4))
+    expect(config.getConcurrency()).toBe(32)
+  })
+
+  it('should return 16 * num of cpu when cpu num is > 4', () => {
+    ;(os.cpus as jest.Mock).mockReturnValue(new Array(6))
+    expect(config.getConcurrency()).toBe(96)
+  })
+
+  it('should return up to 300 max value', () => {
+    ;(os.cpus as jest.Mock).mockReturnValue(new Array(32))
+    expect(config.getConcurrency()).toBe(300)
+  })
+
+  it('should return override value when ACTIONS_ARTIFACT_UPLOAD_CONCURRENCY is set', () => {
+    ;(os.cpus as jest.Mock).mockReturnValue(new Array(4))
+    process.env.ACTIONS_ARTIFACT_UPLOAD_CONCURRENCY = '10'
+    expect(config.getConcurrency()).toBe(10)
+  })
+
+  it('should throw with invalid value of ACTIONS_ARTIFACT_UPLOAD_CONCURRENCY', () => {
+    ;(os.cpus as jest.Mock).mockReturnValue(new Array(4))
+    process.env.ACTIONS_ARTIFACT_UPLOAD_CONCURRENCY = 'abc'
+    expect(() => {
+      config.getConcurrency()
+    }).toThrow()
+  })
+
+  it('should throw if ACTIONS_ARTIFACT_UPLOAD_CONCURRENCY is < 1', () => {
+    ;(os.cpus as jest.Mock).mockReturnValue(new Array(4))
+    process.env.ACTIONS_ARTIFACT_UPLOAD_CONCURRENCY = '0'
+    expect(() => {
+      config.getConcurrency()
+    }).toThrow()
+  })
+
+  it('cannot go over currency cap when override value is greater', () => {
+    ;(os.cpus as jest.Mock).mockReturnValue(new Array(4))
+    process.env.ACTIONS_ARTIFACT_UPLOAD_CONCURRENCY = '40'
+    expect(config.getConcurrency()).toBe(32)
+  })
+})

packages/artifact/src/internal/shared/config.ts

@@ -1,4 +1,5 @@
 import os from 'os'
+import {info} from '@actions/core'
 
 // Used for controlling the highWaterMark value of the zip that is being streamed
 // The same value is used as the chunk size that is use during upload to blob storage
@@ -47,17 +48,49 @@ export function getGitHubWorkspaceDir(): string {
 // Mimics behavior of azcopy: https://learn.microsoft.com/en-us/azure/storage/common/storage-use-azcopy-optimize
 // If your machine has fewer than 5 CPUs, then the value of this variable is set to 32.
 // Otherwise, the default value is equal to 16 multiplied by the number of CPUs. The maximum value of this variable is 300.
+// This value can be lowered with ACTIONS_ARTIFACT_UPLOAD_CONCURRENCY variable.
 export function getConcurrency(): number {
   const numCPUs = os.cpus().length
+  let concurrencyCap = 32
 
-  if (numCPUs <= 4) {
-    return 32
+  if (numCPUs > 4) {
+    const concurrency = 16 * numCPUs
+    concurrencyCap = concurrency > 300 ? 300 : concurrency
   }
 
-  const concurrency = 16 * numCPUs
-  return concurrency > 300 ? 300 : concurrency
+  const concurrencyOverride = process.env['ACTIONS_ARTIFACT_UPLOAD_CONCURRENCY']
+  if (concurrencyOverride) {
+    const concurrency = parseInt(concurrencyOverride)
+    if (isNaN(concurrency) || concurrency < 1) {
+      throw new Error(
+        'Invalid value set for ACTIONS_ARTIFACT_UPLOAD_CONCURRENCY env variable'
+      )
+    }
+
+    if (concurrency < concurrencyCap) {
+      return concurrency
+    }
+
+    info(
+      `ACTIONS_ARTIFACT_UPLOAD_CONCURRENCY is higher than the cap of ${concurrencyCap} based on the number of cpus. Lowering it to the cap.`
+    )
+  }
+
+  return concurrencyCap
 }
 
 export function getUploadChunkTimeout(): number {
-  return 300_000 // 5 minutes
+  const timeoutVar = process.env['ACTIONS_ARTIFACT_UPLOAD_TIMEOUT_MS']
+  if (!timeoutVar) {
+    return 300000 // 5 minutes
+  }
+
+  const timeout = parseInt(timeoutVar)
+  if (isNaN(timeout)) {
+    throw new Error(
+      'Invalid value set for ACTIONS_ARTIFACT_UPLOAD_TIMEOUT_MS env variable'
+    )
+  }
+
+  return timeout
 }