Merge pull request #878 from chromaui/ghengeveld/ap-3912-index-update…

…-cli-to-wait-for-copy-job-using-sentinel-file

Retrieve `sentinelUrls` from `uploadBuild` and wait for all of them before finishing upload task

node-src/index.test.ts

@@ -207,6 +207,7 @@ vi.mock('node-fetch', () => ({
           data: {
             [key]: {
               info: {
+                sentinelUrls: [],
                 targets: variables.files.map(({ filePath }) => ({
                   contentType: contentTypes[filePath.split('.').at(-1)],
                   fileKey: '',
@@ -469,6 +470,7 @@ it('calls out to npm build script passed and uploads files', async () => {
     expect.any(Object),
     [
       {
+        contentHash: 'hash',
         contentLength: 42,
         contentType: 'text/html',
         fileKey: '',
@@ -479,6 +481,7 @@ it('calls out to npm build script passed and uploads files', async () => {
         targetPath: 'iframe.html',
       },
       {
+        contentHash: 'hash',
         contentLength: 42,
         contentType: 'text/html',
         fileKey: '',
@@ -502,6 +505,7 @@ it('skips building and uploads directly with storybook-build-dir', async () => {
     expect.any(Object),
     [
       {
+        contentHash: 'hash',
         contentLength: 42,
         contentType: 'text/html',
         fileKey: '',
@@ -512,6 +516,7 @@ it('skips building and uploads directly with storybook-build-dir', async () => {
         targetPath: 'iframe.html',
       },
       {
+        contentHash: 'hash',
         contentLength: 42,
         contentType: 'text/html',
         fileKey: '',

node-src/lib/FileReaderBlob.ts

@@ -0,0 +1,20 @@
+import { ReadStream, createReadStream } from 'fs';
+
+export class FileReaderBlob {
+  readStream: ReadStream;
+  size: number;
+
+  constructor(filePath: string, contentLength: number, onProgress: (delta: number) => void) {
+    this.size = contentLength;
+    this.readStream = createReadStream(filePath);
+    this.readStream.on('data', (chunk: Buffer | string) => onProgress(chunk.length));
+  }
+
+  stream() {
+    return this.readStream;
+  }
+
+  get [Symbol.toStringTag]() {
+    return 'Blob';
+  }
+}

node-src/lib/upload.ts

@@ -1,7 +1,7 @@
 import makeZipFile from './compress';
-import { Context, FileDesc, TargetInfo } from '../types';
-import { uploadZip, waitForUnpack } from './uploadZip';
+import { uploadZip } from './uploadZip';
 import { uploadFiles } from './uploadFiles';
+import { Context, FileDesc, TargetInfo } from '../types';
 import { maxFileCountExceeded } from '../ui/messages/errors/maxFileCountExceeded';
 import { maxFileSizeExceeded } from '../ui/messages/errors/maxFileSizeExceeded';
 import { skippingEmptyFiles } from '../ui/messages/warnings/skippingEmptyFiles';
@@ -13,6 +13,7 @@ const UploadBuildMutation = `
   mutation UploadBuildMutation($buildId: ObjID!, $files: [FileUploadInput!]!, $zip: Boolean) {
     uploadBuild(buildId: $buildId, files: $files, zip: $zip) {
       info {
+        sentinelUrls
         targets {
           contentType
           fileKey
@@ -26,7 +27,6 @@ const UploadBuildMutation = `
           filePath
           formAction
           formFields
-          sentinelUrl
         }
       }
       userErrors {
@@ -50,8 +50,9 @@ const UploadBuildMutation = `
 interface UploadBuildMutationResult {
   uploadBuild: {
     info?: {
+      sentinelUrls: string[];
       targets: TargetInfo[];
-      zipTarget?: TargetInfo & { sentinelUrl: string };
+      zipTarget?: TargetInfo;
     };
     userErrors: (
       | {
@@ -80,15 +81,16 @@ export async function uploadBuild(
   options: {
     onStart?: () => void;
     onProgress?: (progress: number, total: number) => void;
-    onComplete?: (uploadedBytes: number, uploadedFiles: number) => void;
+    onComplete?: (uploadedBytes: number, uploadedFiles: number, sentinelUrls: string[]) => void;
     onError?: (error: Error, path?: string) => void;
   } = {}
 ) {
+  ctx.sentinelUrls = [];
   ctx.uploadedBytes = 0;
   ctx.uploadedFiles = 0;
 
   const targets: (TargetInfo & FileDesc)[] = [];
-  let zipTarget: (TargetInfo & { sentinelUrl: string }) | undefined;
+  let zipTarget: TargetInfo | undefined;
 
   const batches = files.reduce<typeof files[]>((acc, file, fileIndex) => {
     const batchIndex = Math.floor(fileIndex / MAX_FILES_PER_REQUEST);
@@ -128,6 +130,7 @@ export async function uploadBuild(
       return options.onError?.(new Error('Upload rejected due to user error'));
     }
 
+    ctx.sentinelUrls.push(...uploadBuild.info.sentinelUrls);
     targets.push(
       ...uploadBuild.info.targets.map((target) => {
         const file = batch.find((f) => f.targetPath === target.filePath);
@@ -153,7 +156,6 @@ export async function uploadBuild(
 
       const target = { ...zipTarget, contentLength: size, localPath: path };
       await uploadZip(ctx, target, (progress) => options.onProgress?.(progress, size));
-      await waitForUnpack(ctx, target.sentinelUrl);
       ctx.uploadedBytes += size;
       ctx.uploadedFiles += targets.length;
       return;
@@ -213,7 +215,8 @@ export async function uploadMetadata(ctx: Context, files: FileDesc[]) {
     UploadMetadataMutation,
     {
       buildId: ctx.announcedBuild.id,
-      files: files.map(({ contentLength, targetPath }) => ({
+      files: files.map(({ contentHash, contentLength, targetPath }) => ({
+        contentHash,
         contentLength,
         filePath: targetPath,
       })),

node-src/lib/uploadFiles.ts

@@ -1,10 +1,9 @@
 import retry from 'async-retry';
 import { filesize } from 'filesize';
-import FormData from 'form-data';
-import { createReadStream } from 'fs';
+import { FormData } from 'formdata-node';
 import pLimit from 'p-limit';
-import progress from 'progress-stream';
 import { Context, FileDesc, TargetInfo } from '../types';
+import { FileReaderBlob } from './FileReaderBlob';
 
 export async function uploadFiles(
   ctx: Context,
@@ -28,19 +27,15 @@ export async function uploadFiles(
               return bail(signal.reason || new Error('Aborted'));
             }
 
-            const progressStream = progress();
-
-            progressStream.on('progress', ({ delta }) => {
-              fileProgress += delta; // We upload multiple files so we only care about the delta
+            const blob = new FileReaderBlob(localPath, contentLength, (delta) => {
+              fileProgress += delta;
               totalProgress += delta;
               onProgress?.(totalProgress);
             });
 
             const formData = new FormData();
             Object.entries(formFields).forEach(([k, v]) => formData.append(k, v));
-            formData.append('file', createReadStream(localPath).pipe(progressStream), {
-              knownLength: contentLength,
-            });
+            formData.append('file', blob);
 
             const res = await ctx.http.fetch(
               formAction,

node-src/lib/uploadZip.ts

@@ -1,19 +1,12 @@
 import retry from 'async-retry';
 import { filesize } from 'filesize';
-import FormData from 'form-data';
-import { createReadStream } from 'fs';
-import { Response } from 'node-fetch';
-import progress from 'progress-stream';
+import { FormData } from 'formdata-node';
 import { Context, TargetInfo } from '../types';
-
-// A sentinel file is created by a zip-unpack lambda within the Chromatic infrastructure once the
-// uploaded zip is fully extracted. The contents of this file will consist of 'OK' if the process
-// completed successfully and 'ERROR' if an error occurred.
-const SENTINEL_SUCCESS_VALUE = 'OK';
+import { FileReaderBlob } from './FileReaderBlob';
 
 export async function uploadZip(
   ctx: Context,
-  target: TargetInfo & { contentLength: number; localPath: string; sentinelUrl: string },
+  target: TargetInfo & { contentLength: number; localPath: string },
   onProgress: (progress: number) => void
 ) {
   const { experimental_abortSignal: signal } = ctx.options;
@@ -28,18 +21,14 @@ export async function uploadZip(
         return bail(signal.reason || new Error('Aborted'));
       }
 
-      const progressStream = progress();
-
-      progressStream.on('progress', ({ delta }) => {
+      const blob = new FileReaderBlob(localPath, contentLength, (delta) => {
         totalProgress += delta;
-        onProgress(totalProgress);
+        onProgress?.(totalProgress);
       });
 
       const formData = new FormData();
       Object.entries(formFields).forEach(([k, v]) => formData.append(k, v));
-      formData.append('file', createReadStream(localPath).pipe(progressStream), {
-        knownLength: contentLength,
-      });
+      formData.append('file', blob);
 
       const res = await ctx.http.fetch(
         formAction,
@@ -63,40 +52,3 @@ export async function uploadZip(
     }
   );
 }
-
-export async function waitForUnpack(ctx: Context, url: string) {
-  const { experimental_abortSignal: signal } = ctx.options;
-
-  ctx.log.debug(`Waiting for zip unpack sentinel file to appear at '${url}'`);
-
-  return retry(
-    async (bail) => {
-      if (signal?.aborted) {
-        return bail(signal.reason || new Error('Aborted'));
-      }
-
-      let res: Response;
-      try {
-        res = await ctx.http.fetch(url, { signal }, { retries: 0, noLogErrorBody: true });
-      } catch (e) {
-        const { response = {} } = e;
-        if (response.status === 403) {
-          return bail(new Error('Provided signature expired.'));
-        }
-        throw new Error('Sentinel file not present.');
-      }
-
-      const result = await res.text();
-      if (result !== SENTINEL_SUCCESS_VALUE) {
-        return bail(new Error('Zip file failed to unpack remotely.'));
-      } else {
-        ctx.log.debug(`Sentinel file present, continuing.`);
-      }
-    },
-    {
-      retries: 185, // 3 minutes and some change (matches the lambda timeout with some extra buffer)
-      minTimeout: 1000,
-      maxTimeout: 1000,
-    }
-  );
-}

node-src/lib/utils.ts

@@ -34,8 +34,6 @@ export const activityBar = (n = 0, size = 20) => {
   return `[${track.join('')}]`;
 };
 
-export const baseStorybookUrl = (url: string) => url?.replace(/\/iframe\.html$/, '');
-
 export const rewriteErrorMessage = (err: Error, message: string) => {
   try {
     // DOMException doesn't allow setting the message, so this might fail

node-src/lib/waitForSentinel.ts

@@ -0,0 +1,45 @@
+import retry from 'async-retry';
+import { Response } from 'node-fetch';
+import { Context } from '../types';
+
+// A sentinel file is created by a zip-unpack lambda within the Chromatic infrastructure once the
+// uploaded zip is fully extracted. The contents of this file will consist of 'OK' if the process
+// completed successfully and 'ERROR' if an error occurred.
+const SENTINEL_SUCCESS_VALUE = 'OK';
+
+export async function waitForSentinel(ctx: Context, url: string) {
+  const { experimental_abortSignal: signal } = ctx.options;
+
+  ctx.log.debug(`Waiting for sentinel file to appear at ${url}`);
+
+  return retry(
+    async (bail) => {
+      if (signal?.aborted) {
+        return bail(signal.reason || new Error('Aborted'));
+      }
+
+      let res: Response;
+      try {
+        res = await ctx.http.fetch(url, { signal }, { retries: 0, noLogErrorBody: true });
+      } catch (e) {
+        const { response = {} } = e;
+        if (response.status === 403) {
+          return bail(new Error('Provided signature expired.'));
+        }
+        throw new Error('Sentinel file not present.');
+      }
+
+      const result = await res.text();
+      if (result !== SENTINEL_SUCCESS_VALUE) {
+        ctx.log.debug(`Sentinel file not OK, got ${result}`);
+        return bail(new Error('Sentinel file error.'));
+      }
+      ctx.log.debug(`Sentinel file OK.`);
+    },
+    {
+      retries: 185, // 3 minutes and some change (matches the lambda timeout with some extra buffer)
+      minTimeout: 1000,
+      maxTimeout: 1000,
+    }
+  );
+}