wehub-resource-sync
124 lines
3.9 KiB
TypeScript
124 lines
3.9 KiB
TypeScript
import { z } from "zod";
|
|
import { env } from "~/env.server";
|
|
import {
|
|
createLimiterFromConfig,
|
|
RateLimiterConfig,
|
|
} from "~/services/authorizationRateLimitMiddleware.server";
|
|
import type { Duration } from "~/services/rateLimiter.server";
|
|
import { createRedisRateLimitClient, RateLimiter } from "~/services/rateLimiter.server";
|
|
import { singleton } from "~/utils/singleton";
|
|
|
|
const BatchLimitsConfig = z.object({
|
|
processingConcurrency: z.number().int().default(env.BATCH_CONCURRENCY_LIMIT_DEFAULT),
|
|
});
|
|
|
|
/**
|
|
* Batch limits configuration for a plan type
|
|
*/
|
|
export type BatchLimitsConfig = z.infer<typeof BatchLimitsConfig>;
|
|
|
|
const batchLimitsRedisClient = singleton("batchLimitsRedisClient", createBatchLimitsRedisClient);
|
|
|
|
function createBatchLimitsRedisClient() {
|
|
const redisClient = createRedisRateLimitClient({
|
|
port: env.RATE_LIMIT_REDIS_PORT,
|
|
host: env.RATE_LIMIT_REDIS_HOST,
|
|
username: env.RATE_LIMIT_REDIS_USERNAME,
|
|
password: env.RATE_LIMIT_REDIS_PASSWORD,
|
|
tlsDisabled: env.RATE_LIMIT_REDIS_TLS_DISABLED === "true",
|
|
clusterMode: env.RATE_LIMIT_REDIS_CLUSTER_MODE_ENABLED === "1",
|
|
});
|
|
|
|
return redisClient;
|
|
}
|
|
|
|
// Just the org fields this module reads. Compatible with both the full
|
|
// Prisma `Organization` payload and the slim `AuthenticatedEnvironment`
|
|
// `["organization"]` shape (when passed `batchRateLimitConfig` /
|
|
// `batchQueueConcurrencyConfig` as `unknown`).
|
|
type OrganizationForBatchLimits = {
|
|
batchRateLimitConfig?: unknown;
|
|
batchQueueConcurrencyConfig?: unknown;
|
|
};
|
|
|
|
function createOrganizationRateLimiter(organization: OrganizationForBatchLimits): RateLimiter {
|
|
const limiterConfig = resolveBatchRateLimitConfig(organization.batchRateLimitConfig);
|
|
|
|
const limiter = createLimiterFromConfig(limiterConfig);
|
|
|
|
return new RateLimiter({
|
|
redisClient: batchLimitsRedisClient,
|
|
keyPrefix: "ratelimit:batch",
|
|
limiter,
|
|
logSuccess: false,
|
|
logFailure: true,
|
|
});
|
|
}
|
|
|
|
function resolveBatchRateLimitConfig(batchRateLimitConfig?: unknown): RateLimiterConfig {
|
|
const defaultRateLimiterConfig: RateLimiterConfig = {
|
|
type: "tokenBucket",
|
|
refillRate: env.BATCH_RATE_LIMIT_REFILL_RATE,
|
|
interval: env.BATCH_RATE_LIMIT_REFILL_INTERVAL as Duration,
|
|
maxTokens: env.BATCH_RATE_LIMIT_MAX,
|
|
};
|
|
|
|
if (!batchRateLimitConfig) {
|
|
return defaultRateLimiterConfig;
|
|
}
|
|
|
|
const parsedBatchRateLimitConfig = RateLimiterConfig.safeParse(batchRateLimitConfig);
|
|
|
|
if (!parsedBatchRateLimitConfig.success) {
|
|
return defaultRateLimiterConfig;
|
|
}
|
|
|
|
return parsedBatchRateLimitConfig.data;
|
|
}
|
|
|
|
/**
|
|
* Get the rate limiter and limits for an organization.
|
|
* Internally looks up the plan type, but doesn't expose it to callers.
|
|
*/
|
|
export async function getBatchLimits(
|
|
organization: OrganizationForBatchLimits
|
|
): Promise<{ rateLimiter: RateLimiter; config: BatchLimitsConfig }> {
|
|
const rateLimiter = createOrganizationRateLimiter(organization);
|
|
const config = resolveBatchLimitsConfig(organization.batchQueueConcurrencyConfig);
|
|
return { rateLimiter, config };
|
|
}
|
|
|
|
function resolveBatchLimitsConfig(batchLimitsConfig?: unknown): BatchLimitsConfig {
|
|
const defaultLimitsConfig: BatchLimitsConfig = {
|
|
processingConcurrency: env.BATCH_CONCURRENCY_LIMIT_DEFAULT,
|
|
};
|
|
|
|
if (!batchLimitsConfig) {
|
|
return defaultLimitsConfig;
|
|
}
|
|
|
|
const parsedBatchLimitsConfig = BatchLimitsConfig.safeParse(batchLimitsConfig);
|
|
|
|
if (!parsedBatchLimitsConfig.success) {
|
|
return defaultLimitsConfig;
|
|
}
|
|
|
|
return parsedBatchLimitsConfig.data;
|
|
}
|
|
|
|
/**
|
|
* Error thrown when batch rate limit is exceeded.
|
|
* Contains information for constructing a proper 429 response.
|
|
*/
|
|
export class BatchRateLimitExceededError extends Error {
|
|
constructor(
|
|
public readonly limit: number,
|
|
public readonly remaining: number,
|
|
public readonly resetAt: Date,
|
|
public readonly itemCount: number
|
|
) {
|
|
super(`Batch rate limit exceeded. Limit resets at ${resetAt.toISOString()}`);
|
|
this.name = "BatchRateLimitExceededError";
|
|
}
|
|
}
|