perf: use redis for api key rate limit (anomalyco#29242)

Author: vimtor

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()
     },
   }
 }

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()
     },
   }
 }