From 6fef75d2bb125c488048af95926752262fada80e Mon Sep 17 00:00:00 2001 From: Nathan Flurry Date: Wed, 26 Jun 2024 23:57:19 -0700 Subject: [PATCH] refactor(rate_limit): migrate to actors --- modules/rate_limit/actors/limiter.ts | 46 ++++++++++++++ modules/rate_limit/module.json | 57 +++++++++-------- modules/rate_limit/scripts/throttle.ts | 63 +++++-------------- modules/rate_limit/scripts/throttle_public.ts | 6 +- 4 files changed, 93 insertions(+), 79 deletions(-) create mode 100644 modules/rate_limit/actors/limiter.ts diff --git a/modules/rate_limit/actors/limiter.ts b/modules/rate_limit/actors/limiter.ts new file mode 100644 index 00000000..aea0f643 --- /dev/null +++ b/modules/rate_limit/actors/limiter.ts @@ -0,0 +1,46 @@ +import { ActorBase, ActorContext } from "../module.gen.ts"; + +type Input = undefined; + +interface State { + tokens: number; + lastRefillTimestamp: number; +} + +export interface ThrottleRequest { + requests: number; + period: number; +} + +export interface ThrottleResponse { + success: boolean; + refillAt: number; +} + +export class Actor extends ActorBase { + public initialize(): State { + // Will refill on first call of `throttle` + return { + tokens: 0, + lastRefillTimestamp: 0, + }; + } + + throttle(_ctx: ActorContext, req: ThrottleRequest): ThrottleResponse { + // Reset bucket + const now = Date.now(); + if (now > this.state.lastRefillTimestamp + req.period * 1000) { + this.state.tokens = req.requests; + this.state.lastRefillTimestamp = now; + } + + // Attempt to consume token + const success = this.state.tokens >= 1; + if (success) { + this.state.tokens -= 1; + } + + const refillAt = Math.ceil((1 - this.state.tokens) * (req.period / req.requests)); + return { success, refillAt }; + } +} diff --git a/modules/rate_limit/module.json b/modules/rate_limit/module.json index 12c7a4a3..3210575e 100644 --- a/modules/rate_limit/module.json +++ b/modules/rate_limit/module.json @@ -1,29 +1,32 @@ { - "name": "Rate Limit", - "description": "Prevent abuse by limiting request rate.", - "icon": "gauge-circle-minus", - "tags": [ - "core", - "utility" - ], - "authors": [ - "rivet-gg", - "NathanFlurry" - ], - "status": "stable", - "scripts": { - "throttle": { - "name": "Throttle", - "description": "Limit the amount of times an request can be made by a given key." - }, - "throttle_public": { - "name": "Throttle Public", - "description": "Limit the amount of times a public request can be made by a given key. This will rate limit based off the user's IP address." - } - }, - "errors": { - "rate_limit_exceeded": { - "name": "Rate Limit Exceeded" - } - } + "name": "Rate Limit", + "description": "Prevent abuse by limiting request rate.", + "icon": "gauge-circle-minus", + "tags": [ + "core", + "utility" + ], + "authors": [ + "rivet-gg", + "NathanFlurry" + ], + "status": "stable", + "scripts": { + "throttle": { + "name": "Throttle", + "description": "Limit the amount of times an request can be made by a given key." + }, + "throttle_public": { + "name": "Throttle Public", + "description": "Limit the amount of times a public request can be made by a given key. This will rate limit based off the user's IP address." + } + }, + "errors": { + "rate_limit_exceeded": { + "name": "Rate Limit Exceeded" + } + }, + "actors": { + "limiter": {} + } } \ No newline at end of file diff --git a/modules/rate_limit/scripts/throttle.ts b/modules/rate_limit/scripts/throttle.ts index 8356230e..fa36cac3 100644 --- a/modules/rate_limit/scripts/throttle.ts +++ b/modules/rate_limit/scripts/throttle.ts @@ -1,3 +1,5 @@ +import { assert } from "https://deno.land/std@0.208.0/assert/mod.ts"; +import { ThrottleRequest, ThrottleResponse } from "../actors/limiter.ts"; import { RuntimeError, ScriptContext } from "../module.gen.ts"; export interface Request { @@ -28,58 +30,23 @@ export async function run( ctx: ScriptContext, req: Request, ): Promise { - interface TokenBucket { - tokens: number; - lastRefill: Date; - } + assert(req.requests > 0); + assert(req.period > 0); + + // Create key + const key = `${JSON.stringify(req.type)}.${JSON.stringify(req.key)}`; - // Update the token bucket - // - // `TokenBucket` is an unlogged table which are significantly faster to - // write to than regular tables, but are not durable. This is important - // because this script will be called on every request. - const rows = await ctx.db.$queryRawUnsafe( - ` - WITH - "UpdateBucket" AS ( - UPDATE "${ctx.dbSchema}"."TokenBuckets" b - SET - "tokens" = CASE - -- Reset the bucket and consume 1 token - WHEN now() > b."lastRefill" + make_interval(secs => $4) THEN $3 - 1 - -- Consume 1 token - ELSE b.tokens - 1 - END, - "lastRefill" = CASE - WHEN now() > b."lastRefill" + make_interval(secs => $4) THEN now() - ELSE b."lastRefill" - END - WHERE b."type" = $1 AND b."key" = $2 - RETURNING b."tokens", b."lastRefill" - ), - inserted AS ( - INSERT INTO "${ctx.dbSchema}"."TokenBuckets" ("type", "key", "tokens", "lastRefill") - SELECT $1, $2, $3 - 1, now() - WHERE NOT EXISTS (SELECT 1 FROM "UpdateBucket") - RETURNING "tokens", "lastRefill" - ) - SELECT * FROM "UpdateBucket" - UNION ALL - SELECT * FROM inserted; - `, - req.type, - req.key, - req.requests, - req.period, - ); - const { tokens, lastRefill } = rows[0]; + // Throttle request + const res = await ctx.actors.limiter.getOrCreateAndCall(key, undefined, "throttle", { + requests: req.requests, + period: req.period, + }); - // If the bucket is empty, throw an error - if (tokens < 0) { + // Check if allowed + if (!res.success) { throw new RuntimeError("RATE_LIMIT_EXCEEDED", { meta: { - retryAfter: new Date(lastRefill.getTime() + req.period * 1000) - .toUTCString(), + retryAfter: new Date(res.refillAt).toUTCString(), }, }); } diff --git a/modules/rate_limit/scripts/throttle_public.ts b/modules/rate_limit/scripts/throttle_public.ts index bbe2c7ec..d3e4356f 100644 --- a/modules/rate_limit/scripts/throttle_public.ts +++ b/modules/rate_limit/scripts/throttle_public.ts @@ -20,9 +20,6 @@ export async function run( ctx: ScriptContext, req: Request, ): Promise { - const requests = req.requests || 20; - const period = req.period || 300; - // Find the IP address of the client let key: string | undefined; for (const entry of ctx.trace.entries) { @@ -32,7 +29,8 @@ export async function run( } } - // If no IP address, this request is not coming from a client + // If no IP address, this request is not coming from a client and should not + // be throttled if (!key) { return {}; }