import { Ratelimit } from "@upstash/ratelimit"; import type { RedisWithClusterOptions } from "~/redis.server"; import { createRedisClient } from "~/redis.server"; import { logger } from "./logger.server"; type Options = { redis?: RedisWithClusterOptions; redisClient?: RateLimiterRedisClient; keyPrefix: string; limiter: Limiter; logSuccess?: boolean; logFailure?: boolean; }; export type Limiter = ConstructorParameters[0]["limiter"]; export type Duration = Parameters[1]; export type RateLimitResponse = Awaited>; export type RateLimiterRedisClient = ConstructorParameters[0]["redis"]; export class RateLimiter { #ratelimit: Ratelimit; constructor(private readonly options: Options) { const { redis, redisClient, keyPrefix, limiter } = options; const prefix = `ratelimit:${keyPrefix}`; const resolvedRedisClient = redisClient ?? (redis ? createRedisRateLimitClient(redis) : undefined); if (!resolvedRedisClient) { throw new Error("RateLimiter requires either redis or redisClient options"); } this.#ratelimit = new Ratelimit({ redis: resolvedRedisClient, limiter, ephemeralCache: new Map(), analytics: false, prefix, }); } async limit(identifier: string, rate = 1): Promise { const result = this.#ratelimit.limit(identifier, { rate }); const { success, limit, reset, remaining } = await result; if (success && this.options.logSuccess) { logger.info(`RateLimiter (${this.options.keyPrefix}): under rate limit`, { limit, reset, remaining, identifier, }); } //log these by default if (!success && this.options.logFailure !== false) { logger.info(`RateLimiter (${this.options.keyPrefix}): rate limit exceeded`, { limit, reset, remaining, identifier, }); } return result; } } export function createRedisRateLimitClient( redisOptions: RedisWithClusterOptions ): RateLimiterRedisClient { const redis = createRedisClient("trigger:rateLimiter", redisOptions); return { sadd: async (key: string, ...members: TData[]): Promise => { return redis.sadd(key, members as (string | number | Buffer)[]); }, hset: ( key: string, obj: { [key: string]: TValue; } ): Promise => { return redis.hset(key, obj); }, eval: ( ...args: [script: string, keys: string[], args: TArgs] ): Promise => { const script = args[0]; const keys = args[1]; const argsArray = args[2]; return redis.eval( script, keys.length, ...keys, ...(argsArray as (string | Buffer | number)[]) ) as Promise; }, }; }