diff --git a/apps/api/src/pkg/hono/env.ts b/apps/api/src/pkg/hono/env.ts index 3f84455824..991193f7b7 100644 --- a/apps/api/src/pkg/hono/env.ts +++ b/apps/api/src/pkg/hono/env.ts @@ -27,6 +27,7 @@ export type HonoEnv = { Bindings: Env; Variables: { isolateId: string; + isolateCreatedAt: number; requestId: string; metricsContext: { keyId?: string; diff --git a/apps/api/src/pkg/middleware/init.ts b/apps/api/src/pkg/middleware/init.ts index a1336567d5..70732d3fa1 100644 --- a/apps/api/src/pkg/middleware/init.ts +++ b/apps/api/src/pkg/middleware/init.ts @@ -38,8 +38,10 @@ export function init(): MiddlewareHandler { isolateId = crypto.randomUUID(); } c.set("isolateId", isolateId); + c.set("isolateCreatedAt", Date.now()); const requestId = newId("request"); c.set("requestId", requestId); + c.res.headers.set("Unkey-Request-Id", requestId); const logger = new ConsoleLogger({ diff --git a/apps/api/src/pkg/middleware/metrics.ts b/apps/api/src/pkg/middleware/metrics.ts index cefe2e8985..0e9bb6ff32 100644 --- a/apps/api/src/pkg/middleware/metrics.ts +++ b/apps/api/src/pkg/middleware/metrics.ts @@ -4,8 +4,6 @@ import type { HonoEnv } from "../hono/env"; type DiscriminateMetric = M extends { metric: T } ? M : never; -let coldstartAt: number | null = null; - export function metrics(): MiddlewareHandler { return async (c, next) => { const { metrics, analytics, logger } = c.get("services"); @@ -15,9 +13,10 @@ export function metrics(): MiddlewareHandler { // }); // const start = performance.now(); + const isolateCreatedAt = c.get("isolateCreatedAt"); const m = { isolateId: c.get("isolateId"), - isolateLifetime: coldstartAt ? Date.now() - coldstartAt : 0, + isolateLifetime: isolateCreatedAt ? Date.now() - isolateCreatedAt : 0, metric: "metric.http.request", path: c.req.path, host: new URL(c.req.url).host, @@ -35,7 +34,6 @@ export function metrics(): MiddlewareHandler { context: {}, } as DiscriminateMetric<"metric.http.request">; - coldstartAt = Date.now(); try { const telemetry = { runtime: c.req.header("Unkey-Telemetry-Runtime"), diff --git a/apps/api/src/pkg/ratelimit/client.ts b/apps/api/src/pkg/ratelimit/client.ts index d83b0aff02..40d8b91ba6 100644 --- a/apps/api/src/pkg/ratelimit/client.ts +++ b/apps/api/src/pkg/ratelimit/client.ts @@ -166,7 +166,10 @@ export class AgentRatelimiter implements RateLimiter { })(); // A rollout of the sync rate limiting - const shouldSyncOnNoData = Math.random() < c.env.SYNC_RATELIMIT_ON_NO_DATA; + // Isolates younger than 60s must not sync. It would cause a stampede of requests as the cache is entirely empty + const isolateCreatedAt = c.get("isolateCreatedAt"); + const isOlderThan60s = isolateCreatedAt ? Date.now() - isolateCreatedAt > 60_000 : false; + const shouldSyncOnNoData = isOlderThan60s && Math.random() < c.env.SYNC_RATELIMIT_ON_NO_DATA; const cacheHit = this.cache.has(id); const sync = !req.async || (!cacheHit && shouldSyncOnNoData); this.logger.info("sync rate limiting", { diff --git a/apps/api/src/routes/v1_keys_verifyKey.ratelimit_accuracy.test.ts b/apps/api/src/routes/v1_keys_verifyKey.ratelimit_accuracy.test.ts new file mode 100644 index 0000000000..c898035e8e --- /dev/null +++ b/apps/api/src/routes/v1_keys_verifyKey.ratelimit_accuracy.test.ts @@ -0,0 +1,102 @@ +import { test } from "vitest"; + +import { loadTest } from "@/pkg/testutil/load"; +import { schema } from "@unkey/db"; +import { newId } from "@unkey/id"; +import { IntegrationHarness } from "src/pkg/testutil/integration-harness"; + +import { randomUUID } from "node:crypto"; +import type { V1KeysVerifyKeyRequest, V1KeysVerifyKeyResponse } from "./v1_keys_verifyKey"; + +/** + * As a rule of thumb, the test duration (seconds) should be at least 10x the duration of the rate limit window + */ +const testCases: { + limit: number; + duration: number; + rps: number; + seconds: number; +}[] = [ + { + limit: 200, + duration: 10_000, + rps: 100, + seconds: 60, + }, + { + limit: 10, + duration: 10000, + rps: 15, + seconds: 120, + }, + { + limit: 20, + duration: 1000, + rps: 50, + seconds: 60, + }, + { + limit: 200, + duration: 10000, + rps: 20, + seconds: 20, + }, + { + limit: 500, + duration: 10000, + rps: 100, + seconds: 30, + }, + { + limit: 100, + duration: 5000, + rps: 200, + seconds: 120, + }, +]; + +for (const { limit, duration, rps, seconds } of testCases) { + const name = `[${limit} / ${duration / 1000}s], attacked with ${rps} rps for ${seconds}s`; + test(name, { skip: process.env.TEST_LOCAL, retry: 3, timeout: 600_000 }, async (t) => { + const h = await IntegrationHarness.init(t); + + const { key, keyId } = await h.createKey(); + + const ratelimitName = randomUUID(); + await h.db.primary.insert(schema.ratelimits).values({ + id: newId("test"), + name: ratelimitName, + keyId, + limit, + duration, + workspaceId: h.resources.userWorkspace.id, + }); + + const results = await loadTest({ + rps, + seconds, + fn: () => + h.post({ + url: "/v1/keys.verifyKey", + headers: { + "Content-Type": "application/json", + }, + body: { + key, + ratelimits: [{ name: ratelimitName }], + }, + }), + }); + t.expect(results.length).toBe(rps * seconds); + const passed = results.reduce((sum, res) => { + return res.body.valid ? sum + 1 : sum; + }, 0); + + const exactLimit = (limit / (duration / 1000)) * seconds; + const upperLimit = Math.round(exactLimit * 1.3); + const lowerLimit = Math.round(exactLimit * 0.9); + console.info({ name, passed, exactLimit, upperLimit, lowerLimit }); + t.expect(passed).toBeGreaterThanOrEqual(lowerLimit); + t.expect(passed).toBeLessThanOrEqual(upperLimit); + }); +} diff --git a/apps/api/src/routes/v1_keys_verifyKey.ratelimit_consistency.test.ts.skipped b/apps/api/src/routes/v1_keys_verifyKey.ratelimit_consistency.test.ts.skipped deleted file mode 100644 index ee8be26c8e..0000000000 --- a/apps/api/src/routes/v1_keys_verifyKey.ratelimit_consistency.test.ts.skipped +++ /dev/null @@ -1,67 +0,0 @@ -import { describe, expect, test } from "vitest"; - -import { IntegrationHarness } from "@/pkg/testutil/integration-harness"; -import { schema } from "@unkey/db"; -import { sha256 } from "@unkey/hash"; -import { newId } from "@unkey/id"; -import { KeyV1 } from "@unkey/keys"; -import type { V1KeysVerifyKeyRequest, V1KeysVerifyKeyResponse } from "./v1_keys_verifyKey"; - -describe.each<{ limit: number; duration: number; n: number }>([ - { limit: 10, duration: 1_000, n: 100 }, - { limit: 10, duration: 2_000, n: 100 }, - { limit: 500, duration: 1_000, n: 100 }, - { limit: 500, duration: 60_000, n: 100 }, - // { limit: 1000, duration: 1_000, n: 250 }, -])("$limit per $duration ms @ $n runs", async ({ limit, duration, n }) => { - test( - "counts down monotonically", - async (t) => { - const h = await IntegrationHarness.init(t); - - const key = new KeyV1({ prefix: "test", byteLength: 16 }).toString(); - await h.db.primary.insert(schema.keys).values({ - id: newId("test"), - keyAuthId: h.resources.userKeyAuth.id, - hash: await sha256(key), - start: key.slice(0, 8), - workspaceId: h.resources.userWorkspace.id, - createdAt: new Date(), - ratelimitLimit: limit, - ratelimitDuration: duration, - ratelimitAsync: false, - }); - - let errors = 0; - let lastResponse = limit; - for (let i = 0; i < n; i++) { - const res = await h.post({ - url: "/v1/keys.verifyKey", - headers: { - "Content-Type": "application/json", - }, - body: { - key, - apiId: h.resources.userApi.id, - }, - }); - - expect(res.status, `Received wrong status, res: ${JSON.stringify(res)}`).toEqual(200); - expect(res.body.ratelimit).toBeDefined(); - /** - * It should either be counting down monotonically, or be reset in a new window - */ - const acceptable = [Math.max(0, lastResponse - 1), limit - 1]; - const correct = acceptable.some((acc) => res.body.ratelimit!.remaining === acc); - - if (!correct) { - errors += 1; - console.warn("Inconsistent remaining", { correct, acceptable, body: res.body }); - } - lastResponse = res.body.ratelimit!.remaining; - } - expect(errors).toBeLessThanOrEqual(n / 5); - }, - { timeout: 120_000 }, - ); -}); diff --git a/apps/api/src/routes/v1_ratelimit_limit.accuracy.test.ts b/apps/api/src/routes/v1_ratelimit_limit.accuracy.test.ts index d33e738345..39c441fde0 100644 --- a/apps/api/src/routes/v1_ratelimit_limit.accuracy.test.ts +++ b/apps/api/src/routes/v1_ratelimit_limit.accuracy.test.ts @@ -12,101 +12,94 @@ import type { V1RatelimitLimitRequest, V1RatelimitLimitResponse } from "./v1_rat * As a rule of thumb, the test duration (seconds) should be at least 10x the duration of the rate limit window */ const testCases: { - name: string; limit: number; duration: number; rps: number; seconds: number; - expected: { min: number; max: number }; }[] = [ { - name: "Basic Test", + limit: 200, + duration: 10_000, + rps: 100, + seconds: 60, + }, + { limit: 10, duration: 10000, rps: 15, seconds: 120, - expected: { min: 120, max: 150 }, }, { - name: "High Rate with Short Window", limit: 20, duration: 1000, rps: 50, seconds: 60, - expected: { min: 1200, max: 1500 }, }, { - name: "Constant Rate Equals Limit", limit: 200, duration: 10000, rps: 20, seconds: 20, - expected: { min: 400, max: 400 }, }, { - name: "Rate Lower Than Limit", limit: 500, duration: 10000, rps: 100, seconds: 30, - expected: { min: 1500, max: 2000 }, }, { - name: "Rate Higher Than Limit", limit: 100, duration: 5000, rps: 200, seconds: 120, - expected: { min: 2400, max: 3000 }, }, ]; -for (const { name, limit, duration, rps, seconds, expected } of testCases) { - test( - `${name}, [${limit} / ${duration / 1000}s], passed requests are within [${expected.min} - ${ - expected.max - }]`, - { skip: process.env.TEST_LOCAL, retry: 3, timeout: 600_000 }, - async (t) => { - const h = await IntegrationHarness.init(t); - const namespace = { - id: newId("test"), - workspaceId: h.resources.userWorkspace.id, - createdAt: new Date(), - name: "namespace", - }; - await h.db.primary.insert(schema.ratelimitNamespaces).values(namespace); +for (const { limit, duration, rps, seconds } of testCases) { + const name = `[${limit} / ${duration / 1000}s], attacked with ${rps} rps for ${seconds}s`; + test(name, { skip: process.env.TEST_LOCAL, retry: 3, timeout: 600_000 }, async (t) => { + const h = await IntegrationHarness.init(t); + const namespace = { + id: newId("test"), + workspaceId: h.resources.userWorkspace.id, + createdAt: new Date(), + name: "namespace", + }; + await h.db.primary.insert(schema.ratelimitNamespaces).values(namespace); + + const identifier = randomUUID(); - const identifier = randomUUID(); + const root = await h.createRootKey(["ratelimit.*.limit"]); - const root = await h.createRootKey(["ratelimit.*.limit"]); + const results = await loadTest({ + rps, + seconds, + fn: () => + h.post({ + url: "/v1/ratelimits.limit", + headers: { + "Content-Type": "application/json", + Authorization: `Bearer ${root.key}`, + }, + body: { + identifier, + async: false, + namespace: namespace.name, + limit, + duration, + }, + }), + }); + t.expect(results.length).toBe(rps * seconds); + const passed = results.reduce((sum, res) => { + return res.body.success ? sum + 1 : sum; + }, 0); - const results = await loadTest({ - rps, - seconds, - fn: () => - h.post({ - url: "/v1/ratelimits.limit", - headers: { - "Content-Type": "application/json", - Authorization: `Bearer ${root.key}`, - }, - body: { - identifier, - async: false, - namespace: namespace.name, - limit, - duration, - }, - }), - }); - t.expect(results.length).toBe(rps * seconds); - const passed = results.reduce((sum, res) => { - return res.body.success ? sum + 1 : sum; - }, 0); - console.info({ name, passed }); - t.expect(passed).toBeGreaterThanOrEqual(expected.min); - t.expect(passed).toBeLessThanOrEqual(expected.max); - }, - ); + const exactLimit = (limit / (duration / 1000)) * seconds; + const upperLimit = Math.round(exactLimit * 1.25); + const lowerLimit = Math.round(exactLimit * 0.9); + console.info({ name, passed, exactLimit, upperLimit, lowerLimit }); + t.expect(passed).toBeGreaterThanOrEqual(lowerLimit); + t.expect(passed).toBeLessThanOrEqual(upperLimit); + }); }