UNPKG

@convex-dev/rate-limiter

Version:

A rate limiter component for Convex. Define and use application-layer rate limits. Type-safe, transactional, fair, safe, and configurable sharding to scale.

201 lines 8.46 kB
import { mutationGeneric, queryGeneric, } from "convex/server"; import { ConvexError, v } from "convex/values"; import { getValueArgs, getValueReturns } from "../shared.js"; export { calculateRateLimit } from "../shared.js"; export const SECOND = 1000; export const MINUTE = 60 * SECOND; export const HOUR = 60 * MINUTE; export const DAY = 24 * HOUR; export const WEEK = 7 * DAY; export function isRateLimitError(error) { return (error instanceof ConvexError && error.data["kind"] === "RateLimited"); } /** * Define rate limits for a set of named rate limits. * e.g. * ```ts * import { RateLimiter } from "@convex-dev/rate-limiter"; * import { components } from "./_generated/api.js"; * * const rateLimiter = new RateLimiter(components.rateLimiter, { * // A per-user limit, allowing one every ~6 seconds. * // Allows up to 3 in quick succession if they haven't sent many recently. * sendMessage: { kind: "token bucket", rate: 10, period: MINUTE, capacity: 3 }, * // One global / singleton rate limit * freeTrialSignUp: { kind: "fixed window", rate: 100, period: HOUR }, * }); * //... elsewhere * await rateLimiter.limit(ctx, "sendMessage", { key: ctx.userId, throws: true }); * ``` * * @param component The rate limiter component. Like `components.rateLimiter`. * Imported like `import { components } from "./_generated/api.js";` * @param limits The rate limits to define. The key is the name of the rate limit. * See {@link RateLimitConfig} for more information. * @returns A rate limiter that has types based on the provided limits. * If you provide a different name, you will need to provide the config inline. */ export class RateLimiter { component; limits; constructor(component, limits) { this.component = component; this.limits = limits; } /** * Check a rate limit. * This function will check the rate limit and return whether the request is * allowed, and if not, when it could be retried. * Unlike {@link limit}, this function does not consume any tokens. * * @param ctx The ctx object from a query or mutation, including runQuery. * @param name The name of the rate limit. * @param options The rate limit arguments. `config` is required if the rate * limit was not defined in {@link RateLimiter}. See {@link RateLimitArgs}. * @returns `{ ok, retryAfter }`: `ok` is true if the rate limit is not exceeded. * `retryAfter` is the duration in milliseconds when retrying could succeed. * If `reserve` is true, `ok` is true if there's enough capacity including * reservation. If there is a maxiumum reservation limit, `ok` will be false * when it is exceeded. When `ok` is true and `retryAfter` is defined, it is * the duration you must wait before executing the work. * e.g.: * ```ts * if (status.retryAfter) { * await ctx.scheduler.runAfter(retryAfter, ...) * ``` */ async check(ctx, name, ...options) { return ctx.runQuery(this.component.lib.checkRateLimit, { ...options[0], name, config: this.getConfig(options[0], name), }); } /** * Rate limit a request. * This function will check the rate limit and return whether the request is * allowed, and if not, when it could be retried. * * @param ctx The ctx object from a mutation, including runMutation. * @param name The name of the rate limit. * @param options The rate limit arguments. `config` is required if the rate * limit was not defined in {@link RateLimiter}. See {@link RateLimitArgs}. * @returns `{ ok, retryAfter }`: `ok` is true if the rate limit is not exceeded. * `retryAfter` is the duration in milliseconds when retrying could succeed. * If `reserve` is true, `ok` is true if there's enough capacity including * reservation. If there is a maxiumum reservation limit, `ok` will be false * when it is exceeded. When `ok` is true and `retryAfter` is defined, it is * the duration you must wait before executing the work. * e.g.: * ```ts * if (status.retryAfter) { * await ctx.scheduler.runAfter(retryAfter, ...) * ``` */ async limit(ctx, name, ...options) { return ctx.runMutation(this.component.lib.rateLimit, { ...options[0], name, config: this.getConfig(options[0], name), }); } /** * Reset a rate limit. This will remove the rate limit from the database. * The next request will start fresh. * Note: In the case of a fixed window without a specified `start`, * the new window will be a random time. * @param ctx The ctx object from a mutation, including runMutation. * @param name The name of the rate limit to reset, including all shards. * @param key If a key is provided, it will reset the rate limit for that key. * If not, it will reset the rate limit for the shared value. */ async reset({ runMutation }, name, args) { await runMutation(this.component.lib.resetRateLimit, { ...(args ?? null), name, }); } /** * Get the current value and metadata of a rate limit. * This function returns the current token utilization data without consuming any tokens. * * @param ctx The ctx object from a query, including runQuery. * @param name The name of the rate limit. * @param options The rate limit arguments. `config` is required if the rate * limit was not defined in {@link RateLimiter}. See {@link RateLimitArgs}. * @returns An object containing the current value, timestamp, window start time (for fixed window), * and the rate limit configuration. */ async getValue(ctx, name, ...options) { return ctx.runQuery(this.component.lib.getValue, { ...options[0], name, config: this.getConfig(options[0], name), }); } /** * Creates a public query that can be exported from your API that returns the * current value of a rate limit. * This is a convenience function to re-export the query for client use. * * @param name The name of the rate limit. * @returns An object containing a getRateLimit function that can be exported. * * Example: * ```ts * // In your API file: * export const getRateLimit = rateLimiter.getValueQuery("myLimit"); * * // In your client: * const { status, getValue, retryAt } = useRateLimit(api.getRateLimit, 10); * ``` */ hookAPI(name, ...options) { return { getRateLimit: queryGeneric({ args: getValueArgs, returns: getValueReturns, handler: async (ctx, args) => { const finalName = args.name ?? name; const { key: keyOrFn, ...rest } = options[0] ?? {}; let key; if (args.key && !keyOrFn) { throw new Error("To allow client-provided key, provide a `key` function in the hook options."); } if (typeof keyOrFn === "function") { key = await keyOrFn(ctx, args.key); } else if (keyOrFn !== undefined) { key = keyOrFn; } return ctx.runQuery(this.component.lib.getValue, { ...rest, ...args, key, name: finalName, config: args.config ?? this.getConfig(options[0], finalName), }); }, }), getServerTime: mutationGeneric({ args: {}, returns: v.number(), handler: async () => { return Date.now(); }, }), }; } getConfig(args, name) { const config = (args && "config" in args && args.config) || (this.limits && this.limits[name]); if (!config) { throw new Error(`Rate limit ${name} not defined. ` + `You must provide a config inline or define it in the constructor.`); } return config; } } export default RateLimiter; //# sourceMappingURL=index.js.map