chore: re-enable rate_limit module (#107)

rivet-gg · Jun 10, 2024 · c7a6789 · c7a6789
1 parent 1e2a0ca
commit c7a6789
Show file tree

Hide file tree

Showing 4 changed files with 1,385 additions and 50 deletions.
diff --git a/modules/rate_limit/scripts/throttle.ts b/modules/rate_limit/scripts/throttle.ts
@@ -28,56 +28,56 @@ export async function run(
 	ctx: ScriptContext,
 	req: Request,
 ): Promise<Response> {
-	// interface TokenBucket {
-	// 	tokens: number;
-	// 	lastRefill: Date;
-	// }
-	//
-	// // 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.$queryRaw<TokenBucket[]>`
-	// 	WITH updated_bucket AS (
-	// 		UPDATE "TokenBuckets" b
-	// 		SET
-	// 			"tokens" = CASE
-	// 					-- Reset the bucket and consume 1 token
-	// 					WHEN now() > b."lastRefill" + make_interval(secs => ${req.period}) THEN ${
-	// 	req.requests - 1
-	// }
-	// 					-- Consume 1 token
-	// 					ELSE b.tokens - 1
-	// 				END,
-	// 			"lastRefill" = CASE
-	// 					WHEN now() > b."lastRefill" + make_interval(secs => ${req.period}) THEN now()
-	// 					ELSE b."lastRefill"
-	// 				END
-	// 		WHERE b."type" = ${req.type} AND b."key" = ${req.key}
-	// 		RETURNING b."tokens", b."lastRefill"
-	// 	),
-	// 	inserted AS (
-	// 		INSERT INTO "TokenBuckets" ("type", "key", "tokens", "lastRefill")
-	// 		SELECT ${req.type}, ${req.key}, ${req.requests - 1}, now()
-	// 		WHERE NOT EXISTS (SELECT 1 FROM updated_bucket)
-	// 		RETURNING "tokens", "lastRefill"
-	// 	)
-	// 	SELECT * FROM updated_bucket
-	// 	UNION ALL
-	// 	SELECT * FROM inserted;
-	// `;
-	// const { tokens, lastRefill } = rows[0];
+	interface TokenBucket {
+		tokens: number;
+		lastRefill: Date;
+	}
+
+	// Update the token bucket
 	//
-	// // If the bucket is empty, throw an error
-	// if (tokens < 0) {
-	// 	throw new RuntimeError("RATE_LIMIT_EXCEEDED", {
-	// 		meta: {
-	// 			retryAfter: new Date(lastRefill.getTime() + req.period * 1000)
-	// 				.toUTCString(),
-	// 		},
-	// 	});
-	// }
+	// `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.$queryRaw<TokenBucket[]>`
+		WITH updated_bucket AS (
+			UPDATE "TokenBuckets" b
+			SET
+				"tokens" = CASE
+						-- Reset the bucket and consume 1 token
+						WHEN now() > b."lastRefill" + make_interval(secs => ${req.period}) THEN ${
+		req.requests - 1
+	}
+						-- Consume 1 token
+						ELSE b.tokens - 1
+					END,
+				"lastRefill" = CASE
+						WHEN now() > b."lastRefill" + make_interval(secs => ${req.period}) THEN now()
+						ELSE b."lastRefill"
+					END
+			WHERE b."type" = ${req.type} AND b."key" = ${req.key}
+			RETURNING b."tokens", b."lastRefill"
+		),
+		inserted AS (
+			INSERT INTO "TokenBuckets" ("type", "key", "tokens", "lastRefill")
+			SELECT ${req.type}, ${req.key}, ${req.requests - 1}, now()
+			WHERE NOT EXISTS (SELECT 1 FROM updated_bucket)
+			RETURNING "tokens", "lastRefill"
+		)
+		SELECT * FROM updated_bucket
+		UNION ALL
+		SELECT * FROM inserted;
+	`;
+	const { tokens, lastRefill } = rows[0];
+
+	// If the bucket is empty, throw an error
+	if (tokens < 0) {
+		throw new RuntimeError("RATE_LIMIT_EXCEEDED", {
+			meta: {
+				retryAfter: new Date(lastRefill.getTime() + req.period * 1000)
+					.toUTCString(),
+			},
+		});
+	}
 
 	return {};
 }