refactor(rate_limit): migrate to actors (#116)

Author: NathanFlurry

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<undefined, State> {
+  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 };
+  }
+}

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": {}
+	}
 }

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<Response> {
-	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<TokenBucket[]>(
-    `
-		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<undefined, ThrottleRequest, ThrottleResponse>(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(),
 			},
 		});
 	}

modules/rate_limit/scripts/throttle_public.ts

@@ -20,9 +20,6 @@ export async function run(
 	ctx: ScriptContext,
 	req: Request,
 ): Promise<Response> {
-	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 {};
 	}