backoff with helper

2024-06-05 08:33:31 +01:00
parent 141ca4dd0b
commit 4541f2e004
2 changed files with 350 additions and 69 deletions
@@ -0,0 +1,241 @@
+type ExponentialBackoffType = "NoJitter" | "FullJitter" | "EqualJitter";
+
+type ExponentialBackoffOptions = {
+  base: number;
+  factor: number;
+  min: number;
+  max: number;
+  maxRetries: number;
+  maxElapsed: number;
+};
+
+export class StopRetrying extends Error {
+  constructor(message?: string) {
+    super(message);
+    this.name = "StopRetrying";
+  }
+}
+
+export class ExponentialBackoff {
+  #retries: number = 0;
+
+  #type: ExponentialBackoffType;
+  #base: number;
+  #factor: number;
+
+  #min: number;
+  #max: number;
+
+  #maxRetries: number;
+  #maxElapsed: number;
+
+  constructor(type?: ExponentialBackoffType, opts: Partial<ExponentialBackoffOptions> = {}) {
+    this.#type = type ?? "NoJitter";
+    this.#base = opts.base ?? 2;
+    this.#factor = opts.factor ?? 1;
+
+    this.#min = opts.min ?? -Infinity;
+    this.#max = opts.max ?? Infinity;
+
+    this.#maxRetries = opts.maxRetries ?? Infinity;
+    this.#maxElapsed = opts.maxElapsed ?? Infinity;
+  }
+
+  #clone() {
+    return new ExponentialBackoff(this.#type, {
+      base: this.#base,
+      factor: this.#factor,
+      min: this.#min,
+      max: this.#max,
+      maxRetries: this.#maxRetries,
+      maxElapsed: this.#maxElapsed,
+    });
+  }
+
+  type(type?: ExponentialBackoffType) {
+    if (typeof type !== "undefined") {
+      this.#type = type;
+    }
+    return this.#clone();
+  }
+
+  base(base?: number) {
+    if (typeof base !== "undefined") {
+      this.#base = base;
+    }
+    return this.#clone();
+  }
+
+  factor(factor?: number) {
+    if (typeof factor !== "undefined") {
+      this.#factor = factor;
+    }
+    return this.#clone();
+  }
+
+  min(min?: number) {
+    if (typeof min !== "undefined") {
+      this.#min = min;
+    }
+    return this.#clone();
+  }
+
+  max(max?: number) {
+    if (typeof max !== "undefined") {
+      this.#max = max;
+    }
+    return this.#clone();
+  }
+
+  maxRetries(maxRetries?: number) {
+    if (typeof maxRetries !== "undefined") {
+      this.#maxRetries = maxRetries;
+    }
+    return this.#clone();
+  }
+
+  maxElapsed(maxElapsed?: number) {
+    if (typeof maxElapsed !== "undefined") {
+      this.#maxElapsed = maxElapsed;
+    }
+    return this.#clone();
+  }
+
+  retries(retries?: number) {
+    if (typeof retries !== "undefined") {
+      if (retries > this.#maxRetries) {
+        console.error(
+          `Can't set retries ${retries} higher than maxRetries (${
+            this.#maxRetries
+          }), setting to maxRetries instead.`
+        );
+        this.#retries = this.#maxRetries;
+      } else {
+        this.#retries = retries;
+      }
+    }
+    return this.#clone();
+  }
+
+  async *retryAsync(maxRetries: number = this.#maxRetries ?? Infinity) {
+    let elapsed = 0;
+    let retry = 0;
+
+    while (retry <= maxRetries) {
+      const delay = this.delay(retry);
+      elapsed += delay;
+
+      if (elapsed > this.#maxElapsed) {
+        break;
+      }
+
+      yield {
+        delay: {
+          seconds: delay,
+          milliseconds: delay * 1000,
+        },
+        retry,
+      };
+
+      retry++;
+    }
+  }
+
+  async *[Symbol.asyncIterator]() {
+    yield* this.retryAsync();
+  }
+
+  delay(retries: number = this.#retries, jitter: boolean = true) {
+    if (retries > this.#maxRetries) {
+      console.error(
+        `Can't set retries ${retries} higher than maxRetries (${
+          this.#maxRetries
+        }), setting to maxRetries instead.`
+      );
+      retries = this.#maxRetries;
+    }
+
+    let delay = this.#factor * this.#base ** retries;
+
+    switch (this.#type) {
+      case "NoJitter": {
+        break;
+      }
+      case "FullJitter": {
+        if (!jitter) {
+          delay = 0;
+          break;
+        }
+
+        delay *= Math.random();
+        break;
+      }
+      case "EqualJitter": {
+        if (!jitter) {
+          delay *= 0.5;
+          break;
+        }
+
+        delay *= 0.5 * (1 + Math.random());
+        break;
+      }
+      default: {
+        throw new Error(`Unknown backoff type: ${this.#type}`);
+      }
+    }
+
+    delay = Math.min(delay, this.#max);
+    delay = Math.max(delay, this.#min);
+    delay = Math.round(delay);
+
+    return delay;
+  }
+
+  elapsed(retries: number = this.#retries, jitter: boolean = true) {
+    let elapsed = 0;
+
+    for (let i = 0; i <= retries; i++) {
+      elapsed += this.delay(i, jitter);
+    }
+
+    const total = elapsed;
+
+    let days = 0;
+    if (elapsed > 3600 * 24) {
+      days = Math.floor(elapsed / 3600 / 24);
+      elapsed -= days * 3600 * 24;
+    }
+
+    let hours = 0;
+    if (elapsed > 3600) {
+      hours = Math.floor(elapsed / 3600);
+      elapsed -= hours * 3600;
+    }
+
+    let minutes = 0;
+    if (elapsed > 60) {
+      minutes = Math.floor(elapsed / 60);
+      elapsed -= minutes * 60;
+    }
+
+    const seconds = elapsed;
+
+    return {
+      seconds,
+      minutes,
+      hours,
+      days,
+      total,
+    };
+  }
+
+  reset() {
+    this.#retries = 0;
+    return this;
+  }
+
+  next() {
+    this.#retries++;
+    return this.delay();
+  }
+}
@@ -13,6 +13,7 @@ import {
 import { ZodNamespace } from "@trigger.dev/core/v3/zodNamespace";
 import { ZodSocketConnection } from "@trigger.dev/core/v3/zodSocket";
 import { HttpReply, getTextBody, SimpleLogger } from "@trigger.dev/core-apps";
+import { ExponentialBackoff, StopRetrying } from "./backoff";

 import { collectDefaultMetrics, register, Gauge } from "prom-client";
 collectDefaultMetrics();
@@ -33,6 +34,13 @@ const SIMULATE_PUSH_FAILURE_SECONDS = parseInt(
  process.env.SIMULATE_PUSH_FAILURE_SECONDS ?? "300",
  10
 );
+const SIMULATE_CHECKPOINT_FAILURE = ["1", "true"].includes(
+  process.env.SIMULATE_CHECKPOINT_FAILURE ?? "false"
+);
+const SIMULATE_CHECKPOINT_FAILURE_SECONDS = parseInt(
+  process.env.SIMULATE_CHECKPOINT_FAILURE_SECONDS ?? "300",
+  10
+);

 const REGISTRY_HOST = process.env.REGISTRY_HOST || "localhost:5000";
 const CHECKPOINT_PATH = process.env.CHECKPOINT_PATH || "/checkpoints";
@@ -258,84 +266,100 @@ class Checkpointer {
    return true;
  }

-  async #checkpointAndPushWithBackoff(
-    {
-      runId,
-      leaveRunning = true, // This mirrors kubernetes behaviour more accurately
-      projectRef,
-      deploymentVersion,
-    }: CheckpointAndPushOptions,
-    retryCount = 0
-  ): Promise<CheckpointAndPushResult> {
-    const MAX_RETRIES = 10;
-
-    logger.log("Checkpointing with backoff", { runId, retryCount });
-
-    const result = await this.#checkpointAndPush({
+  async #checkpointAndPushWithBackoff({
+    runId,
+    leaveRunning = true, // This mirrors kubernetes behaviour more accurately
+    projectRef,
+    deploymentVersion,
+  }: CheckpointAndPushOptions): Promise<CheckpointAndPushResult> {
+    this.#logger.log("Checkpointing with backoff", {
      runId,
      leaveRunning,
      projectRef,
      deploymentVersion,
    });

