fix: Add retries to all idempotent endpoints and added parallel file writing (#140)

Author: christianalfoni

src/Sandboxes.ts

@@ -13,7 +13,7 @@ import {
   getStartOptions,
   getStartResponse,
   handleResponse,
-  withCustomTimeout,
+  retryWithDelay,
 } from "./utils/api";
 
 import {
@@ -33,22 +33,24 @@ export async function startVm(
   sandboxId: string,
   startOpts?: StartSandboxOpts
 ): Promise<PitcherManagerResponse> {
-  const startResult = await withCustomTimeout((signal) =>
-    vmStart({
-      client: apiClient,
-      body: startOpts
-        ? {
-            ipcountry: startOpts.ipcountry,
-            tier: startOpts.vmTier?.name,
-            hibernation_timeout_seconds: startOpts.hibernationTimeoutSeconds,
-            automatic_wakeup_config: startOpts.automaticWakeupConfig,
-          }
-        : undefined,
-      path: {
-        id: sandboxId,
-      },
-      signal,
-    })
+  const startResult = await retryWithDelay(
+    () =>
+      vmStart({
+        client: apiClient,
+        body: startOpts
+          ? {
+              ipcountry: startOpts.ipcountry,
+              tier: startOpts.vmTier?.name,
+              hibernation_timeout_seconds: startOpts.hibernationTimeoutSeconds,
+              automatic_wakeup_config: startOpts.automaticWakeupConfig,
+            }
+          : undefined,
+        path: {
+          id: sandboxId,
+        },
+      }),
+    3,
+    200
   );
 
   const response = handleResponse(
@@ -91,7 +93,6 @@ export class Sandboxes {
         description: opts?.description,
         tags: tagsWithSdk,
         path,
-        start_options: getStartOptions(opts),
       },
       path: {
         id: templateId,
@@ -100,11 +101,13 @@ export class Sandboxes {
 
     const sandbox = handleResponse(result, "Failed to create sandbox");
 
-    return new Sandbox(
-      sandbox.id,
-      this.apiClient,
-      getStartResponse(sandbox.start_response)
+    const startResponse = await retryWithDelay(
+      () => startVm(this.apiClient, sandbox.id, getStartOptions(opts)),
+      3,
+      200
     );
+
+    return new Sandbox(sandbox.id, this.apiClient, startResponse);
   }
 
   /**
@@ -125,14 +128,16 @@ export class Sandboxes {
    * Shuts down a sandbox. Files will be saved, and the sandbox will be stopped.
    */
   async shutdown(sandboxId: string): Promise<void> {
-    const response = await withCustomTimeout((signal) =>
-      vmShutdown({
-        client: this.apiClient,
-        path: {
-          id: sandboxId,
-        },
-        signal,
-      })
+    const response = await retryWithDelay(
+      () =>
+        vmShutdown({
+          client: this.apiClient,
+          path: {
+            id: sandboxId,
+          },
+        }),
+      3,
+      200
     );
 
     handleResponse(response, `Failed to shutdown sandbox ${sandboxId}`);
@@ -156,53 +161,39 @@ export class Sandboxes {
    * Will resolve once the sandbox is restarted with its setup running.
    */
   public async restart(sandboxId: string, opts?: StartSandboxOpts) {
-    let didRestart = false;
-
-    for (let attempt = 1; attempt <= 3; attempt++) {
-      try {
-        await this.shutdown(sandboxId);
-        didRestart = true;
-        break;
-      } catch (e) {
-        await sleep(500);
-      }
+    try {
+      await this.shutdown(sandboxId);
+    } catch (e) {
+      throw new Error("Failed to shutdown VM, " + String(e));
     }
 
-    if (!didRestart) {
-      throw new Error("Failed to shutdown VM after 3 attempts");
-    }
-
-    let startResponse: PitcherManagerResponse | undefined;
-
-    for (let attempt = 1; attempt <= 3; attempt++) {
-      try {
-        startResponse = await startVm(this.apiClient, sandboxId, opts);
-        break;
-      } catch (e) {
-        await sleep(500);
-      }
-    }
+    try {
+      const startResponse = await retryWithDelay(
+        () => startVm(this.apiClient, sandboxId, opts),
+        3,
+        200
+      );
 
-    if (!startResponse) {
-      throw new Error("Failed to start VM after 3 attempts");
+      return new Sandbox(sandboxId, this.apiClient, startResponse);
+    } catch (e) {
+      throw new Error("Failed to start VM, " + String(e));
     }
-
-    return new Sandbox(sandboxId, this.apiClient, startResponse);
   }
-
   /**
    * Hibernates a sandbox. Files will be saved, and the sandbox will be put to sleep. Next time
    * you resume the sandbox it will continue from the last state it was in.
    */
   async hibernate(sandboxId: string): Promise<void> {
-    const response = await withCustomTimeout((signal) =>
-      vmHibernate({
-        client: this.apiClient,
-        path: {
-          id: sandboxId,
-        },
-        signal,
-      })
+    const response = await retryWithDelay(
+      () =>
+        vmHibernate({
+          client: this.apiClient,
+          path: {
+            id: sandboxId,
+          },
+        }),
+      3,
+      200
     );
 
     handleResponse(response, `Failed to hibernate sandbox ${sandboxId}`);

src/bin/commands/build.ts

@@ -16,11 +16,13 @@ import {
   createApiClient,
   getDefaultTemplateId,
   handleResponse,
+  retryWithDelay,
 } from "../../utils/api";
 import { getInferredApiKey } from "../../utils/constants";
 import { hashDirectory as getFilePaths } from "../utils/files";
 import { startVm } from "../../Sandboxes";
 import { mkdir, writeFile } from "fs/promises";
+import { sleep } from "../../utils/sleep";
 
 export type BuildCommandArgs = {
   directory: string;
@@ -196,49 +198,58 @@ export const buildCommand: yargs.CommandModule<
         try {
           spinner.start(updateSpinnerMessage(index, "Starting sandbox..."));
 
-          const startResponse = await withCustomError(
-            startVm(apiClient, id),
-            "Failed to start sandbox"
+          const startResponse = await retryWithDelay(() =>
+            withCustomError(
+              startVm(apiClient, id),
+              "Failed to start sandbox at all"
+            )
           );
           let sandboxVM = new Sandbox(id, apiClient, startResponse);
 
-          let session = await sandboxVM.connect();
+          let session = await retryWithDelay(() => sandboxVM.connect(), 3, 100);
 
           spinner.start(
             updateSpinnerMessage(index, "Writing files to sandbox...")
           );
 
-          let i = 0;
-          for (const filePath of filePaths) {
-            i++;
-            try {
-              const fullPath = path.join(argv.directory, filePath);
-              const content = await fs.readFile(fullPath);
-              const dirname = path.dirname(filePath);
-              await session.fs.mkdir(dirname, true);
-              await session.fs.writeFile(filePath, content, {
-                create: true,
-                overwrite: true,
-              });
-            } catch (error) {
-              throw new Error(
-                `Failed to write "${filePath}" to sandbox: ${error}`
-              );
-            }
-          }
+          await Promise.all(
+            filePaths.map((filePath) =>
+              retryWithDelay(
+                async () => {
+                  const fullPath = path.join(argv.directory, filePath);
+                  const content = await fs.readFile(fullPath);
+                  const dirname = path.dirname(filePath);
+                  await session.fs.mkdir(dirname, true);
+                  await session.fs.writeFile(filePath, content, {
+                    create: true,
+                    overwrite: true,
+                  });
+                },
+                3,
+                200
+              )
+            )
+          ).catch((error) => {
+            throw new Error(`Failed to write files to sandbox: ${error}`);
+          });
 
           spinner.start(updateSpinnerMessage(index, "Building sandbox..."));
 
           sandboxVM = await withCustomError(
             sdk.sandboxes.restart(id, {
               vmTier: buildTier,
             }),
-            "Failed to restart sandbox"
+            "Failed to restart sandbox after building"
           );
 
-          session = await withCustomError(
-            sandboxVM.connect(),
-            "Failed to connect to sandbox"
+          session = await retryWithDelay(
+            () =>
+              withCustomError(
+                sandboxVM.connect(),
+                "Failed to connect to sandbox after building"
+              ),
+            3,
+            100
           );
 
           await waitForSetup(session, index);
@@ -250,12 +261,17 @@ export const buildCommand: yargs.CommandModule<
             sdk.sandboxes.restart(id, {
               vmTier: sandboxTier,
             }),
-            "Failed to restart sandbox"
+            "Failed to restart sandbox after optimizing initial state"
           );
 
-          session = await withCustomError(
-            sandboxVM.connect(),
-            "Failed to connect to sandbox"
+          session = await retryWithDelay(
+            () =>
+              withCustomError(
+                sandboxVM.connect(),
+                "Failed to connect to sandbox after optimizing initial state"
+              ),
+            3,
+            100
           );
 
           await waitForSetup(session, index);
@@ -289,7 +305,7 @@ export const buildCommand: yargs.CommandModule<
                     break;
                   }
 
-                  await new Promise((resolve) => setTimeout(resolve, 1000));
+                  await sleep(1000);
                 }
 
                 updatePortSpinner();
@@ -308,7 +324,7 @@ export const buildCommand: yargs.CommandModule<
           spinner.start(updateSpinnerMessage(index, "Creating snapshot..."));
           await withCustomError(
             sdk.sandboxes.hibernate(id),
-            "Failed to hibernate"
+            "Failed to hibernate after building and optimizing sandbox"
           );
           spinner.start(updateSpinnerMessage(index, "Snapshot created"));
 

src/utils/api.ts

@@ -98,35 +98,6 @@ export function getStartResponse(
   };
 }
 
-/**
- * Our infra has 2 min timeout, so we use that as default
- */
-export async function withCustomTimeout<T>(
-  cb: (signal: AbortSignal) => Promise<T>,
-  timeoutSeconds: number = 120
-) {
-  const controller = new AbortController();
-  const signal = controller.signal;
-  const timeoutHandle = setTimeout(() => {
-    controller.abort();
-  }, timeoutSeconds * 1000);
-
-  try {
-    // We have to await for the finally to run
-    return await cb(signal);
-  } catch (err) {
-    if (err instanceof Error && err.name === "AbortError") {
-      throw new Error(
-        `Request took longer than ${timeoutSeconds}s, so we aborted.`
-      );
-    }
-
-    throw err;
-  } finally {
-    clearTimeout(timeoutHandle);
-  }
-}
-
 export function getDefaultTemplateTag(apiClient: Client): string {
   if (apiClient.getConfig().baseUrl?.includes("codesandbox.stream")) {
     return "7ngcrf";
@@ -145,6 +116,30 @@ export function getDefaultTemplateId(apiClient: Client): string {
   return "pcz35m";
 }
 
+export async function retryWithDelay<T>(
+  callback: () => Promise<T>,
+  retries: number = 3,
+  delay: number = 500
+): Promise<T> {
+  let lastError: Error;
+
+  for (let attempt = 1; attempt <= retries; attempt++) {
+    try {
+      return await callback();
+    } catch (error) {
+      lastError = error as Error;
+
+      if (attempt === retries) {
+        throw lastError;
+      }
+
+      await new Promise((resolve) => setTimeout(resolve, delay));
+    }
+  }
+
+  throw lastError!;
+}
+
 export function handleResponse<D, E>(
   result: Awaited<{ data?: { data?: D }; error?: E; response: Response }>,
   errorPrefix: string