Skip to content
This repository was archived by the owner on Sep 17, 2024. It is now read-only.

chore: re-enable rate_limit module #107

Merged
merged 1 commit into from
Jun 10, 2024
Merged

chore: re-enable rate_limit module #107

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
98 changes: 49 additions & 49 deletions modules/rate_limit/scripts/throttle.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 {};
}
Loading