-    if (result.success) {
-      return result;
+    const backoff = new ExponentialBackoff()
+      .type("EqualJitter")
+      .base(3)
+      .max(3 * 3600)
+      .maxElapsed(48 * 3600);
+
+    for await (const { delay, retry } of backoff) {
+      try {
+        if (retry > 0) {
+          this.#logger.error("Retrying checkpoint", {
+            runId,
+            retry,
+            delay,
+          });
+
+          this.#waitingForRetry.add(runId);
+          await new Promise((resolve) => setTimeout(resolve, delay.milliseconds));
+
+          if (!this.#waitingForRetry.has(runId)) {
+            this.#logger.log("Checkpoint canceled while waiting for retry", { runId });
+            return { success: false, reason: "CANCELED" };
+          } else {
+            this.#waitingForRetry.delete(runId);
+          }
+        }
+
+        const result = await this.#checkpointAndPush({
+          runId,
+          leaveRunning,
+          projectRef,
+          deploymentVersion,
+        });
+
+        if (result.success) {
+          return result;
+        }
+
+        if (result.reason === "CANCELED") {
+          this.#logger.log("Checkpoint canceled, won't retry", { runId });
+          // Don't fail the checkpoint, as it was canceled
+          return result;
+        }
+
+        if (result.reason === "IN_PROGRESS") {
+          this.#logger.log("Checkpoint already in progress, won't retry", { runId });
+          this.#failCheckpoint(runId, result.reason);
+          return result;
+        }
+
+        if (result.reason === "NO_SUPPORT") {
+          this.#logger.log("No checkpoint support, won't retry", { runId });
+          this.#failCheckpoint(runId, result.reason);
+          return result;
+        }
+
+        if (result.reason === "DISABLED") {
+          this.#logger.log("Checkpoint support disabled, won't retry", { runId });
+          this.#failCheckpoint(runId, result.reason);
+          return result;
+        }
+
+        continue;
+      } catch (error) {
+        this.#logger.error("Checkpoint error", {
+          retry,
+          runId,
+          delay,
+          error: error instanceof Error ? error.message : error,
+        });
+      }
    }

