From c6f1224d30a062385d872421e3a5c0efab0923e7 Mon Sep 17 00:00:00 2001
From: Bassem Dghaidi <568794+Link-@users.noreply.github.com>
Date: Mon, 2 Dec 2024 02:33:27 -0800
Subject: [PATCH] Add progress tracking for blob uploads

---
 packages/cache/__tests__/uploadUtils.test.ts |  58 +++++++
 packages/cache/src/internal/uploadUtils.ts   | 155 +++++++++++++++++--
 2 files changed, 199 insertions(+), 14 deletions(-)
 create mode 100644 packages/cache/__tests__/uploadUtils.test.ts

diff --git a/packages/cache/__tests__/uploadUtils.test.ts b/packages/cache/__tests__/uploadUtils.test.ts
new file mode 100644
index 00000000..6a4876d1
--- /dev/null
+++ b/packages/cache/__tests__/uploadUtils.test.ts
@@ -0,0 +1,58 @@
+import {UploadProgress} from '../src/internal/uploadUtils'
+import {TransferProgressEvent} from '@azure/ms-rest-js'
+
+test('upload progress tracked correctly', () => {
+  const progress = new UploadProgress(1000)
+
+  expect(progress.contentLength).toBe(1000)
+  expect(progress.sentBytes).toBe(0)
+  expect(progress.displayedComplete).toBe(false)
+  expect(progress.timeoutHandle).toBeUndefined()
+  expect(progress.getTransferredBytes()).toBe(0)
+  expect(progress.isDone()).toBe(false)
+
+  progress.onProgress()({loadedBytes: 0} as TransferProgressEvent)
+
+  expect(progress.contentLength).toBe(1000)
+  expect(progress.sentBytes).toBe(0)
+  expect(progress.displayedComplete).toBe(false)
+  expect(progress.timeoutHandle).toBeUndefined()
+  expect(progress.getTransferredBytes()).toBe(0)
+  expect(progress.isDone()).toBe(false)
+
+  progress.onProgress()({loadedBytes: 250} as TransferProgressEvent)
+
+  expect(progress.contentLength).toBe(1000)
+  expect(progress.sentBytes).toBe(250)
+  expect(progress.displayedComplete).toBe(false)
+  expect(progress.timeoutHandle).toBeUndefined()
+  expect(progress.getTransferredBytes()).toBe(250)
+  expect(progress.isDone()).toBe(false)
+
+  progress.onProgress()({loadedBytes: 500} as TransferProgressEvent)
+
+  expect(progress.contentLength).toBe(1000)
+  expect(progress.sentBytes).toBe(500)
+  expect(progress.displayedComplete).toBe(false)
+  expect(progress.timeoutHandle).toBeUndefined()
+  expect(progress.getTransferredBytes()).toBe(500)
+  expect(progress.isDone()).toBe(false)
+
+  progress.onProgress()({loadedBytes: 750} as TransferProgressEvent)
+
+  expect(progress.contentLength).toBe(1000)
+  expect(progress.sentBytes).toBe(750)
+  expect(progress.displayedComplete).toBe(false)
+  expect(progress.timeoutHandle).toBeUndefined()
+  expect(progress.getTransferredBytes()).toBe(750)
+  expect(progress.isDone()).toBe(false)
+
+  progress.onProgress()({loadedBytes: 1000} as TransferProgressEvent)
+
+  expect(progress.contentLength).toBe(1000)
+  expect(progress.sentBytes).toBe(1000)
+  expect(progress.displayedComplete).toBe(false)
+  expect(progress.timeoutHandle).toBeUndefined()
+  expect(progress.getTransferredBytes()).toBe(1000)
+  expect(progress.isDone()).toBe(true)
+})
diff --git a/packages/cache/src/internal/uploadUtils.ts b/packages/cache/src/internal/uploadUtils.ts
index a3376d9d..5ba98f91 100644
--- a/packages/cache/src/internal/uploadUtils.ts
+++ b/packages/cache/src/internal/uploadUtils.ts
@@ -5,35 +5,162 @@ import {
   BlockBlobClient,
   BlockBlobParallelUploadOptions
 } from '@azure/storage-blob'
+import {TransferProgressEvent} from '@azure/ms-rest-js'
 import {InvalidResponseError} from './shared/errors'
 import {UploadOptions} from '../options'
 
