boki/stock-bot
1
0
Fork 0
Code Issues 1 Pull requests Projects Releases Packages Wiki Activity Actions

fixed up ratelimiting

Browse source
This commit is contained in:
Boki 2025-07-06 18:53:02 -04:00
parent a616c92656
commit a7146a3f57
15 changed files with 912 additions and 186 deletions
Download patch file Download diff file Expand all files Collapse all files

187
libs/core/queue/src/rate-limiter.ts
View file

@ -1,5 +1,5 @@
import { RateLimiterRedis, RateLimiterRes } from 'rate-limiter-flexible';
import type { RateLimitConfig as BaseRateLimitConfig, RateLimitRule } from './types';
import type { RateLimitWindow, RateLimitConfig, RateLimitRule } from './types';
// Logger interface for type safety
interface Logger {
@ -9,13 +9,16 @@ interface Logger {
debug(message: string, meta?: Record<string, unknown>): void;
}
// Extend the base config to add rate-limiter specific fields
export interface RateLimitConfig extends BaseRateLimitConfig {
keyPrefix?: string;
interface RateLimiterEntry {
windows: Array<{
limiter: RateLimiterRedis;
window: RateLimitWindow;
}>;
cost: number;
}
export class QueueRateLimiter {
private limiters = new Map<string, RateLimiterRedis>();
private limiters = new Map<string, RateLimiterEntry>();
private rules: RateLimitRule[] = [];
private readonly logger: Logger;
@ -33,23 +36,33 @@ export class QueueRateLimiter {
this.rules.push(rule);
const key = this.getRuleKey(rule.level, rule.queueName, rule.handler, rule.operation);
const limiter = new RateLimiterRedis({
storeClient: this.redisClient,
keyPrefix: `rl:${key}`,
points: rule.config.points,
duration: rule.config.duration,
blockDuration: rule.config.blockDuration || 0,
// Extract limits and cost from config
const limits = rule.config.limits || [];
const cost = rule.config.cost || 1;
// Create rate limiters for each window
const windows = limits.map((window, index) => {
const limiter = new RateLimiterRedis({
storeClient: this.redisClient,
keyPrefix: `rl:${key}:${index}`,
points: window.points,
duration: window.duration,
blockDuration: window.blockDuration || 0,
});
return { limiter, window };
});
this.limiters.set(key, limiter);
this.limiters.set(key, { windows, cost });
this.logger.info('Rate limit rule added', {
level: rule.level,
queueName: rule.queueName,
handler: rule.handler,
operation: rule.operation,
points: rule.config.points,
duration: rule.config.duration,
windows: limits.length,
cost,
});
}
@ -57,6 +70,7 @@ export class QueueRateLimiter {
* Check if a job can be processed based on rate limits
* Uses hierarchical precedence: operation > handler > queue > global
* The most specific matching rule takes precedence
* Returns the longest wait time if multiple windows are hit
*/
async checkLimit(
queueName: string,
@ -67,6 +81,7 @@ export class QueueRateLimiter {
retryAfter?: number;
remainingPoints?: number;
appliedRule?: RateLimitRule;
cost?: number;
}> {
const applicableRule = this.getMostSpecificRule(queueName, handler, operation);
@ -80,28 +95,74 @@ export class QueueRateLimiter {
applicableRule.handler,
applicableRule.operation
);
const limiter = this.limiters.get(key);
const limiterEntry = this.limiters.get(key);
if (!limiter) {
if (!limiterEntry || limiterEntry.windows.length === 0) {
this.logger.warn('Rate limiter not found for rule', { key, rule: applicableRule });
return { allowed: true };
}
try {
const result = await this.consumePoint(
limiter,
this.getConsumerKey(queueName, handler, operation)
);
const consumerKey = this.getConsumerKey(queueName, handler, operation);
const cost = limiterEntry.cost;
// Check all windows and collect results
const windowResults = await Promise.all(
limiterEntry.windows.map(async ({ limiter, window }) => {
try {
// Try to consume points for this window
const result = await limiter.consume(consumerKey, cost);
return {
allowed: true,
remainingPoints: result.remainingPoints,
retryAfter: 0,
};
} catch (rejRes) {
if (rejRes instanceof RateLimiterRes) {
return {
allowed: false,
remainingPoints: rejRes.remainingPoints,
retryAfter: rejRes.msBeforeNext,
};
}
throw rejRes;
}
})
);
// Find if any window rejected the request
const rejectedWindow = windowResults.find(r => !r.allowed);
if (rejectedWindow) {
// Find the longest wait time among all rejected windows
const maxRetryAfter = Math.max(
...windowResults.filter(r => !r.allowed).map(r => r.retryAfter || 0)
);
this.logger.warn('Rate limit exceeded', {
handler,
operation,
cost,
retryAfter: maxRetryAfter,
});
return {
...result,
allowed: false,
retryAfter: maxRetryAfter,
remainingPoints: rejectedWindow.remainingPoints,
appliedRule: applicableRule,
cost,
};
} catch (error) {
this.logger.error('Rate limit check failed', { queueName, handler, operation, error });
// On error, allow the request to proceed
return { allowed: true };
}
// All windows allowed - return the minimum remaining points
const minRemainingPoints = Math.min(...windowResults.map(r => r.remainingPoints || 0));
return {
allowed: true,
remainingPoints: minRemainingPoints,
appliedRule: applicableRule,
cost,
};
}
/**
@ -144,35 +205,6 @@ export class QueueRateLimiter {
return rule;
}
/**
* Consume a point from the rate limiter
*/
private async consumePoint(
limiter: RateLimiterRedis,
key: string
): Promise<{ allowed: boolean; retryAfter?: number; remainingPoints?: number }> {
try {
const result = await limiter.consume(key);
return {
allowed: true,
remainingPoints: result.remainingPoints,
};
} catch (rejRes) {
if (rejRes instanceof RateLimiterRes) {
this.logger.warn('Rate limit exceeded', {
key,
retryAfter: rejRes.msBeforeNext,
});
return {
allowed: false,
retryAfter: rejRes.msBeforeNext,
remainingPoints: rejRes.remainingPoints,
};
}
throw rejRes;
}
}
/**
* Get rule key for storing rate limiter
@ -216,13 +248,13 @@ export class QueueRateLimiter {
handler: string;
operation: string;
appliedRule?: RateLimitRule;
limit?: {
level: string;
cost?: number;
windows?: Array<{
points: number;
duration: number;
remaining: number;
resetIn: number;
};
}>;
}> {
const applicableRule = this.getMostSpecificRule(queueName, handler, operation);
@ -240,9 +272,9 @@ export class QueueRateLimiter {
applicableRule.handler,
applicableRule.operation
);
const limiter = this.limiters.get(key);
const limiterEntry = this.limiters.get(key);
if (!limiter) {
if (!limiterEntry || limiterEntry.windows.length === 0) {
return {
queueName,
handler,
@ -253,22 +285,27 @@ export class QueueRateLimiter {
try {
const consumerKey = this.getConsumerKey(queueName, handler, operation);
const result = await limiter.get(consumerKey);
const limit = {
level: applicableRule.level,
points: limiter.points,
duration: limiter.duration,
remaining: result?.remainingPoints ?? limiter.points,
resetIn: result?.msBeforeNext ?? 0,
};
// Get status for all windows
const windows = await Promise.all(
limiterEntry.windows.map(async ({ limiter, window }) => {
const result = await limiter.get(consumerKey);
return {
points: window.points,
duration: window.duration,
remaining: result?.remainingPoints ?? window.points,
resetIn: result?.msBeforeNext ?? 0,
};
})
);
return {
queueName,
handler,
operation,
appliedRule: applicableRule,
limit,
cost: limiterEntry.cost,
windows,
};
} catch (error) {
this.logger.error('Failed to get rate limit status', {
@ -282,6 +319,7 @@ export class QueueRateLimiter {
handler,
operation,
appliedRule: applicableRule,
cost: limiterEntry.cost,
};
}
}
@ -297,9 +335,12 @@ export class QueueRateLimiter {
if (rule) {
const key = this.getRuleKey(rule.level, rule.queueName, rule.handler, rule.operation);
const limiter = this.limiters.get(key);
if (limiter) {
await limiter.delete(consumerKey);
const limiterEntry = this.limiters.get(key);
if (limiterEntry) {
// Reset all windows for this consumer
await Promise.all(
limiterEntry.windows.map(({ limiter }) => limiter.delete(consumerKey))
);
}
}
} else {