circut_breaker.ts

import {
  ObjectContext,
  Rand,
  service,
  TerminalError,
  type Context,
} from "@restatedev/restate-sdk";

class CircuitBreakerState {
  constructor(
    private state: "closed" | "open" = "closed",
    private rand: Rand
  ) {}

  getState() {
    return this.state;
  }

  tryCall() {
    const isClosed = this.state === "closed";
    return { call: isClosed || this.rand.random() < 0.7 };
  }

  onCallSuccess() {
    if (this.state === "open") {
      this.state = "closed";
    }
  }
  onCallFailure() {
    this.state = "open";
  }
}

class CircuitBreaker<T> {
  constructor(
    public state: "closed" | "open" = "closed",
    private terminalOnCircuitOpen: boolean = false
  ) {}

  public async call(action: () => Promise<T>): Promise<T> {
    switch (this.state) {
      case "closed":
        try {
          return await action();
        } catch (error) {
          this.state = "open";
          throw error;
        }
      case "open":
        if (Math.random() < 0.7) {
          const err = this.terminalOnCircuitOpen
            ? new TerminalError("Circuit is open. Try again later.")
            : new Error("Circuit is open. Try again later.");
          throw err;
        }
        try {
          const result = await action();
          this.state = "closed";
          return result;
        } catch (error) {
          this.state = "open";
          throw error;
        }
    }
  }
}

/***
 *  OPTION #1
 *  ---------
 *
 *  This is a global circuit breaker *per node process*.
 *  It is shared across all handlers that run in the same node process.
 *  Use this when you have multiple handlers that call the same 3rd party service
 *  and they are running in the same node process, or you have relatively few node processes running,
 *  which approximate *global* circuit breaker.
 *
 */

const GLOBAL_PER_NODE_PROCESS = new CircuitBreaker<string>();

const someHandler = async (ctx: Context, arg: string) => {
  await ctx.run(
    "call a 3rd party service",
    () =>
      GLOBAL_PER_NODE_PROCESS.call(async () => {
        return `Hello ${arg}`;
      }),
    { initialRetryIntervalMillis: 5000 }
  );
};

/***
 *
 * OPTION #2
 * ---------
 *
 * Virtual object
 *
 */
const handler = async (ctx: ObjectContext, awekable: string) => {
  const state = (await ctx.get<"open" | "closed">("state")) ?? "closed";
  const breaker = new CircuitBreakerState(state, ctx.rand);

  if (!breaker.tryCall()) {
    ctx.rejectAwakeable(awekable, "closed");
  }

  ctx.set("state", breaker.getState());

  
};