+/**
+ * Class for tracking the upload state and displaying stats.
+ */
+export class UploadProgress {
+  contentLength: number
+  sentBytes: number
+  startTime: number
+  displayedComplete: boolean
+  timeoutHandle?: ReturnType<typeof setTimeout>
+
+  constructor(contentLength: number) {
+    this.contentLength = contentLength
+    this.sentBytes = 0
+    this.displayedComplete = false
+    this.startTime = Date.now()
+  }
+
+  /**
+   * Sets the number of bytes sent
+   *
+   * @param sentBytes the number of bytes sent
+   */
+  setSentBytes(sentBytes: number): void {
+    this.sentBytes = sentBytes
+  }
+
+  /**
+   * Returns the total number of bytes transferred.
+   */
+  getTransferredBytes(): number {
+    return this.sentBytes
+  }
+
+  /**
+   * Returns true if the upload is complete.
+   */
+  isDone(): boolean {
+    return this.getTransferredBytes() === this.contentLength
+  }
+
+  /**
+   * Prints the current upload stats. Once the upload completes, this will print one
+   * last line and then stop.
+   */
+  display(): void {
+    if (this.displayedComplete) {
+      return
+    }
+
+    const transferredBytes = this.sentBytes
+    const percentage = (100 * (transferredBytes / this.contentLength)).toFixed(
+      1
+    )
+    const elapsedTime = Date.now() - this.startTime
+    const uploadSpeed = (
+      transferredBytes /
+      (1024 * 1024) /
+      (elapsedTime / 1000)
+    ).toFixed(1)
+
+    core.info(
+      `Sent ${transferredBytes} of ${this.contentLength} (${percentage}%), ${uploadSpeed} MBs/sec`
+    )
+
+    if (this.isDone()) {
+      this.displayedComplete = true
+    }
+  }
+
+  /**
+   * Returns a function used to handle TransferProgressEvents.
+   */
+  onProgress(): (progress: TransferProgressEvent) => void {
+    return (progress: TransferProgressEvent) => {
+      this.setSentBytes(progress.loadedBytes)
+    }
+  }
+
+  /**
+   * Starts the timer that displays the stats.
+   *
+   * @param delayInMs the delay between each write
+   */
+  startDisplayTimer(delayInMs = 1000): void {
+    const displayCallback = (): void => {
+      this.display()
+
+      if (!this.isDone()) {
+        this.timeoutHandle = setTimeout(displayCallback, delayInMs)
+      }
+    }
+
+    this.timeoutHandle = setTimeout(displayCallback, delayInMs)
+  }
+
+  /**
+   * Stops the timer that displays the stats. As this typically indicates the upload
+   * is complete, this will display one last line, unless the last line has already
+   * been written.
+   */
+  stopDisplayTimer(): void {
+    if (this.timeoutHandle) {
+      clearTimeout(this.timeoutHandle)
+      this.timeoutHandle = undefined
+    }
+
+    this.display()
+  }
+}
+
 export async function uploadCacheArchiveSDK(
   signedUploadURL: string,
   archivePath: string,
   options?: UploadOptions
 ): Promise<BlobUploadCommonResponse> {
+  const blobClient: BlobClient = new BlobClient(signedUploadURL)
+  const blockBlobClient: BlockBlobClient = blobClient.getBlockBlobClient()
+
+  const properties = await blobClient.getProperties()
+  const contentLength = properties.contentLength ?? -1
+
+  const uploadProgress = new UploadProgress(contentLength)
+
   // Specify data transfer options
   const uploadOptions: BlockBlobParallelUploadOptions = {
     blockSize: options?.uploadChunkSize,
     concurrency: options?.uploadConcurrency, // maximum number of parallel transfer workers
-    maxSingleShotSize: 128 * 1024 * 1024 // 128 MiB initial transfer size
+    maxSingleShotSize: 128 * 1024 * 1024, // 128 MiB initial transfer size
+    onProgress: uploadProgress.onProgress()
   }
 
-  const blobClient: BlobClient = new BlobClient(signedUploadURL)
-  const blockBlobClient: BlockBlobClient = blobClient.getBlockBlobClient()
+  try {
+    uploadProgress.startDisplayTimer()
 
-  core.debug(
-    `BlobClient: ${blobClient.name}:${blobClient.accountName}:${blobClient.containerName}`
-  )
-
-  const resp = await blockBlobClient.uploadFile(archivePath, uploadOptions)
-
-  if (resp._response.status >= 400) {
-    throw new InvalidResponseError(
-      `Upload failed with status code ${resp._response.status}`
+    core.debug(
+      `BlobClient: ${blobClient.name}:${blobClient.accountName}:${blobClient.containerName}`
     )
-  }
 
-  return resp
+    const response = await blockBlobClient.uploadFile(
+      archivePath,
+      uploadOptions
+    )
+
+    // TODO: better management of non-retryable errors
+    if (response._response.status >= 400) {
+      throw new InvalidResponseError(
+        `Upload failed with status code ${response._response.status}`
+      )
+    }
+
+    return response
+  } finally {
+    uploadProgress.stopDisplayTimer()
+  }
 }