Rate Limiting

import { Context, TerminalError } from "@restatedev/restate-sdk";
import type { Limiter as LimiterObject, Reservation as ReservationResponse } from "./limiter";

export interface Reservation extends ReservationResponse {
  // cancel indicates that the reservation holder will not perform the reserved action
  // and reverses the effects of this Reservation on the rate limit as much as possible,
  // considering that other reservations may have already been made.
  cancel(): void;
}

export interface Limiter {
  // limit returns the maximum overall event rate.
  limit(): Promise<number>;
  // burst returns the maximum burst size. Burst is the maximum number of tokens
  // that can be consumed in a single call to allow, reserve, or wait, so higher
  // Burst values allow more events to happen at once.
  // A zero Burst allows no events, unless limit == Inf.
  burst(): Promise<number>;
  // tokens returns the number of tokens available at time t (defaults to now).
  tokens(): Promise<number>;
  // allow reports whether n events may happen at time t.
  // Use this method if you intend to drop / skip events that exceed the rate limit.
  // Otherwise use reserve or wait.
  allow(n?: number): Promise<boolean>;
  // reserve returns a Reservation that indicates how long the caller must wait before n events happen.
  // The limiter takes this Reservation into account when allowing future events.
  // The returned Reservation’s ok parameter is false if n exceeds the limiter's burst size, or provided waitLimitMillis.
  // Use this method if you wish to wait and slow down in accordance with the rate limit without dropping events.
  // If you need to cancel the delay, use wait instead.
  // To drop or skip events exceeding rate limit, use allow instead.
  reserve(n?: number, waitLimitMillis?: number): Promise<Reservation>;
  // setLimit sets a new limit for the limiter. The new limit, and burst, may be violated
  // or underutilized by those which reserved (using reserve or wait) but did not yet act
  // before setLimit was called.
  setLimit(newLimit: number): Promise<void>;
  // setBurst sets a new burst size for the limiter.
  setBurst(newBurst: number): Promise<void>;
  // setRate sets a new limit and burst size for the limiter.
  setRate(newLimit: number, newBurst: number): Promise<void>;
  // waitN blocks until the limiter permits n events to happen.
  // It returns an error if n exceeds the limiter's burst size, the invocation is canceled,
  // or the wait would be longer than the deadline.
  // The burst limit is ignored if the rate limit is Inf.
  wait(n?: number, waitLimitMillis?: number): Promise<void>;
}

export namespace Limiter {
  export function fromContext(ctx: Context, limiterID: string): Limiter {
    const client = ctx.objectClient<LimiterObject>({ name: "limiter" }, limiterID);
    return {
      async limit() {
        return (await client.state()).limit;
      },
      async burst() {
        return (await client.state()).burst;
      },
      async tokens() {
        return client.tokens();
      },
      async allow(n?: number) {
        const r = await client.reserve({
          n,
          waitLimitMillis: 0,
        });
        return r.ok;
      },
      async reserve(n?: number, waitLimitMillis?: number) {
        const r = await client.reserve({
          n,
          waitLimitMillis,
        });
        return {
          cancel() {
            ctx
              .objectSendClient<LimiterObject>({ name: "limiter" }, limiterID)
              .cancelReservation(r);
          },
          ...r,
        };
      },
      async setLimit(newLimit: number) {
        return client.setRate({
          newLimit,
        });
      },
      async setBurst(newBurst: number) {
        return client.setRate({
          newBurst,
        });
      },
      async setRate(newLimit: number, newBurst: number) {
        return client.setRate({
          newLimit,
          newBurst,
        });
      },
      async wait(n: number = 1, waitLimitMillis?: number) {
        // Reserve
        const r = await this.reserve(n, waitLimitMillis);
        if (!r.ok) {
          if (waitLimitMillis === undefined) {
            throw new TerminalError(`rate: Wait(n=${n}) would exceed the limiters burst`, {
              errorCode: 429,
            });
          } else {
            throw new TerminalError(
              `rate: Wait(n=${n}) would either exceed the limiters burst or the provided waitLimitMillis`,
              { errorCode: 429 },
            );
          }
        }
        // Wait if necessary
        const delay = delayFrom(r, r.creationDate);
        if (delay == 0) {
          return;
        }

        try {
          await ctx.sleep(delay);
        } catch (e) {
          // this only happens on invocation cancellation - cancel the reservation in the background
          r.cancel();
          throw e;
        }
      },
    };
  }
}

// delayFrom returns the duration in millis for which the reservation holder must wait
// before taking the reserved action.  Zero duration means act immediately.
// Infinity means the limiter cannot grant the tokens requested in this
// Reservation within the maximum wait time.
function delayFrom(r: ReservationResponse, date: number): number {
  if (!r.ok) {
    return Infinity;
  }
  const delay = r.dateToAct - date;
  if (delay < 0) {
    return 0;
  }
  return Math.floor(delay);
}

// a faithful reimplementation of https://pkg.go.dev/golang.org/x/time/rate#Limiter
// using virtual object state

import { object, ObjectContext } from "@restatedev/restate-sdk";

type LimiterState = {
  state: LimiterStateInner;
};
type LimiterStateInner = {
  limit: number;
  burst: number;
  tokens: number;
  // last is the last time the limiter's tokens field was updated, in unix millis
  last: number;
  // lastEvent is the latest time of a rate-limited event (past or future), in unix millis
  lastEvent: number;
};

