From 5acc368ef4595dd69d5538cccab1c01f99e83cf3 Mon Sep 17 00:00:00 2001 From: Victor Navarro Date: Wed, 27 May 2026 11:27:06 +0200 Subject: [PATCH] perf: use redis for api key rate limit (#29242) --- .../app/src/routes/zen/util/ipRateLimiter.ts | 43 ++----------------- .../app/src/routes/zen/util/keyRateLimiter.ts | 23 ++++------ 2 files changed, 12 insertions(+), 54 deletions(-) diff --git a/packages/console/app/src/routes/zen/util/ipRateLimiter.ts b/packages/console/app/src/routes/zen/util/ipRateLimiter.ts index 81f73a4e5a..7461fa6313 100644 --- a/packages/console/app/src/routes/zen/util/ipRateLimiter.ts +++ b/packages/console/app/src/routes/zen/util/ipRateLimiter.ts @@ -1,5 +1,3 @@ -import { Database, eq, and, sql, inArray } from "@opencode-ai/console-core/drizzle/index.js" -import { IpRateLimitTable } from "@opencode-ai/console-core/schema/ip.sql.js" import { FreeUsageLimitError } from "./error" import { logger } from "./logger" import { buildRateLimitKey, getRedis } from "./redis" @@ -22,7 +20,6 @@ export function createRateLimiter(modelId: string, rateLimit: number | undefined const ip = !rawIp.length ? "unknown" : rawIp const now = Date.now() - const lifetimeInterval = "" const dailyInterval = rateLimit ? `${buildYYYYMMDD(now)}${modelId.substring(0, 2)}` : buildYYYYMMDD(now) const retryAfter = getRetryAfterDay(now) const redis = getRedis() @@ -32,33 +29,12 @@ export function createRateLimiter(modelId: string, rateLimit: number | undefined return { check: async () => { - const [counts, rows] = await Promise.all([ - redis.mget<(string | number | null)[]>(isDefaultModel ? [lifetimeKey, dailyKey] : [dailyKey]).catch(() => []), - Database.use((tx) => - tx - .select({ interval: IpRateLimitTable.interval, count: IpRateLimitTable.count }) - .from(IpRateLimitTable) - .where( - and( - eq(IpRateLimitTable.ip, ip), - isDefaultModel - ? inArray(IpRateLimitTable.interval, [lifetimeInterval, dailyInterval]) - : inArray(IpRateLimitTable.interval, [dailyInterval]), - ), - ), - ), - ]) - const redisLifetimeCount = isDefaultModel ? Number(counts[0] ?? 0) : 0 - const redisDailyCount = Number(counts[isDefaultModel ? 1 : 0] ?? 0) - const databaseLifetimeCount = rows.find((r) => r.interval === lifetimeInterval)?.count ?? 0 - const databaseDailyCount = rows.find((r) => r.interval === dailyInterval)?.count ?? 0 - const lifetimeCount = Math.max(redisLifetimeCount, databaseLifetimeCount) - const dailyCount = Math.max(redisDailyCount, databaseDailyCount) + const counts = await redis.mget<(string | number | null)[]>(isDefaultModel ? [lifetimeKey, dailyKey] : [dailyKey]) + const lifetimeCount = isDefaultModel ? Number(counts[0] ?? 0) : 0 + const dailyCount = Number(counts[isDefaultModel ? 1 : 0] ?? 0) logger.debug(`rate limit lifetime: ${lifetimeCount}, daily: ${dailyCount}`) isNew = isDefaultModel && lifetimeCount < dailyLimit * 7 - if (isDefaultModel && databaseLifetimeCount > redisLifetimeCount) - await redis.set(lifetimeKey, databaseLifetimeCount).catch(() => {}) if ((isNew && dailyCount >= dailyLimit * 2) || (!isNew && dailyCount >= dailyLimit)) throw new FreeUsageLimitError(dict["zen.api.error.rateLimitExceeded"], retryAfter) @@ -68,18 +44,7 @@ export function createRateLimiter(modelId: string, rateLimit: number | undefined pipeline.incr(dailyKey) pipeline.expire(dailyKey, retryAfter) if (isNew) pipeline.incr(lifetimeKey) - await Promise.all([ - pipeline.exec().catch(() => {}), - Database.use((tx) => - tx - .insert(IpRateLimitTable) - .values([ - { ip, interval: dailyInterval, count: 1 }, - ...(isNew ? [{ ip, interval: lifetimeInterval, count: 1 }] : []), - ]) - .onDuplicateKeyUpdate({ set: { count: sql`${IpRateLimitTable.count} + 1` } }), - ), - ]) + await pipeline.exec() }, } } diff --git a/packages/console/app/src/routes/zen/util/keyRateLimiter.ts b/packages/console/app/src/routes/zen/util/keyRateLimiter.ts index c8cc413a86..fdf5925299 100644 --- a/packages/console/app/src/routes/zen/util/keyRateLimiter.ts +++ b/packages/console/app/src/routes/zen/util/keyRateLimiter.ts @@ -1,6 +1,5 @@ -import { Database, eq, and, sql } from "@opencode-ai/console-core/drizzle/index.js" -import { KeyRateLimitTable } from "@opencode-ai/console-core/schema/ip.sql.js" import { RateLimitError } from "./error" +import { buildRateLimitKey, getRedis } from "./redis" import { i18n } from "~/i18n" import { localeFromRequest } from "~/lib/language" @@ -19,26 +18,20 @@ export function createRateLimiter( .replace(/[^0-9]/g, "") .substring(0, 12) const interval = `${modelId.substring(0, 27)}-${yyyyMMddHHmm}` + const redis = getRedis() + const key = buildRateLimitKey("key", zenApiKey, interval) return { check: async () => { - const rows = await Database.use((tx) => - tx - .select({ interval: KeyRateLimitTable.interval, count: KeyRateLimitTable.count }) - .from(KeyRateLimitTable) - .where(and(eq(KeyRateLimitTable.key, zenApiKey), eq(KeyRateLimitTable.interval, interval))), - ).then((rows) => rows[0]) - const count = rows?.count ?? 0 + const count = Number((await redis.mget<(string | number | null)[]>([key]))[0] ?? 0) if (count >= LIMIT) throw new RateLimitError(dict["zen.api.error.rateLimitExceeded"], 60) }, track: async () => { - await Database.use((tx) => - tx - .insert(KeyRateLimitTable) - .values({ key: zenApiKey, interval, count: 1 }) - .onDuplicateKeyUpdate({ set: { count: sql`${KeyRateLimitTable.count} + 1` } }), - ) + const pipeline = redis.pipeline() + pipeline.incr(key) + pipeline.expire(key, 60) + await pipeline.exec() }, } }