wehub-resource-sync
101 lines
2.9 KiB
TypeScript
101 lines
2.9 KiB
TypeScript
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<typeof Ratelimit>[0]["limiter"];
|
|
export type Duration = Parameters<typeof Ratelimit.slidingWindow>[1];
|
|
export type RateLimitResponse = Awaited<ReturnType<Ratelimit["limit"]>>;
|
|
export type RateLimiterRedisClient = ConstructorParameters<typeof Ratelimit>[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<RateLimitResponse> {
|
|
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 <TData>(key: string, ...members: TData[]): Promise<number> => {
|
|
return redis.sadd(key, members as (string | number | Buffer)[]);
|
|
},
|
|
hset: <TValue>(
|
|
key: string,
|
|
obj: {
|
|
[key: string]: TValue;
|
|
}
|
|
): Promise<number> => {
|
|
return redis.hset(key, obj);
|
|
},
|
|
eval: <TArgs extends unknown[], TData = unknown>(
|
|
...args: [script: string, keys: string[], args: TArgs]
|
|
): Promise<TData> => {
|
|
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<TData>;
|
|
},
|
|
};
|
|
}
|