export interface Reservation {
  ok: boolean;
  tokens: number;
  creationDate: number;
  dateToAct: number;
  // This is the Limit at reservation time, it can change later.
  limit: number;
}

export const limiter = object({
  name: "limiter",
  handlers: {
    state: async (ctx: ObjectContext<LimiterState>): Promise<LimiterStateInner> => {
      return getState(ctx);
    },
    tokens: async (ctx: ObjectContext<LimiterState>): Promise<number> => {
      // deterministic date not needed, as there is only an output entry
      const tokens = advance(await getState(ctx), Date.now());
      return tokens;
    },
    reserve: async (
      ctx: ObjectContext<LimiterState>,
      { n = 1, waitLimitMillis = Infinity }: { n?: number; waitLimitMillis?: number },
    ): Promise<Reservation> => {
      let lim = await getState(ctx);

      if (lim.limit == Infinity) {
        // deterministic date is not necessary, as this is part of a response body, which won't be replayed.
        const now = Date.now();
        return {
          ok: true,
          tokens: n,
          creationDate: now,
          dateToAct: now,
          limit: 0,
        };
      }

      let r: Reservation;
      ({ lim, r } = await ctx.run(() => {
        const now = Date.now();
        let tokens = advance(lim, now);

        // Calculate the remaining number of tokens resulting from the request.
        tokens -= n;

        // Calculate the wait duration
        let waitDurationMillis = 0;
        if (tokens < 0) {
          waitDurationMillis = durationFromTokens(lim.limit, -tokens);
        }

        // Decide result
        const ok = n <= lim.burst && waitDurationMillis <= waitLimitMillis;

        // Prepare reservation
        const r = {
          ok,
          tokens: 0,
          creationDate: now,
          dateToAct: 0,
          limit: lim.limit,
        } satisfies Reservation;

        if (ok) {
          r.tokens = n;
          r.dateToAct = now + waitDurationMillis;

          // Update state
          lim.last = now;
          lim.tokens = tokens;
          lim.lastEvent = r.dateToAct;
        }

        return { lim, r };
      }));

      setState(ctx, lim);

      return r;
    },
    setRate: async (
      ctx: ObjectContext<LimiterState>,
      { newLimit, newBurst }: { newLimit?: number; newBurst?: number },
    ) => {
      if (newLimit === undefined && newBurst === undefined) {
        return;
      }

      let lim = await getState(ctx);

      lim = await ctx.run(() => {
        const now = Date.now();
        const tokens = advance(lim, now);

        lim.last = now;
        lim.tokens = tokens;
        if (newLimit !== undefined) lim.limit = newLimit;
        if (newBurst !== undefined) lim.burst = newBurst;

        return lim;
      });

      setState(ctx, lim);
    },
    cancelReservation: async (ctx: ObjectContext<LimiterState>, r: Reservation) => {
      let lim = await getState(ctx);

      lim = await ctx.run(() => {
        const now = Date.now();

        if (lim.limit == Infinity || r.tokens == 0 || r.dateToAct < now) {
          return lim;
        }

        // calculate tokens to restore
        // The duration between lim.lastEvent and r.timeToAct tells us how many tokens were reserved
        // after r was obtained. These tokens should not be restored.
        const restoreTokens = r.tokens - tokensFromDuration(r.limit, lim.lastEvent - r.dateToAct);
        if (restoreTokens <= 0) {
          return lim;
        }
        // advance time to now
        let tokens = advance(lim, now);
        // calculate new number of tokens
        tokens += restoreTokens;
        if (tokens > lim.burst) {
          tokens = lim.burst;
        }
        // update state
        lim.last = now;
        lim.tokens = tokens;
        if (r.dateToAct == lim.lastEvent) {
          const prevEvent = r.dateToAct + durationFromTokens(r.limit, -r.tokens);
          if (prevEvent >= now) {
            lim.lastEvent = prevEvent;
          }
        }

        return lim;
      });

      setState(ctx, lim);
    },
  },
});

function advance(lim: LimiterStateInner, date: number): number {
  let last = lim.last;
  if (date <= last) {
    last = date;
  }

  // Calculate the new number of tokens, due to time that passed.
  const elapsedMillis = date - last;
  const delta = tokensFromDuration(lim.limit, elapsedMillis);
  let tokens = lim.tokens + delta;
  if (tokens > lim.burst) {
    tokens = lim.burst;
  }

  return tokens;
}

async function getState(ctx: ObjectContext<LimiterState>): Promise<LimiterStateInner> {
  return (
    (await ctx.get("state")) ?? {
      limit: 0,
      burst: 0,
      tokens: 0,
      last: 0,
      lastEvent: 0,
    }
  );
}

async function setState(ctx: ObjectContext<LimiterState>, lim: LimiterStateInner) {
  ctx.set("state", lim);
}

function durationFromTokens(limit: number, tokens: number): number {
  if (limit <= 0) {
    return Infinity;
  }

  return (tokens / limit) * 1000;
}

function tokensFromDuration(limit: number, durationMillis: number): number {
  if (limit <= 0) {
    return 0;
  }
  return (durationMillis / 1000) * limit;
}

export type Limiter = typeof limiter;