Retrieve sentinelUrls from uploadBuild and wait for all of them befor…

…e finishing upload task

node-src/lib/upload.ts

@@ -1,14 +1,15 @@
 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';
 
 const UploadBuildMutation = `
   mutation UploadBuildMutation($buildId: ObjID!, $files: [FileUploadInput!]!, $zip: Boolean) {
     uploadBuild(buildId: $buildId, files: $files, zip: $zip) {
       info {
+        sentinelUrls
         targets {
           contentType
           fileKey
@@ -22,7 +23,6 @@ const UploadBuildMutation = `
           filePath
           formAction
           formFields
-          sentinelUrl
         }
       }
       userErrors {
@@ -46,8 +46,9 @@ const UploadBuildMutation = `
 interface UploadBuildMutationResult {
   uploadBuild: {
     info?: {
+      sentinelUrls: string[];
       targets: TargetInfo[];
-      zipTarget?: TargetInfo & { sentinelUrl: string };
+      zipTarget?: TargetInfo;
     };
     userErrors: (
       | {
@@ -76,7 +77,7 @@ 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;
   } = {}
 ) {
@@ -106,14 +107,16 @@ export async function uploadBuild(
     return options.onError?.(new Error('Upload rejected due to user error'));
   }
 
+  const { sentinelUrls } = uploadBuild.info;
+
   const targets = uploadBuild.info.targets.map((target) => {
     const file = files.find((f) => f.targetPath === target.filePath);
     return { ...file, ...target };
   });
 
   if (!targets.length) {
     ctx.log.debug('No new files to upload, continuing');
-    return options.onComplete?.(0, 0);
+    return options.onComplete?.(0, 0, sentinelUrls);
   }
 
   options.onStart?.();
@@ -127,16 +130,15 @@ export async function uploadBuild(
 
       const target = { ...uploadBuild.info.zipTarget, contentLength: size, localPath: path };
       await uploadZip(ctx, target, (progress) => options.onProgress?.(progress, size));
-      await waitForUnpack(ctx, target.sentinelUrl);
-      return options.onComplete?.(size, targets.length);
+      return options.onComplete?.(size, targets.length, sentinelUrls);
     } catch (err) {
       ctx.log.debug({ err }, 'Error uploading zip, falling back to uploading individual files');
     }
   }
 
   try {
     await uploadFiles(ctx, targets, (progress) => options.onProgress?.(progress, total));
-    return options.onComplete?.(total, targets.length);
+    return options.onComplete?.(total, targets.length, sentinelUrls);
   } catch (e) {
     return options.onError?.(e, files.some((f) => f.localPath === e.message) && e.message);
   }

node-src/lib/uploadZip.ts

@@ -1,18 +1,12 @@
 import retry from 'async-retry';
 import { filesize } from 'filesize';
 import { FormData } from 'formdata-node';
-import { Response } from 'node-fetch';
 import { Context, TargetInfo } from '../types';
 import { FileReaderBlob } from './FileReaderBlob';
 
-// 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 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;
@@ -58,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/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,
+    }
+  );
+}

node-src/tasks/upload.ts

@@ -21,6 +21,7 @@ import {
   uploading,
   success,
   hashing,
+  finalizing,
 } from '../ui/tasks/upload';
 import { Context, FileDesc, Task } from '../types';
 import { readStatsFile } from './read-stats-file';
@@ -29,6 +30,7 @@ import { findChangedPackageFiles } from '../lib/findChangedPackageFiles';
 import { findChangedDependencies } from '../lib/findChangedDependencies';
 import { uploadBuild } from '../lib/upload';
 import { getFileHashes } from '../lib/getFileHashes';
+import { waitForSentinel } from '../lib/waitForSentinel';
 
 interface PathSpec {
   pathname: string;
@@ -225,7 +227,8 @@ export const uploadStorybook = async (ctx: Context, task: Task) => {
       // Avoid spamming the logs with progress updates in non-interactive mode
       ctx.options.interactive ? 100 : ctx.env.CHROMATIC_OUTPUT_INTERVAL
     ),
-    onComplete: (uploadedBytes: number, uploadedFiles: number) => {
+    onComplete: (uploadedBytes: number, uploadedFiles: number, sentinelUrls: string[]) => {
+      ctx.sentinelUrls = sentinelUrls;
       ctx.uploadedBytes = uploadedBytes;
       ctx.uploadedFiles = uploadedFiles;
     },
@@ -235,6 +238,13 @@ export const uploadStorybook = async (ctx: Context, task: Task) => {
   });
 };
 
+export const waitForSentinels = async (ctx: Context, task: Task) => {
+  if (ctx.skip || !ctx.sentinelUrls?.length) return;
+  transitionTo(finalizing)(ctx, task);
+
+  await Promise.all(ctx.sentinelUrls.map((url) => waitForSentinel(ctx, url)));
+};
+
 export default createTask({
   name: 'upload',
   title: initial.title,
@@ -249,6 +259,7 @@ export default createTask({
     traceChangedFiles,
     calculateFileHashes,
     uploadStorybook,
+    waitForSentinels,
     transitionTo(success, true),
   ],
 });

node-src/types.ts

@@ -302,6 +302,7 @@ export interface Context {
     }[];
     total: number;
   };
+  sentinelUrls?: string[];
   uploadedBytes?: number;
   uploadedFiles?: number;
   turboSnap?: Partial<{

node-src/ui/tasks/upload.stories.ts

@@ -3,6 +3,7 @@ import {
   bailed,
   dryRun,
   failed,
+  finalizing,
   hashing,
   initial,
   invalid,
@@ -57,6 +58,8 @@ export const Starting = () => starting();
 
 export const Uploading = () => uploading({ percentage: 42 });
 
+export const Finalizing = () => finalizing();
+
 export const Success = () =>
   success({
     now: 0,

node-src/ui/tasks/upload.ts

@@ -99,6 +99,12 @@ export const uploading = ({ percentage }: { percentage: number }) => ({
   output: `${progressBar(percentage)} ${percentage}%`,
 });
 
+export const finalizing = () => ({
+  status: 'pending',
+  title: 'Publishing your built Storybook',
+  output: `Finalizing upload`,
+});
+
 export const success = (ctx: Context) => {
   const files = pluralize('file', ctx.uploadedFiles, true);
   const bytes = filesize(ctx.uploadedBytes || 0);