-    if (result.reason === "CANCELED") {
-      logger.log("Checkpoint canceled, won't retry", { runId });
-      // Don't fail the checkpoint, as it was canceled
-      return result;
-    }
+    this.#logger.error(`Checkpoint failed after exponential backoff`, {
+      runId,
+      leaveRunning,
+      projectRef,
+      deploymentVersion,
+    });
+    this.#failCheckpoint(runId, "ERROR");

-    if (result.reason === "IN_PROGRESS") {
-      logger.log("Checkpoint already in progress, won't retry", { runId });
-      this.#failCheckpoint(runId, result.reason);
-      return result;
-    }
-
-    if (result.reason === "NO_SUPPORT") {
-      logger.log("No checkpoint support, won't retry", { runId });
-      this.#failCheckpoint(runId, result.reason);
-      return result;
-    }
-
-    if (result.reason === "DISABLED") {
-      logger.log("Checkpoint support disabled, won't retry", { runId });
-      this.#failCheckpoint(runId, result.reason);
-      return result;
-    }
-
-    if (retryCount >= MAX_RETRIES) {
-      logger.error(`Checkpoint failed after ${MAX_RETRIES} retries`, { runId });
-      this.#failCheckpoint(runId, result.reason);
-      return result;
-    }
-
-    const retry = retryCount + 1;
-    const delay = exponentialBackoffMs(retry);
-
-    logger.log("Retrying checkpoint", { runId, retry, delay });
-
-    this.#waitingForRetry.add(runId);
-    await new Promise((resolve) => setTimeout(resolve, delay));
-
-    if (!this.#waitingForRetry.has(runId)) {
-      logger.log("Checkpoint canceled while waiting for retry", { runId });
-      return { success: false, reason: "CANCELED" };
-    } else {
-      this.#waitingForRetry.delete(runId);
-    }
-
-    return this.#checkpointAndPushWithBackoff(
-      {
-        runId,
-        leaveRunning,
-        projectRef,
-        deploymentVersion,
-      },
-      retry
-    );
+    return { success: false, reason: "ERROR" };
  }

  async #checkpointAndPush({
@@ -424,6 +448,13 @@ class Checkpointer {
            this.#logger.log("Simulating checkpoint");
            this.#logger.debug(await $$`docker pause ${containterName}`);
          } else {
+            if (SIMULATE_CHECKPOINT_FAILURE) {
+              if (performance.now() < SIMULATE_CHECKPOINT_FAILURE_SECONDS * 1000) {
+                this.#logger.error("Simulating checkpoint failure", { options });
+                throw new Error("SIMULATE_CHECKPOINT_FAILURE");
+              }
+            }
+
            if (leaveRunning) {
              this.#logger.debug(
                await $$`docker checkpoint create --leave-running ${containterName} ${exportLocation}`
@@ -471,6 +502,13 @@ class Checkpointer {

      const start = performance.now();

+      if (SIMULATE_CHECKPOINT_FAILURE) {
+        if (performance.now() < SIMULATE_CHECKPOINT_FAILURE_SECONDS * 1000) {
+          this.#logger.error("Simulating checkpoint failure", { options });
+          throw new Error("SIMULATE_CHECKPOINT_FAILURE");
+        }
+      }
+
      // Create checkpoint
      this.#logger.debug(await $$`crictl checkpoint --export=${exportLocation} ${containerId}`);
      const postCheckpoint = performance.now();
@@ -699,6 +737,8 @@ class TaskCoordinator {
            return;
          }

+          this.#checkpointer.cancelCheckpoint(message.runId);
+
          if (message.delayInMs) {
            taskSocket.emit("REQUEST_EXIT", {
              version: "v2",