saas: async auth/rate-limit path for high-freq keys

src/api/searcher-serper.ts

@@ -19,9 +19,11 @@ import { AsyncLocalContext } from '../services/async-context';
 import { Context, Ctx, Method, Param, RPCReflect } from '../services/registry';
 import { OutputServerEventStream } from '../lib/transform-server-event-stream';
 import { JinaEmbeddingsAuthDTO } from '../dto/jina-embeddings-auth';
-import { InsufficientBalanceError } from '../services/errors';
+import { InsufficientBalanceError, RateLimitTriggeredError } from '../services/errors';
 import { SerperImageSearchResponse, SerperNewsSearchResponse, SerperSearchQueryParams, SerperSearchResponse, SerperWebSearchResponse, WORLD_COUNTRIES, WORLD_LANGUAGES } from '../shared/3rd-party/serper-search';
 import { toAsyncGenerator } from '../utils/misc';
+import type { JinaEmbeddingsTokenAccount } from '../shared/db/jina-embeddings-token-account';
+import { LRUCache } from 'lru-cache';
 
 const WORLD_COUNTRY_CODES = Object.keys(WORLD_COUNTRIES).map((x) => x.toLowerCase());
 
@@ -30,6 +32,11 @@ interface FormattedPage extends RealFormattedPage {
     date?: string;
 }
 
+type RateLimitCache = {
+    blockedUntil?: Date;
+    user?: JinaEmbeddingsTokenAccount;
+};
+
 @singleton()
 export class SearcherHost extends RPCHost {
     logger = this.globalLogger.child({ service: this.constructor.name });
@@ -42,6 +49,13 @@ export class SearcherHost extends RPCHost {
 
     targetResultCount = 5;
 
+    highFreqKeyCache = new LRUCache<string, RateLimitCache>({
+        max: 256,
+        ttl: 60 * 60 * 1000,
+        updateAgeOnGet: false,
+        updateAgeOnHas: false,
+    });
+
     constructor(
         protected globalLogger: GlobalLogger,
         protected rateLimitControl: RateLimitControl,
@@ -105,6 +119,14 @@ export class SearcherHost extends RPCHost {
         // Here we combine 'count' and 'num' to 'count' for the rest of the function.
         count = (num !== undefined ? num : count) ?? 10;
 
+        const authToken = auth.bearerToken;
+        let highFreqKey: RateLimitCache | undefined;
+        if (authToken && this.highFreqKeyCache.has(authToken)) {
+            highFreqKey = this.highFreqKeyCache.get(authToken)!;
+            auth.user = highFreqKey.user;
+            auth.uid = highFreqKey.user?.user_id;
+        }
+
         const uid = await auth.solveUID();
         // Return content by default
         const crawlWithoutContent = crawlerOptions.respondWith.includes('no-content');
@@ -134,6 +156,17 @@ export class SearcherHost extends RPCHost {
             throw new InsufficientBalanceError(`Account balance not enough to run this query, please recharge.`);
         }
 
+        if (highFreqKey?.blockedUntil) {
+            const now = new Date();
+            const blockedTimeRemaining = (highFreqKey.blockedUntil.valueOf() - now.valueOf());
+            if (blockedTimeRemaining > 0) {
+                throw RateLimitTriggeredError.from({
+                    message: `Per UID rate limit exceeded (async)`,
+                    retryAfter: Math.ceil(blockedTimeRemaining / 1000),
+                });
+            }
+        }
+
         const rateLimitPolicy = auth.getRateLimits(rpcReflect.name.toUpperCase()) || [
             parseInt(user.metadata?.speed_level) >= 2 ?
                 RateLimitDesc.from({
@@ -146,16 +179,75 @@ export class SearcherHost extends RPCHost {
                 })
         ];
 
-        const apiRoll = await this.rateLimitControl.simpleRPCUidBasedLimit(
+        const apiRollPromise = this.rateLimitControl.simpleRPCUidBasedLimit(
             rpcReflect, uid!, [rpcReflect.name.toUpperCase()],
             ...rateLimitPolicy
         );
 
-        rpcReflect.finally(() => {
+        if (!highFreqKey) {
+            // Normal path
+            await apiRollPromise;
+
+            if (rateLimitPolicy.some(
+                (x) => {
+                    const rpm = x.occurrence / (x.periodSeconds / 60);
+                    if (rpm >= 400) {
+                        return true;
+                    }
+
+                    return false;
+                })
+            ) {
+                this.highFreqKeyCache.set(auth.bearerToken!, {
+                    user,
+                });
+            }
+
+        } else {
+            // High freq key path
+            apiRollPromise.then(
+                // Rate limit not triggered, make sure not blocking.
+                () => {
+                    delete highFreqKey.blockedUntil;
+                },
+                // Rate limit triggered
+                (err) => {
+                    if (!(err instanceof RateLimitTriggeredError)) {
+                        return;
+                    }
+                    const now = Date.now();
+                    let tgtDate;
+                    if (err.retryAfter) {
+                        tgtDate = new Date(now + err.retryAfter * 1000);
+                    } else if (err.retryAfterDate) {
+                        tgtDate = err.retryAfterDate;
+                    }
+
+                    if (tgtDate) {
+                        const dt = tgtDate.valueOf() - now;
+                        highFreqKey.blockedUntil = tgtDate;
+                        setTimeout(() => {
+                            if (highFreqKey.blockedUntil === tgtDate) {
+                                delete highFreqKey.blockedUntil;
+                            }
+                        }, dt).unref();
+                    }
+                }
+            ).finally(async () => {
+                // Always asynchronously update user(wallet);
+                const user = await auth.getBrief().catch(() => undefined);
+                if (user) {
+                    highFreqKey.user = user;
+                }
+            });
+        }
+
+        rpcReflect.finally(async () => {
             if (chargeAmount) {
                 auth.reportUsage(chargeAmount, `reader-${rpcReflect.name}`).catch((err) => {
                     this.logger.warn(`Unable to report usage for ${uid}`, { err: marshalErrorLike(err) });
                 });
+                const apiRoll = await apiRollPromise;
                 apiRoll.chargeAmount = chargeAmount;
             }
         });

thinapps-shared

@@ -1 +1 @@
-Subproject commit 07d23193d85b1d3c8bbd5d0b024a6884ecfe17fd
+Subproject commit 492ed4cac38958c3715013da49a9e73b1d1ef8cb