revert: use durable objects for sync ratelimiting

apps/api/src/pkg/env.ts

@@ -18,6 +18,7 @@ export const zEnv = z.object({
   CLOUDFLARE_API_KEY: z.string().optional(),
   CLOUDFLARE_ZONE_ID: z.string().optional(),
   ENVIRONMENT: z.enum(["development", "preview", "canary", "production"]).default("development"),
+  DO_RATELIMIT: z.custom<DurableObjectNamespace>((ns) => typeof ns === "object"), // pretty loose check but it'll do I think
   DO_USAGELIMIT: z.custom<DurableObjectNamespace>((ns) => typeof ns === "object"),
   KEY_MIGRATIONS: z.custom<Queue<MessageBody>>((q) => typeof q === "object").optional(),
   EMIT_METRICS_LOGS: z

apps/api/src/pkg/middleware/init.ts

@@ -2,7 +2,7 @@ import { Analytics } from "@/pkg/analytics";
 import { createConnection } from "@/pkg/db";
 
 import { KeyService } from "@/pkg/keys/service";
-import { AgentRatelimiter } from "@/pkg/ratelimit";
+import { DurableRateLimiter } from "@/pkg/ratelimit";
 import { DurableUsageLimiter, NoopUsageLimiter } from "@/pkg/usagelimit";
 import { RBAC } from "@unkey/rbac";
 import { ConsoleLogger } from "@unkey/worker-logging";
@@ -100,8 +100,8 @@ export function init(): MiddlewareHandler<HonoEnv> {
       clickhouseUrl: c.env.CLICKHOUSE_URL,
       clickhouseInsertUrl: c.env.CLICKHOUSE_INSERT_URL,
     });
-    const rateLimiter = new AgentRatelimiter({
-      agent: { url: c.env.AGENT_URL, token: c.env.AGENT_TOKEN },
+    const rateLimiter = new DurableRateLimiter({
+      namespace: c.env.DO_RATELIMIT,
       cache: rlMap,
       logger,
       metrics,

apps/api/src/pkg/ratelimit/do_client.ts

@@ -0,0 +1,240 @@
+import { Err, Ok, type Result } from "@unkey/error";
+import type { Logger } from "@unkey/worker-logging";
+import type { Context } from "hono";
+import { z } from "zod";
+import type { Metrics } from "../metrics";
+import {
+  type RateLimiter,
+  RatelimitError,
+  type RatelimitRequest,
+  type RatelimitResponse,
+} from "./interface";
+
+export class DurableRateLimiter implements RateLimiter {
+  private readonly namespace: DurableObjectNamespace;
+  private readonly domain: string;
+  private readonly logger: Logger;
+  private readonly metrics: Metrics;
+  private readonly cache: Map<string, number>;
+  constructor(opts: {
+    namespace: DurableObjectNamespace;
+
+    domain?: string;
+    logger: Logger;
+    metrics: Metrics;
+    cache: Map<string, number>;
+  }) {
+    this.namespace = opts.namespace;
+    this.domain = opts.domain ?? "unkey.dev";
+    this.logger = opts.logger;
+    this.metrics = opts.metrics;
+    this.cache = opts.cache;
+  }
+
+  private getId(req: RatelimitRequest): string {
+    const now = Date.now();
+    const window = Math.floor(now / req.interval);
+
+    return [req.identifier, window, req.shard].join("::");
+  }
+
+  private setCacheMax(id: string, i: number): number {
+    const current = this.cache.get(id) ?? 0;
+    if (i > current) {
+      this.cache.set(id, i);
+      return i;
+    }
+    return current;
+  }
+
+  /**
+   * Do not use
+   */
+  public async multiLimit(
+    c: Context,
+    req: Array<RatelimitRequest>,
+  ): Promise<Result<RatelimitResponse, RatelimitError>> {
+    const res = await Promise.all(req.map((r) => this.limit(c, r)));
+    for (const r of res) {
+      if (r.err) {
+        return r;
+      }
+      if (!r.val.passed) {
+        return r;
+      }
+    }
+    if (res.length > 0) {
+      return Ok(res[0].val!);
+    }
+
+    return Ok({
+      current: -1,
+      passed: true,
+      reset: -1,
+      remaining: -1,
+      triggered: null,
+    });
+  }
+
+  public async limit(
+    c: Context,
+    req: RatelimitRequest,
+  ): Promise<Result<RatelimitResponse, RatelimitError>> {
+    const start = performance.now();
+    const res = await this._limit(c, req);
+    this.metrics.emit({
+      metric: "metric.ratelimit",
+      workspaceId: req.workspaceId,
+      namespaceId: req.namespaceId,
+      latency: performance.now() - start,
+      identifier: req.identifier,
+      mode: req.async ? "async" : "sync",
+      error: !!res.err,
+      success: res?.val?.passed,
+      source: "durable_object",
+    });
+    return res;
+  }
+
+  private async _limit(
+    c: Context,
+    req: RatelimitRequest,
+  ): Promise<Result<RatelimitResponse, RatelimitError>> {
+    const window = Math.floor(Date.now() / req.interval);
+    const reset = (window + 1) * req.interval;
+    const cost = req.cost ?? 1;
+    const id = this.getId(req);
+
+    /**
+     * Catching identifiers that exceeded the limit already
+     *
+     * This might not happen too often, but in extreme cases the cache should hit and we can skip
+     * the request to the durable object entirely, which speeds everything up and is cheper for us
+     */
+    let current = this.cache.get(id) ?? 0;
+    if (current >= req.limit) {
+      return Ok({
+        remaining: -1,
+        passed: false,
+        current,
+        reset,
+        triggered: req.name,
+      });
+    }
+
+    const p = this.callDurableObject({
+      trigger: req.name,
+      identifier: req.identifier,
+      objectName: id,
+      window,
+      reset,
+      cost,
+      limit: req.limit,
+    });
+
+    if (!req.async) {
+      const res = await p;
+      if (res.val) {
+        this.setCacheMax(id, res.val.current);
+      }
+      return res;
+    }
+
+    c.executionCtx.waitUntil(
+      p.then(async (res) => {
+        if (res.err) {
+          console.error(res.err.message);
+          return;
+        }
+        this.setCacheMax(id, res.val.current);
+
+        this.metrics.emit({
+          workspaceId: req.workspaceId,
+          metric: "metric.ratelimit.accuracy",
+          identifier: req.identifier,
+          namespaceId: req.namespaceId,
+          responded: current + cost <= req.limit,
+          correct: res.val.current + cost <= req.limit,
+        });
+        await this.metrics.flush();
+      }),
+    );
+    if (current + cost > req.limit) {
+      return Ok({
+        current,
+        passed: false,
+        reset,
+        remaining: req.limit - current,
+        triggered: req.name,
+      });
+    }
+    current += cost;
+    this.cache.set(id, current);
+
+    return Ok({
+      passed: true,
+      current,
+      reset,
+      remaining: req.limit - current,
+      triggered: null,
+    });
+  }
+
+  private getStub(name: string): DurableObjectStub {
+    return this.namespace.get(this.namespace.idFromName(name));
+  }
+
+  private async callDurableObject(req: {
+    identifier: string;
+    objectName: string;
+    window: number;
+    reset: number;
+    cost: number;
+    limit: number;
+    trigger: string;
+  }): Promise<Result<RatelimitResponse, RatelimitError>> {
+    try {
+      const url = `https://${this.domain}/limit`;
+      const res = await this.getStub(req.objectName)
+        .fetch(url, {
+          method: "POST",
+          headers: { "Content-Type": "application/json" },
+          body: JSON.stringify({ reset: req.reset, cost: req.cost, limit: req.limit }),
+        })
+        .catch(async (e) => {
+          this.logger.warn("calling the ratelimit DO failed, retrying ...", {
+            identifier: req.identifier,
+            error: (e as Error).message,
+          });
+
+          return this.getStub(req.objectName).fetch(url, {
+            method: "POST",
+            headers: { "Content-Type": "application/json" },
+            body: JSON.stringify({ reset: req.reset }),
+          });
+        });
+
+      const json = await res.json();
+      const { current, success } = z
+        .object({ current: z.number(), success: z.boolean() })
+        .parse(json);
+
+      return Ok({
+        current,
+        reset: req.reset,
+        passed: success,
+        remaining: req.limit - current,
+        triggered: success ? null : req.trigger,
+      });
+    } catch (e) {
+      const err = e as Error;
+      this.logger.error("ratelimit failed", {
+        identifier: req.identifier,
+        error: err.message,
+        stack: err.stack,
+        cause: err.cause,
+      });
+      return Err(new RatelimitError({ message: err.message }));
+    }
+  }
+}

apps/api/src/pkg/ratelimit/durable_object.ts

@@ -0,0 +1,72 @@
+import { zValidator } from "@hono/zod-validator";
+import { Hono } from "hono";
+import { z } from "zod";
+
+type Memory = {
+  current: number;
+  alarmScheduled?: number;
+};
+
+export class DurableObjectRatelimiter {
+  private state: DurableObjectState;
+  private memory: Memory;
+  private readonly storageKey = "rl";
+  private readonly hono = new Hono();
+  constructor(state: DurableObjectState) {
+    this.state = state;
+    this.state.blockConcurrencyWhile(async () => {
+      const m = await this.state.storage.get<Memory>(this.storageKey);
+      if (m) {
+        this.memory = m;
+      }
+    });
+    this.memory ??= {
+      current: 0,
+    };
+
+    this.hono.post(
+      "/limit",
+      zValidator(
+        "json",
+        z.object({
+          reset: z.number().int(),
+          cost: z.number().int().default(1),
+          limit: z.number().int(),
+        }),
+      ),
+      async (c) => {
+        const { reset, cost, limit } = c.req.valid("json");
+        if (!this.memory.alarmScheduled) {
+          this.memory.alarmScheduled = reset;
+          await this.state.storage.setAlarm(this.memory.alarmScheduled);
+        }
+        if (this.memory.current + cost > limit) {
+          return c.json({
+            success: false,
+            current: this.memory.current,
+          });
+        }
+        this.memory.current += cost;
+
+        await this.state.storage.put(this.storageKey, this.memory);
+
+        return c.json({
+          success: true,
+          current: this.memory.current,
+        });
+      },
+    );
+  }
+
+  // Handle HTTP requests from clients.
+  async fetch(request: Request) {
+    return this.hono.fetch(request);
+  }
+
+  /**
+   * alarm is called to clean up all state, which will remove the durable object from existence.
+   */
+  public async alarm(): Promise<void> {
+    await this.state.storage.deleteAll();
+  }
+}

apps/api/src/pkg/ratelimit/index.ts

@@ -1,3 +1,5 @@
 export * from "./client";
 export * from "./interface";
 export * from "./noop";
+export * from "./durable_object";
+export * from "./do_client";