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; 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"; } }