Refine streaming cache upload handling

Author: Copilot

rush-plugins/rush-amazon-s3-build-cache-plugin/src/AmazonS3Client.ts

@@ -74,13 +74,16 @@ const storageRetryOptions: IStorageRetryOptions = {
  * Computes the SHA-256 hash of a file on disk using streaming reads.
  */
 async function _hashFileAsync(filePath: string): Promise<string> {
-  return await new Promise<string>((resolve, reject) => {
-    const hash: crypto.Hash = crypto.createHash(HASH_ALGORITHM);
-    const stream: FileSystemReadStream = FileSystem.createReadStream(filePath);
-    stream.on('data', (chunk: string | Buffer) => hash.update(chunk));
-    stream.on('end', () => resolve(hash.digest('hex')));
-    stream.on('error', reject);
-  });
+  const hash: crypto.Hash = crypto.createHash(HASH_ALGORITHM);
+  const stream: FileSystemReadStream = FileSystem.createReadStream(filePath);
+
+  // If this becomes a hotspot, we can move the hashing work to a worker thread
+  // that reuses a preallocated buffer for the file reads.
+  for await (const chunk of stream) {
+    hash.update(chunk);
+  }
+
+  return hash.digest('hex');
 }
 /**
  * A helper for reading and updating objects on Amazon S3
@@ -223,11 +226,10 @@ export class AmazonS3Client {
       true,
       contentHash
     );
+    response.stream.resume();
     if (!response.ok) {
-      response.stream.resume();
       throw new Error(`Amazon S3 responded with status code ${response.status} (${response.statusText})`);
     }
-    response.stream.resume();
   }
 
   private _writeDebugLine(...messageParts: string[]): void {

rush-plugins/rush-http-build-cache-plugin/src/HttpBuildCacheProvider.ts

@@ -202,12 +202,16 @@ export class HttpBuildCacheProvider implements ICloudBuildCacheProvider {
     }
 
     try {
+      const { size } = await FileSystem.getStatisticsAsync(localFilePath);
       const entryStream: FileSystemReadStream = FileSystem.createReadStream(localFilePath);
       const result: IWebClientStreamResponse | false = await this._makeHttpStreamRequestAsync({
         terminal,
         relUrl: `${this._cacheKeyPrefix}${cacheId}`,
         method: this._uploadMethod,
         body: entryStream,
+        headers: {
+          'Content-Length': `${size}`
+        },
         warningText: 'Could not write cache entry',
         // maxAttempts is 1 because the file read stream is consumed after the first attempt
         // and cannot be replayed. Downloads use MAX_HTTP_CACHE_ATTEMPTS since each retry
@@ -353,6 +357,7 @@ export class HttpBuildCacheProvider implements ICloudBuildCacheProvider {
     relUrl: string;
     method: 'GET' | UploadMethod;
     body: Readable | undefined;
+    headers?: Record<string, string>;
     warningText: string;
     maxAttempts: number;
     credentialOptions?: CredentialsOptions;
@@ -383,12 +388,22 @@ export class HttpBuildCacheProvider implements ICloudBuildCacheProvider {
     relUrl: string;
     method: 'GET' | UploadMethod;
     body: Buffer | Readable | undefined;
+    headers?: Record<string, string>;
     warningText: string;
     maxAttempts: number;
     credentialOptions?: CredentialsOptions;
     stream: boolean;
   }): Promise<IWebClientResponse | IWebClientStreamResponse | false> {
-    const { terminal, relUrl, method, body, warningText, credentialOptions, stream } = options;
+    const {
+      terminal,
+      relUrl,
+      method,
+      body,
+      headers: requestHeaders,
+      warningText,
+      credentialOptions,
+      stream
+    } = options;
     const safeCredentialOptions: CredentialsOptions = credentialOptions ?? CredentialsOptions.Optional;
     const credentials: string | undefined = await this._tryGetCredentialsAsync(safeCredentialOptions);
     const url: string = new URL(relUrl, this._url).href;
@@ -404,6 +419,12 @@ export class HttpBuildCacheProvider implements ICloudBuildCacheProvider {
       }
     }
 
+    if (requestHeaders) {
+      for (const [key, value] of Object.entries(requestHeaders)) {
+        headers[key] = value;
+      }
+    }
+
     const bodyLengthDesc: string = Buffer.isBuffer(body) ? `${body.length} bytes` : 'unknown length';
 
     terminal.writeDebugLine(`[http-build-cache] request: ${method} ${url} ${bodyLengthDesc}`);

rush-plugins/rush-http-build-cache-plugin/src/test/HttpBuildCacheProvider.test.ts

@@ -57,6 +57,9 @@ describe('HttpBuildCacheProvider', () => {
     jest
       .spyOn(FileSystem, 'createReadStream')
       .mockReturnValue({ pipe: jest.fn() } as unknown as ReturnType<typeof FileSystem.createReadStream>);
+    jest
+      .spyOn(FileSystem, 'getStatisticsAsync')
+      .mockResolvedValue({ size: 123 } as Awaited<ReturnType<typeof FileSystem.getStatisticsAsync>>);
     jest
       .spyOn(FileSystem, 'createWriteStreamAsync')
       .mockResolvedValue({} as unknown as Awaited<ReturnType<typeof FileSystem.createWriteStreamAsync>>);
@@ -413,7 +416,10 @@ Array [
       expect(streamFetchFn).toHaveBeenCalledWith(
         'https://buildcache.example.acme.com/some-key',
         expect.objectContaining({
-          method: 'POST'
+          method: 'POST',
+          headers: expect.objectContaining({
+            'Content-Length': '123'
+          })
         })
       );
     });