From 02ebe1ed80cb4768e914fafeadadfb241e498e3f Mon Sep 17 00:00:00 2001 From: jlenon7 Date: Sat, 6 Sep 2025 13:21:08 -0300 Subject: [PATCH 1/6] feat: implement retry strategy and API Targets --- package-lock.json | 4 +- package.json | 3 +- src/exceptions/MissingRuleException.ts | 18 +- src/index.ts | 4 +- src/ratelimiter/RateLimitStore.ts | 148 +++++ src/ratelimiter/RateLimiterBuilder.ts | 524 +++++++++++++++++- src/ratelimiter/stores/MemoryStore.ts | 91 --- src/ratelimiter/stores/RateLimitStore.ts | 34 -- src/ratelimiter/stores/RedisStore.ts | 87 --- src/types/QueueItem.ts | 8 +- src/types/RateLimitApiTarget.ts | 37 ++ src/types/RateLimitRetryClosure.ts | 14 + src/types/RateLimitRetryCtx.ts | 32 ++ src/types/RateLimitRetryDecision.ts | 18 + src/types/RateLimitScheduleCtx.ts | 15 + src/types/RateLimiterOptions.ts | 35 +- src/types/ScheduleOptions.ts | 1 - src/types/index.ts | 5 + .../ratelimiter/RateLimiterBuilderTest.ts | 224 +++++++- 19 files changed, 1023 insertions(+), 279 deletions(-) create mode 100644 src/ratelimiter/RateLimitStore.ts delete mode 100644 src/ratelimiter/stores/MemoryStore.ts delete mode 100644 src/ratelimiter/stores/RateLimitStore.ts delete mode 100644 src/ratelimiter/stores/RedisStore.ts create mode 100644 src/types/RateLimitApiTarget.ts create mode 100644 src/types/RateLimitRetryClosure.ts create mode 100644 src/types/RateLimitRetryCtx.ts create mode 100644 src/types/RateLimitRetryDecision.ts create mode 100644 src/types/RateLimitScheduleCtx.ts diff --git a/package-lock.json b/package-lock.json index b249fc8..2e8bff9 100644 --- a/package-lock.json +++ b/package-lock.json @@ -1,12 +1,12 @@ { "name": "@athenna/ratelimiter", - "version": "5.2.0", + "version": "5.3.0", "lockfileVersion": 3, "requires": true, "packages": { "": { "name": "@athenna/ratelimiter", - "version": "5.2.0", + "version": "5.3.0", "license": "MIT", "devDependencies": { "@athenna/cache": "^5.2.0", diff --git a/package.json b/package.json index 4a7f1f6..8f182c8 100644 --- a/package.json +++ b/package.json @@ -1,6 +1,6 @@ { "name": "@athenna/ratelimiter", - "version": "5.2.0", + "version": "5.3.0", "description": "Respect the rate limit rules of API's you need to consume.", "license": "MIT", "author": "João Lenon ", @@ -152,6 +152,7 @@ "camelcase": "off", "dot-notation": "off", "prettier/prettier": "error", + "no-case-declarations": "off", "no-useless-constructor": "off", "@typescript-eslint/no-explicit-any": "off", "@typescript-eslint/no-empty-function": "off", diff --git a/src/exceptions/MissingRuleException.ts b/src/exceptions/MissingRuleException.ts index 84cde0d..e7d676d 100644 --- a/src/exceptions/MissingRuleException.ts +++ b/src/exceptions/MissingRuleException.ts @@ -10,11 +10,23 @@ import { Exception } from '@athenna/common' export class MissingRuleException extends Exception { - public constructor() { + public constructor(apiTargets?: string[]) { + let message = 'Missing rules value for rate limiter.' + let help = + 'This error happens when you forget to define rules for your RateLimiter instance.' + + if (apiTargets) { + message = `Missing rules value for your API Targets: ${apiTargets.join( + ', ' + )}` + help = + 'This error happens when you forget to define rules for you API Target or for your RateLimtier instance. You have two options, define a custom rule for all your API Targets or define a default rule in your RateLimiter to be used by default by API Targets.' + } + super({ code: 'E_MISSING_RULE_ERROR', - help: 'This errors happens when you forget to define rules for your RateLimiter instance.', - message: 'Missing rules value for rate limiter.' + help, + message }) } } diff --git a/src/index.ts b/src/index.ts index 6ba2cf7..6e63b54 100644 --- a/src/index.ts +++ b/src/index.ts @@ -10,7 +10,5 @@ export * from '#src/types' export * from '#src/ratelimiter/RateLimiter' +export * from '#src/ratelimiter/RateLimitStore' export * from '#src/ratelimiter/RateLimiterBuilder' -export * from '#src/ratelimiter/stores/RedisStore' -export * from '#src/ratelimiter/stores/MemoryStore' -export * from '#src/ratelimiter/stores/RateLimitStore' diff --git a/src/ratelimiter/RateLimitStore.ts b/src/ratelimiter/RateLimitStore.ts new file mode 100644 index 0000000..d90e262 --- /dev/null +++ b/src/ratelimiter/RateLimitStore.ts @@ -0,0 +1,148 @@ +/** + * @athenna/ratelimiter + * + * (c) João Lenon + * + * For the full copyright and license information, please view the LICENSE + * file that was distributed with this source code. + */ + +import { debug } from '#src/debug' +import { Cache } from '@athenna/cache' +import { Macroable } from '@athenna/common' +import { WINDOW_MS } from '#src/constants/window' +import type { Reserve, RateLimitRule, RateLimitStoreOptions } from '#src/types' + +export class RateLimitStore extends Macroable { + /** + * Holds the options that will be used to build the rate limiter + * store. + */ + public options: RateLimitStoreOptions + + public constructor(options: RateLimitStoreOptions) { + super() + + options.windowMs = options.windowMs ?? WINDOW_MS + + this.options = options + } + + public async getOrInit(key: string, rules: RateLimitRule[]) { + const cache = Cache.store(this.options.store) + + let buckets = await cache.get(key) + + if (!buckets) { + buckets = JSON.stringify(rules.map(() => [])) + + await cache.set(key, buckets) + } + + return JSON.parse(buckets) as number[][] + } + + /** + * Get the defined cooldown if it exists in the cache. + * If it cannot be found, return 0. + */ + public async getCooldown(key: string) { + const cdKey = `${key}:cooldown` + + debug('getting cooldown in %s store for key %s', this.options.store, cdKey) + + const cooldown = await Cache.store(this.options.store).get(cdKey) + + if (!cooldown) { + return 0 + } + + return Number(cooldown) + } + + /** + * Put the key in cooldown for some milliseconds. Also saves + * the timestamp into the cache for when it will be available + * again. + */ + public async setCooldown(key: string, ms: number) { + if (!ms || ms <= 0) { + return + } + + const cdKey = `${key}:cooldown` + const cdMs = `${Date.now() + ms}` + + debug( + 'setting cooldown of %s ms in %s store for key %s', + cdMs, + this.options.store, + cdKey + ) + + await Cache.store(this.options.store).set(cdKey, cdMs, { ttl: ms }) + } + + /** + * Try to reserve a token for all rules of the key. If not + * allowed to reserve, return the maximum waitMs necessary. + */ + public async tryReserve(key: string, rules: RateLimitRule[]) { + debug( + 'running %s store tryReserve for key %s with rules %o', + this.options.store, + key, + rules + ) + + let wait = 0 + const now = Date.now() + const cache = Cache.store(this.options.store) + const cooldown = await this.getCooldown(key) + + if (Number.isFinite(cooldown) && cooldown > now) { + return { allowed: false, waitMs: cooldown - now } + } + + await cache.delete(`${key}:cooldown`) + + const buckets = await this.getOrInit(key, rules) + + for (let i = 0; i < rules.length; i++) { + const bucket = buckets[i] + const window = this.options.windowMs[rules[i].type] + + while (bucket.length && bucket[0] <= now - window) { + bucket.shift() + } + + if (bucket.length >= rules[i].limit) { + const earliest = bucket[0] + const rem = earliest + window - now + + if (rem > wait) { + wait = rem + } + } + } + + const reserve: Reserve = { allowed: false, waitMs: wait } + + if (wait > 0) { + await cache.set(key, JSON.stringify(buckets)) + + return reserve + } + + for (let i = 0; i < rules.length; i++) { + buckets[i].push(now) + } + + await cache.set(key, JSON.stringify(buckets)) + + reserve.waitMs = 0 + reserve.allowed = true + + return reserve + } +} diff --git a/src/ratelimiter/RateLimiterBuilder.ts b/src/ratelimiter/RateLimiterBuilder.ts index efbb363..592fed7 100644 --- a/src/ratelimiter/RateLimiterBuilder.ts +++ b/src/ratelimiter/RateLimiterBuilder.ts @@ -11,21 +11,28 @@ import type { QueueItem, RateLimitRule, ScheduleOptions, - RateLimiterOptions + RateLimitRetryCtx, + RateLimiterOptions, + RateLimitApiTarget, + RateLimitScheduleCtx, + RateLimitStoreOptions, + RateLimitRetryDecision } from '#src/types' +import { Macroable } from '@athenna/common' +import { RateLimitStore } from '#src/ratelimiter/RateLimitStore' import { MissingKeyException } from '#src/exceptions/MissingKeyException' import { MissingRuleException } from '#src/exceptions/MissingRuleException' -import type { RateLimitStore } from '#src/ratelimiter/stores/RateLimitStore' import { MissingStoreException } from '#src/exceptions/MissingStoreException' -export class RateLimiterBuilder { +export class RateLimiterBuilder extends Macroable { /** * Holds the options that will be used to build the rate limiter. */ private options: RateLimiterOptions = { maxConcurrent: 1, - jitterMs: 0 + jitterMs: 0, + apiTargetSelectionStrategy: 'first_available' } /** @@ -46,6 +53,11 @@ export class RateLimiterBuilder { */ private queue: QueueItem[] = [] + /** + * Index for when using round_robin selection strategy. + */ + private rrIndex = 0 + /** * Holds the setTimeout id to be able to disable it * later on. @@ -57,6 +69,26 @@ export class RateLimiterBuilder { */ private nextWakeUpAt: number = 0 + /** + * Cooldown per API Target for when an error happens. + */ + private cooldownUntil = new Map() + + /** + * Map the key inside the store. + */ + private createApiTargetKey(apiTarget: RateLimitApiTarget) { + let host = '' + + try { + host = new URL(apiTarget.baseUrl).host + } catch { + host = apiTarget.baseUrl + } + + return `${this.options.key}:${host}:${apiTarget.id}` + } + /** * Logical key that will be used by store to save buckets. */ @@ -70,8 +102,15 @@ export class RateLimiterBuilder { * Define the store that will be responsible to save the * rate limit buckets. */ - public store(value: RateLimitStore) { - this.options.store = value + public store( + store: 'memory' | 'redis' | string, + options: Omit = {} + ) { + // eslint-disable-next-line + // @ts-ignore + options.store = store + + this.options.store = new RateLimitStore(options) return this } @@ -108,14 +147,65 @@ export class RateLimiterBuilder { return this } + /** + * Add a new rate limit API target. + */ + public addApiTarget(apiTarget: RateLimitApiTarget) { + if (!this.options.apiTargets) { + this.options.apiTargets = [] + } + + this.options.apiTargets.push(apiTarget) + + return this + } + + /** + * Set multiple rate limit rules with one method call. + */ + public setRules(rules: RateLimitRule[]) { + rules.forEach(rule => this.addRule(rule)) + + return this + } + + /** + * Set multiple rate limit API targets with one method call. + */ + public setApiTargets(apiTargets: RateLimitApiTarget[]) { + apiTargets.forEach(apiTarget => this.addApiTarget(apiTarget)) + + return this + } + + /** + * Define the API target selection strategy that will be used + * to select the next one when an API fails. + */ + public apiTargetSelectionStrategy(value: 'first_available' | 'round_robin') { + this.options.apiTargetSelectionStrategy = value + + return this + } + + public retryStrategy( + fn: ( + ctx: RateLimitRetryCtx + ) => RateLimitRetryDecision | Promise + ) { + this.options.retryStrategy = fn + + return this + } + /** * Return the current number of active tasks. * * @example * ```ts * const limiter = RateLimiter.build() - * .store(new MemoryStore()) - * .key('request:api-key:/profile') + * .store('memory') + * .key('request:/profile') * .addRule({ type: 'second', limit: 1 }) * * limiter.getActiveCount() // 0 @@ -132,8 +222,8 @@ export class RateLimiterBuilder { * @example * ```ts * const limiter = RateLimiter.build() - * .store(new MemoryStore()) - * .key('request:api-key:/profile') + * .store('memory') + * .key('request:/profile') * .addRule({ type: 'second', limit: 1 }) * * limiter.getQueuedCount() // 0 @@ -144,14 +234,14 @@ export class RateLimiterBuilder { } /** - * Stimate when the next slot will be available based on the + * Estimate when the next slot will be available based on the * next retry defined. * * @example * ```ts * const limiter = RateLimiter.build() - * .store(new MemoryStore()) - * .key('request:api-key:/profile') + * .store('memory') + * .key('request:/profile') * .addRule({ type: 'second', limit: 1 }) * * limiter.getAvailableInMs() // 0 @@ -171,6 +261,7 @@ export class RateLimiterBuilder { this.active = 0 this.nextWakeUpAt = 0 this.storeErrorCount = 0 + this.cooldownUntil.clear() if (this.timer) { clearTimeout(this.timer) @@ -185,7 +276,7 @@ export class RateLimiterBuilder { * the rate limit rules. */ public schedule( - fn: (signal?: AbortSignal) => T | Promise, + closure: (ctx: RateLimitScheduleCtx) => T | Promise, opts: ScheduleOptions = {} ): Promise { if (!this.options.key) { @@ -197,7 +288,19 @@ export class RateLimiterBuilder { } if (!this.options.rules?.length) { - throw new MissingRuleException() + if (!this.options.apiTargets?.length) { + throw new MissingRuleException() + } + + const missingRuleApiTargets = this.options.apiTargets.filter( + apiTarget => !apiTarget.rules + ) + + if (missingRuleApiTargets.length) { + throw new MissingRuleException( + missingRuleApiTargets.map(t => t.baseUrl) + ) + } } if (opts.signal?.aborted) { @@ -206,11 +309,12 @@ export class RateLimiterBuilder { return new Promise((resolve, reject) => { const item: QueueItem = { - run: fn, + run: closure, resolve, reject, started: false, - opts + signal: opts.signal, + attempt: 1 } this.queue.push(item) @@ -264,6 +368,225 @@ export class RateLimiterBuilder { const now = Date.now() + if (this.options.apiTargets?.length) { + let minWait = Number.POSITIVE_INFINITY + let apiTargetChosen: RateLimitApiTarget = null + + const nextItem = this.queue[0] + const pinnedApiTargetId = nextItem?.pinnedApiTargetId + + for (const i of this.createIdxBySelectionStrategy(pinnedApiTargetId)) { + const apiTarget = this.options.apiTargets[i] + const cooldown = this.cooldownUntil.get(apiTarget.id) + + if (cooldown && cooldown > now) { + minWait = Math.min(minWait, cooldown - now) + continue + } + + const key = this.createApiTargetKey(apiTarget) + const rules = apiTarget.rules?.length + ? apiTarget.rules + : this.options.rules + + try { + const res = await this.options.store.tryReserve(key, rules) + + this.storeErrorCount = 0 + + if (res.allowed) { + apiTargetChosen = apiTarget + + if (this.options.apiTargetSelectionStrategy === 'round_robin') { + this.rrIndex = (i + 1) % this.options.apiTargets.length + } + + break + } else { + minWait = Math.min(minWait, res.waitMs) + } + } catch (error) { + this.storeErrorCount++ + + if (this.storeErrorCount > 10) { + while (this.queue.length) { + this.queue.shift().reject(error) + } + + throw error + } + + minWait = Math.min(minWait, 100) + } + } + + if (!apiTargetChosen) { + const delay = + (isFinite(minWait) ? minWait : 100) + this.randomJitter() + + this.nextWakeUpAt = now + delay + this.timer = setTimeout(tryRun, delay) + + return + } + + const item = this.queue.shift() + + if (item.signal?.aborted) { + item.reject(new DOMException('Aborted', 'AbortError')) + this.pump() + return + } + + item.started = true + + if (item.signal && item.abortHandler) { + item.signal.removeEventListener('abort', item.abortHandler) + item.abortHandler = undefined + } + + this.active++ + + Promise.resolve() + .then(() => + item.run({ signal: item.signal, apiTarget: apiTargetChosen }) + ) + .then(result => { + this.release({ isToPump: true }) + + item.resolve(result) + }) + .catch(async error => { + if (!this.options.retryStrategy) { + this.release({ isToPump: true }) + + item.reject(error) + + return + } + + const key = this.createApiTargetKey(apiTargetChosen) + + const ctx: RateLimitRetryCtx = { + key, + error, + attempt: item.attempt, + apiTarget: apiTargetChosen + } + + const decision = await this.options.retryStrategy(ctx) + + switch (decision.type) { + case 'retry_same': + this.release({ isToPump: false }) + + item.attempt++ + item.started = false + item.pinnedApiTargetId = apiTargetChosen.id + + if (item.signal?.aborted) { + item.reject(new DOMException('Aborted', 'AbortError')) + break + } + + this.queue.unshift(item) + + const delay = Math.max(0, decision.delayMs ?? 0) + + if (delay > 0) { + this.nextWakeUpAt = Date.now() + delay + this.timer = setTimeout(() => this.pump(), delay) + } else { + this.pump() + } + + return + + case 'retry_other': { + if (decision.cooldownMs > 0) { + await this.options.store!.setCooldown( + key, + decision.cooldownMs + ) + } + + this.release({ isToPump: false }) + + item.attempt++ + item.started = false + item.pinnedApiTargetId = undefined + + if (item.signal?.aborted) { + item.reject(new DOMException('Aborted', 'AbortError')) + break + } + + this.queue.unshift(item) + + const delay = Math.max(0, decision.delayMs ?? 0) + + if (delay > 0) { + this.nextWakeUpAt = Date.now() + delay + this.timer = setTimeout(() => this.pump(), delay) + } else { + this.pump() + } + + return + } + + case 'cooldown': + await this.options.store!.setCooldown( + key, + Math.max(0, decision.cooldownMs) + ) + + if ( + decision.then === 'retry_same' || + decision.then === 'retry_other' + ) { + this.release({ isToPump: false }) + + item.attempt++ + item.started = false + item.pinnedApiTargetId = + decision.then === 'retry_same' + ? apiTargetChosen.id + : undefined + + if (item.signal?.aborted) { + item.reject(new DOMException('Aborted', 'AbortError')) + + break + } + + this.queue.unshift(item) + + const delay = decision.cooldownMs + + this.nextWakeUpAt = Date.now() + delay + this.timer = setTimeout(() => this.pump(), delay) + + return + } + + this.release({ isToPump: true }) + item.reject(error) + break + default: + this.release({ isToPump: true }) + + item.reject(error) + } + }) + + if (this.active < this.options.maxConcurrent) { + this.timer = setTimeout(tryRun, 0) + } + + return + } + let waitMs = 0 let allowed = false @@ -286,9 +609,7 @@ export class RateLimiterBuilder { */ if (this.storeErrorCount > 10) { while (this.queue.length) { - const item = this.queue.shift()! - - item.reject(error) + this.queue.shift().reject(error) } throw error @@ -307,9 +628,9 @@ export class RateLimiterBuilder { return } - const item = this.queue.shift()! + const item = this.queue.shift() - if (item.opts.signal?.aborted) { + if (item.signal?.aborted) { item.reject(new DOMException('Aborted', 'AbortError')) this.pump() @@ -319,19 +640,105 @@ export class RateLimiterBuilder { item.started = true - if (item.opts.signal && item.abortHandler) { - item.opts.signal.removeEventListener('abort', item.abortHandler) + if (item.signal && item.abortHandler) { + item.signal.removeEventListener('abort', item.abortHandler) item.abortHandler = undefined } this.active++ Promise.resolve() - .then(() => item.run(item.opts.signal)) - .then(item.resolve, item.reject) - .finally(() => { - this.active-- - this.pump() + .then(() => item.run({ signal: item.signal })) + .then(result => { + this.release({ isToPump: true }) + + item.resolve(result) + }) + .catch(async error => { + if (!this.options.retryStrategy) { + this.release({ isToPump: true }) + + item.reject(error) + + return + } + + const ctx: RateLimitRetryCtx = { + error, + key: this.options.key, + attempt: item.attempt + } + + const decision = await this.options.retryStrategy(ctx) + + switch (decision.type) { + case 'retry_same': + case 'retry_other': + const delay = Math.max(0, decision.delayMs ?? 0) + + this.release({ isToPump: false }) + + item.attempt++ + item.started = false + + if (item.signal?.aborted) { + item.reject(new DOMException('Aborted', 'AbortError')) + break + } + + this.queue.unshift(item) + + if (delay > 0) { + this.nextWakeUpAt = Date.now() + delay + this.timer = setTimeout(() => this.pump(), delay) + } else { + this.pump() + } + + return + + case 'cooldown': { + await this.options.store!.setCooldown( + this.options.key, + Math.max(0, decision.cooldownMs) + ) + + if ( + decision.then === 'retry_same' || + decision.then === 'retry_other' + ) { + const delay = decision.cooldownMs + + this.release({ isToPump: false }) + + item.attempt++ + item.started = false + + if (item.signal?.aborted) { + item.reject(new DOMException('Aborted', 'AbortError')) + + break + } + + this.queue.unshift(item) + this.nextWakeUpAt = Date.now() + delay + this.timer = setTimeout(() => this.pump(), delay) + + return + } + + this.release({ isToPump: true }) + + item.reject(error) + + break + } + + default: + this.release({ isToPump: true }) + + item.reject(error) + } }) if (this.active < this.options.maxConcurrent) { @@ -353,4 +760,63 @@ export class RateLimiterBuilder { return Math.floor(Math.random() * this.options.jitterMs) } + + /** + * Read the API Target selection strategy and defines which is + * going to be used. + */ + private createIdxBySelectionStrategy(pinnedApiTargetId?: string) { + let indexes = [] + + switch (this.options.apiTargetSelectionStrategy) { + case 'round_robin': + indexes = this.createRoundRobinIdx() + break + case 'first_available': + indexes = this.createFirstAvailableIdx() + break + default: + indexes = this.createFirstAvailableIdx() + } + + if (pinnedApiTargetId) { + const i = this.options.apiTargets.findIndex( + a => a.id === pinnedApiTargetId + ) + + if (i >= 0) { + indexes = [i, ...indexes.filter(x => x !== i)] + } + } + + return indexes + } + + /** + * Create the indexes for when using round_robin selection strategy. + */ + private createRoundRobinIdx() { + return Array.from( + { length: this.options.apiTargets.length }, + (_, k) => (this.rrIndex + k) % this.options.apiTargets.length + ) + } + + /** + * Create the indexes for when using first_available selection strategy. + */ + private createFirstAvailableIdx() { + return Array.from({ length: this.options.apiTargets.length }, (_, k) => k) + } + + /** + * Release the rate limit task. + */ + private release(options: { isToPump: boolean }) { + this.active-- + + if (options.isToPump) { + this.pump() + } + } } diff --git a/src/ratelimiter/stores/MemoryStore.ts b/src/ratelimiter/stores/MemoryStore.ts deleted file mode 100644 index a35b606..0000000 --- a/src/ratelimiter/stores/MemoryStore.ts +++ /dev/null @@ -1,91 +0,0 @@ -/** - * @athenna/ratelimiter - * - * (c) João Lenon - * - * For the full copyright and license information, please view the LICENSE - * file that was distributed with this source code. - */ - -import { debug } from '#src/debug' -import { Cache } from '@athenna/cache' -import { Options } from '@athenna/common' -import { RateLimitStore } from '#src/ratelimiter/stores/RateLimitStore' -import type { Reserve, RateLimitRule, RateLimitStoreOptions } from '#src/types' - -export class MemoryStore extends RateLimitStore { - public constructor(options: RateLimitStoreOptions = {}) { - options = Options.create(options, { - store: 'memory' - }) - - super(options) - } - - private async getOrInit(key: string, rules: RateLimitRule[]) { - const cache = Cache.store(this.options.store) - - let buckets = await cache.get(key) - - if (!buckets) { - buckets = JSON.stringify(rules.map(() => [])) - - await cache.set(key, buckets) - } - - return JSON.parse(buckets) as number[][] - } - - /** - * Try to reserve a token for all rules of the key. If not - * allowed to reserve, return the maximum waitMs necessary. - */ - public async tryReserve(key: string, rules: RateLimitRule[]) { - debug( - 'running memory store tryReserve for key %s with rules %o', - key, - rules - ) - - let wait = 0 - const now = Date.now() - const buckets = await this.getOrInit(key, rules) - - for (let i = 0; i < rules.length; i++) { - const bucket = buckets[i] - const window = this.options.windowMs[rules[i].type] - - while (bucket.length && bucket[0] <= now - window) { - bucket.shift() - } - - if (bucket.length >= rules[i].limit) { - const earliest = bucket[0] - const rem = earliest + window - now - - if (rem > wait) { - wait = rem - } - } - } - - const reserve: Reserve = { allowed: false, waitMs: wait } - - if (wait > 0) { - await Cache.store(this.options.store).set(key, JSON.stringify(buckets)) - - return reserve - } - - for (let i = 0; i < rules.length; i++) { - buckets[i].push(now) - } - - await Cache.store(this.options.store).set(key, JSON.stringify(buckets)) - - reserve.waitMs = 0 - reserve.allowed = true - - return reserve - } -} diff --git a/src/ratelimiter/stores/RateLimitStore.ts b/src/ratelimiter/stores/RateLimitStore.ts deleted file mode 100644 index 24982d3..0000000 --- a/src/ratelimiter/stores/RateLimitStore.ts +++ /dev/null @@ -1,34 +0,0 @@ -/** - * @athenna/ratelimiter - * - * (c) João Lenon - * - * For the full copyright and license information, please view the LICENSE - * file that was distributed with this source code. - */ - -import { WINDOW_MS } from '#src/constants/window' -import type { Reserve, RateLimitRule, RateLimitStoreOptions } from '#src/types' - -export abstract class RateLimitStore { - /** - * Holds the options that will be used to build the rate limiter - * store. - */ - public options: RateLimitStoreOptions - - public constructor(options: RateLimitStoreOptions = {}) { - options.windowMs = options.windowMs ?? WINDOW_MS - - this.options = options - } - - /** - * Try to reserve a token para all rules of the key. - * If allowed is false, return the maximum waitMs necessary. - */ - public abstract tryReserve( - key: string, - rules: RateLimitRule[] - ): Promise -} diff --git a/src/ratelimiter/stores/RedisStore.ts b/src/ratelimiter/stores/RedisStore.ts deleted file mode 100644 index 4a46a24..0000000 --- a/src/ratelimiter/stores/RedisStore.ts +++ /dev/null @@ -1,87 +0,0 @@ -/** - * @athenna/ratelimiter - * - * (c) João Lenon - * - * For the full copyright and license information, please view the LICENSE - * file that was distributed with this source code. - */ - -import { debug } from '#src/debug' -import { Cache } from '@athenna/cache' -import { Options } from '@athenna/common' -import { RateLimitStore } from '#src/ratelimiter/stores/RateLimitStore' -import type { Reserve, RateLimitRule, RateLimitStoreOptions } from '#src/types' - -export class RedisStore extends RateLimitStore { - public constructor(options: RateLimitStoreOptions = {}) { - options = Options.create(options, { - store: 'redis' - }) - - super(options) - } - - private async getOrInit(key: string, rules: RateLimitRule[]) { - const cache = Cache.store(this.options.store) - - let buckets = await cache.get(key) - - if (!buckets) { - buckets = JSON.stringify(rules.map(() => [])) - - await cache.set(key, buckets) - } - - return JSON.parse(buckets) as number[][] - } - - /** - * Try to reserve a token for all rules of the key. If not - * allowed to reserve, return the maximum waitMs necessary. - */ - public async tryReserve(key: string, rules: RateLimitRule[]) { - debug('running redis store tryReserve for key %s with rules %o', key, rules) - - let wait = 0 - const now = Date.now() - const buckets = await this.getOrInit(key, rules) - - for (let i = 0; i < rules.length; i++) { - const bucket = buckets[i] - const window = this.options.windowMs[rules[i].type] - - while (bucket.length && bucket[0] <= now - window) { - bucket.shift() - } - - if (bucket.length >= rules[i].limit) { - const earliest = bucket[0] - const rem = earliest + window - now - - if (rem > wait) { - wait = rem - } - } - } - - const reserve: Reserve = { allowed: false, waitMs: wait } - - if (wait > 0) { - await Cache.store(this.options.store).set(key, JSON.stringify(buckets)) - - return reserve - } - - for (let i = 0; i < rules.length; i++) { - buckets[i].push(now) - } - - await Cache.store(this.options.store).set(key, JSON.stringify(buckets)) - - reserve.waitMs = 0 - reserve.allowed = true - - return reserve - } -} diff --git a/src/types/QueueItem.ts b/src/types/QueueItem.ts index 122993d..475769d 100644 --- a/src/types/QueueItem.ts +++ b/src/types/QueueItem.ts @@ -7,13 +7,15 @@ * file that was distributed with this source code. */ -import type { ScheduleOptions } from '#src/types' +import type { RateLimitScheduleCtx } from '#src/types' export type QueueItem = { - run: (signal?: AbortSignal) => T | Promise + run: (ctx: RateLimitScheduleCtx) => T | Promise resolve: (v: T) => void reject: (e: unknown) => void abortHandler?: () => void started: boolean - opts: ScheduleOptions + signal?: AbortSignal + attempt?: number + pinnedApiTargetId?: string } diff --git a/src/types/RateLimitApiTarget.ts b/src/types/RateLimitApiTarget.ts new file mode 100644 index 0000000..165d718 --- /dev/null +++ b/src/types/RateLimitApiTarget.ts @@ -0,0 +1,37 @@ +/** + * @athenna/ratelimiter + * + * (c) João Lenon + * + * For the full copyright and license information, please view the LICENSE + * file that was distributed with this source code. + */ + +import type { RateLimitRule } from '#src/types' + +export type RateLimitApiTarget = { + /** + * The rate limit target ID. Useful for logs and metrics + * and to create a unique key in your store only for this + * API Target. + */ + id: string + + /** + * API target base URL that will be used to fetch the request. + */ + baseUrl: string + + /** + * Custom rate limit rules for this API target. If not defined, + * the default defined in RateLimiter will be used. + */ + rules?: RateLimitRule[] + + /** + * Define any kind of metadata for this API target. Metadata is + * useful to define informations such as API Keys to not only + * create API rotations but API Keys rotations at the same time. + */ + metadata?: Record +} diff --git a/src/types/RateLimitRetryClosure.ts b/src/types/RateLimitRetryClosure.ts new file mode 100644 index 0000000..dae49a5 --- /dev/null +++ b/src/types/RateLimitRetryClosure.ts @@ -0,0 +1,14 @@ +/** + * @athenna/ratelimiter + * + * (c) João Lenon + * + * For the full copyright and license information, please view the LICENSE + * file that was distributed with this source code. + */ + +import type { RateLimitRetryCtx, RateLimitRetryDecision } from '#src/types' + +export type RateLimitRetryClosure = ( + ctx: RateLimitRetryCtx +) => RateLimitRetryDecision | Promise diff --git a/src/types/RateLimitRetryCtx.ts b/src/types/RateLimitRetryCtx.ts new file mode 100644 index 0000000..8597802 --- /dev/null +++ b/src/types/RateLimitRetryCtx.ts @@ -0,0 +1,32 @@ +/** + * @athenna/ratelimiter + * + * (c) João Lenon + * + * For the full copyright and license information, please view the LICENSE + * file that was distributed with this source code. + */ + +import type { RateLimitApiTarget } from '#src/types' + +export type RateLimitRetryCtx = { + /** + * The error that has happened while trying to make the request. + */ + error: unknown + + /** + * The cache key that was used to store the rate limit rules. + */ + key: string + + /** + * Define the number of attempts that have run so far. + */ + attempt: number + + /** + * The API Target that this retry is currently using. + */ + apiTarget?: RateLimitApiTarget +} diff --git a/src/types/RateLimitRetryDecision.ts b/src/types/RateLimitRetryDecision.ts new file mode 100644 index 0000000..c68d38b --- /dev/null +++ b/src/types/RateLimitRetryDecision.ts @@ -0,0 +1,18 @@ +/** + * @athenna/ratelimiter + * + * (c) João Lenon + * + * For the full copyright and license information, please view the LICENSE + * file that was distributed with this source code. + */ + +export type RateLimitRetryDecision = + | { type: 'fail' } + | { type: 'retry_same'; delayMs?: number } + | { type: 'retry_other'; delayMs?: number; cooldownMs?: number } + | { + type: 'cooldown' + cooldownMs: number + then?: 'fail' | 'retry_same' | 'retry_other' + } diff --git a/src/types/RateLimitScheduleCtx.ts b/src/types/RateLimitScheduleCtx.ts new file mode 100644 index 0000000..dd82939 --- /dev/null +++ b/src/types/RateLimitScheduleCtx.ts @@ -0,0 +1,15 @@ +/** + * @athenna/ratelimiter + * + * (c) João Lenon + * + * For the full copyright and license information, please view the LICENSE + * file that was distributed with this source code. + */ + +import type { RateLimitApiTarget } from '#src/types' + +export type RateLimitScheduleCtx = { + signal?: AbortSignal + apiTarget?: RateLimitApiTarget +} diff --git a/src/types/RateLimiterOptions.ts b/src/types/RateLimiterOptions.ts index 01ca715..d3d88e0 100644 --- a/src/types/RateLimiterOptions.ts +++ b/src/types/RateLimiterOptions.ts @@ -7,8 +7,12 @@ * file that was distributed with this source code. */ -import type { RateLimitRule } from '#src/types' -import type { RateLimitStore } from '#src/ratelimiter/stores/RateLimitStore' +import type { + RateLimitRule, + RateLimitApiTarget, + RateLimitRetryClosure +} from '#src/types' +import type { RateLimitStore } from '#src/ratelimiter/RateLimitStore' export type RateLimiterOptions = { /** @@ -18,9 +22,24 @@ export type RateLimiterOptions = { /** * The logical key that will be used by store to save buckets. + * If targets are defined, it will be used as a prefix: + * `${key}:${target.baseUrl}`. */ key?: string + /** + * The api targets that will be used to create API rotations when + * some of them fails. + */ + apiTargets?: RateLimitApiTarget[] + + /** + * The retry strategy for this rate limiter. This is useful to + * give the power to the user when and how we should proceed with + * the retry of API Targets. + */ + retryStrategy?: RateLimitRetryClosure + /** * The store responsible to save the rate limit buckets. */ @@ -28,11 +47,23 @@ export type RateLimiterOptions = { /** * Max number of tasks that could run concurrently. + * + * @default 1 */ maxConcurrent?: number /** * Random jitter in milliseconds to avoid thundering herd in distributed envs. + * + * @default 0 */ jitterMs?: number + + /** + * Define the selection strategy that will be used to select which API target + * will be used next when some of them fails. + * + * @default 'first_available' + */ + apiTargetSelectionStrategy: 'first_available' | 'round_robin' } diff --git a/src/types/ScheduleOptions.ts b/src/types/ScheduleOptions.ts index bcb35ea..b767f9d 100644 --- a/src/types/ScheduleOptions.ts +++ b/src/types/ScheduleOptions.ts @@ -8,6 +8,5 @@ */ export type ScheduleOptions = { - id?: string signal?: AbortSignal } diff --git a/src/types/index.ts b/src/types/index.ts index 7e07fef..221879f 100644 --- a/src/types/index.ts +++ b/src/types/index.ts @@ -11,5 +11,10 @@ export * from '#src/types/Reserve' export * from '#src/types/QueueItem' export * from '#src/types/RateLimitRule' export * from '#src/types/ScheduleOptions' +export * from '#src/types/RateLimitRetryCtx' +export * from '#src/types/RateLimitApiTarget' export * from '#src/types/RateLimiterOptions' +export * from '#src/types/RateLimitScheduleCtx' +export * from '#src/types/RateLimitRetryClosure' export * from '#src/types/RateLimitStoreOptions' +export * from '#src/types/RateLimitRetryDecision' diff --git a/tests/unit/ratelimiter/RateLimiterBuilderTest.ts b/tests/unit/ratelimiter/RateLimiterBuilderTest.ts index bfb7da3..4c0bf60 100644 --- a/tests/unit/ratelimiter/RateLimiterBuilderTest.ts +++ b/tests/unit/ratelimiter/RateLimiterBuilderTest.ts @@ -1,6 +1,6 @@ -import { Path, Sleep } from '@athenna/common' +import { RateLimiter } from '#src' +import { Path, Sleep, Uuid } from '@athenna/common' import { CacheProvider } from '@athenna/cache' -import { MemoryStore, RateLimiter } from '#src' import { AfterEach, BeforeEach, Test, type Context } from '@athenna/test' import { MissingKeyException } from '#src/exceptions/MissingKeyException' import { MissingRuleException } from '#src/exceptions/MissingRuleException' @@ -81,7 +81,7 @@ export class RateLimiterBuilderTest { assert.throws( () => RateLimiter.build() - .store(new MemoryStore()) + .store('memory') .addRule({ type: 'second', limit: 1 }) .schedule(() => {}), MissingKeyException @@ -94,7 +94,7 @@ export class RateLimiterBuilderTest { () => RateLimiter.build() .key('request:api-key:/profile') - .store(new MemoryStore()) + .store('memory') .schedule(() => {}), MissingRuleException ) @@ -115,7 +115,7 @@ export class RateLimiterBuilderTest { @Test() public async shouldBeAbleToBuildARateLimiterWithARuleOfOneRequestPerSecond({ assert }: Context) { const limiter = RateLimiter.build() - .store(new MemoryStore({ windowMs: { second: 100 } })) + .store('memory', { windowMs: { second: 100 } }) .key('request:api-key:/profile') .addRule({ type: 'second', limit: 1 }) @@ -135,7 +135,7 @@ export class RateLimiterBuilderTest { @Test() public async shouldBeAbleToHaveErrorsHappeningInsideTheRateLimiterHandler({ assert }: Context) { const limiter = RateLimiter.build() - .store(new MemoryStore({ windowMs: { second: 100 } })) + .store('memory', { windowMs: { second: 100 } }) .key('request:api-key:/profile') .addRule({ type: 'second', limit: 1 }) @@ -149,7 +149,7 @@ export class RateLimiterBuilderTest { @Test() public async shouldBeAbleToBuildARateLimiterWithARuleOfOneRequestPerMinute({ assert }: Context) { const limiter = RateLimiter.build() - .store(new MemoryStore({ windowMs: { minute: 100 } })) + .store('memory', { windowMs: { minute: 100 } }) .key('request:api-key:/profile') .addRule({ type: 'minute', limit: 1 }) @@ -169,7 +169,7 @@ export class RateLimiterBuilderTest { @Test() public async shouldBeAbleToBuildARateLimiterWithARuleOfOneRequestPerHour({ assert }: Context) { const limiter = RateLimiter.build() - .store(new MemoryStore({ windowMs: { hour: 100 } })) + .store('memory', { windowMs: { hour: 100 } }) .key('request:api-key:/profile') .addRule({ type: 'hour', limit: 1 }) @@ -189,7 +189,7 @@ export class RateLimiterBuilderTest { @Test() public async shouldBeAbleToBuildARateLimiterWithARuleOfOneRequestPerDay({ assert }: Context) { const limiter = RateLimiter.build() - .store(new MemoryStore({ windowMs: { day: 100 } })) + .store('memory', { windowMs: { day: 100 } }) .key('request:api-key:/profile') .addRule({ type: 'day', limit: 1 }) @@ -209,7 +209,7 @@ export class RateLimiterBuilderTest { @Test() public async shouldBeAbleToBuildARateLimiterWithARuleOfOneRequestPerMonth({ assert }: Context) { const limiter = RateLimiter.build() - .store(new MemoryStore({ windowMs: { month: 100 } })) + .store('memory', { windowMs: { month: 100 } }) .key('request:api-key:/profile') .addRule({ type: 'month', limit: 1 }) @@ -229,7 +229,7 @@ export class RateLimiterBuilderTest { @Test() public async shouldBeAbleToBuildARateLimiterWithSecondAndMinutesRules({ assert }: Context) { const limiter = RateLimiter.build() - .store(new MemoryStore({ windowMs: { second: 300, minute: 400 } })) + .store('memory', { windowMs: { second: 300, minute: 400 } }) .key('request:api-key:/profile') .addRule({ type: 'second', limit: 1 }) .addRule({ type: 'minute', limit: 2 }) @@ -268,7 +268,7 @@ export class RateLimiterBuilderTest { @Test() public async shouldBeAbleToCleanTheLimiterExecutionByTruncatingIt({ assert }: Context) { const limiter = RateLimiter.build() - .store(new MemoryStore({ windowMs: { second: 100 } })) + .store('memory', { windowMs: { second: 100 } }) .key('request:api-key:/profile') .addRule({ type: 'second', limit: 1 }) @@ -287,7 +287,7 @@ export class RateLimiterBuilderTest { @Test() public async shouldBeAbleToGetTheQueuedCountOfTheLimiter({ assert }: Context) { const limiter = RateLimiter.build() - .store(new MemoryStore({ windowMs: { second: 100 } })) + .store('memory', { windowMs: { second: 100 } }) .key('request:api-key:/profile') .addRule({ type: 'second', limit: 1 }) @@ -306,7 +306,7 @@ export class RateLimiterBuilderTest { @Test() public async shouldBeAbleToGetTheActiveCountOfTheLimiter({ assert }: Context) { const limiter = RateLimiter.build() - .store(new MemoryStore({ windowMs: { second: 100 } })) + .store('memory', { windowMs: { second: 100 } }) .key('request:api-key:/profile') .addRule({ type: 'second', limit: 1 }) @@ -338,7 +338,7 @@ export class RateLimiterBuilderTest { @Test() public async shouldBeAbleToGetTheActiveCountOfTheLimiterWithConcurrentRequests({ assert }: Context) { const limiter = RateLimiter.build() - .store(new MemoryStore({ windowMs: { second: 100 } })) + .store('memory', { windowMs: { second: 100 } }) .key('request:api-key:/profile') .jitterMs(0) .maxConcurrent(5) @@ -374,7 +374,7 @@ export class RateLimiterBuilderTest { @Test() public async shouldBeAbleToGetWhenTheRateLimiterWillBeAvailable({ assert }: Context) { const limiter = RateLimiter.build() - .store(new MemoryStore({ windowMs: { second: 100 } })) + .store('memory', { windowMs: { second: 100 } }) .key('request:api-key:/profile') .jitterMs(0) .maxConcurrent(10) @@ -396,7 +396,7 @@ export class RateLimiterBuilderTest { @Test() public async shouldBeAbleToAbortARateLimiterTaskThatIsEnqueuedUsingAnAbortController({ assert }: Context) { const limiter = RateLimiter.build() - .store(new MemoryStore({ windowMs: { second: 100 } })) + .store('memory', { windowMs: { second: 100 } }) .key('request:api-key:/profile') .addRule({ type: 'second', limit: 1 }) @@ -428,7 +428,7 @@ export class RateLimiterBuilderTest { @Test() public async shouldNotBeAbleToCancelAlreadyStartedRateLimiterTask({ assert }: Context) { const limiter = RateLimiter.build() - .store(new MemoryStore({ windowMs: { second: 100 } })) + .store('memory', { windowMs: { second: 100 } }) .key('request:api-key:/profile') .addRule({ type: 'second', limit: 1 }) @@ -462,7 +462,7 @@ export class RateLimiterBuilderTest { @Test() public async shouldBeAbleToCancelAlreadyStartedRateLimiterTaskIfUserUsesTheAbortController({ assert }: Context) { const limiter = RateLimiter.build() - .store(new MemoryStore({ windowMs: { second: 100 } })) + .store('memory', { windowMs: { second: 100 } }) .key('request:api-key:/profile') .addRule({ type: 'second', limit: 1 }) @@ -470,10 +470,10 @@ export class RateLimiterBuilderTest { const abortController = new AbortController() const p = limiter.schedule( - async signal => { + async ctx => { started = true - await this.cancellableSleep(1_000, signal) + await this.cancellableSleep(1_000, ctx.signal) return 'ok' }, @@ -502,7 +502,7 @@ export class RateLimiterBuilderTest { @Test() public async shouldBeAbleToAbortARateLimiterTaskBeforeItEvenStarts({ assert }: Context) { const limiter = RateLimiter.build() - .store(new MemoryStore({ windowMs: { second: 100 } })) + .store('memory', { windowMs: { second: 100 } }) .key('request:api-key:/profile') .addRule({ type: 'second', limit: 1 }) @@ -518,7 +518,7 @@ export class RateLimiterBuilderTest { @Test() public async shouldNotBeAbleToAbortTheTaskIfItHasAlreadyStartedRunning({ assert }: Context) { const limiter = RateLimiter.build() - .store(new MemoryStore({ windowMs: { second: 100 } })) + .store('memory', { windowMs: { second: 100 } }) .key('request:api-key:/profile') .addRule({ type: 'second', limit: 1 }) @@ -532,4 +532,182 @@ export class RateLimiterBuilderTest { await assert.doesNotReject(() => p) } + + @Test() + public async shouldBeAbleToBuildARateLimiterWithARuleOfOneRequestPerSecondWithAnApiTarget({ assert }: Context) { + assert.plan(6) + + const limiter = RateLimiter.build() + .store('memory', { windowMs: { second: 100 } }) + .key('request:api-key:/profile') + .addRule({ type: 'second', limit: 1 }) + .addApiTarget({ id: Uuid.generate(), baseUrl: 'http://api1.com' }) + + const promises = [] + const dateStart = Date.now() + const numberOfRequests = 5 + + for (let i = 0; i < numberOfRequests; i++) { + promises.push( + limiter.schedule(({ apiTarget }) => { + assert.isDefined(apiTarget) + + return 'ok' + i + }) + ) + } + + await Promise.all(promises) + + assert.isAtLeast(Date.now() - dateStart, 400) + } + + @Test() + public async shouldBeAbleToHaveErrorsHappeningInsideTheRateLimiterHandlerEvenWithAnApiTargetSet({ assert }: Context) { + const limiter = RateLimiter.build() + .store('memory', { windowMs: { second: 100 } }) + .key('request:api-key:/profile') + .addRule({ type: 'second', limit: 1 }) + .addApiTarget({ id: Uuid.generate(), baseUrl: 'http://api1.com' }) + + await assert.rejects(() => { + return limiter.schedule(() => { + throw new Error('failed') + }) + }) + } + + @Test() + public async shouldBeAbleToTryWithTheSecondApiTargetIfTheFirstApiTargetIsAtFullCapacityInASequentialScenario({ + assert + }: Context) { + const limiter = RateLimiter.build() + .store('memory', { windowMs: { second: 100 } }) + .key('request:api-key:/profile') + .addRule({ type: 'second', limit: 1 }) + .addApiTarget({ id: 't1', baseUrl: 'http://api1.com' }) + .addApiTarget({ id: 't2', baseUrl: 'http://api2.com' }) + + const first = await limiter.schedule(({ apiTarget }) => apiTarget.baseUrl) + + assert.equal(first, 'http://api1.com') + + const second = await limiter.schedule(({ apiTarget }) => apiTarget.baseUrl) + + assert.equal(second, 'http://api2.com') + } + + @Test() + public async shouldBeAbleToTryWithTheSecondApiTargetIfTheFirstApiTargetIsAtFullCapacityInAConcurrentScenario({ + assert + }: Context) { + const limiter = RateLimiter.build() + .maxConcurrent(2) + .store('memory', { windowMs: { second: 100 } }) + .key('request:api-key:/profile') + .addRule({ type: 'second', limit: 1 }) + .addApiTarget({ id: 't1', baseUrl: 'http://api1.com' }) + .addApiTarget({ id: 't2', baseUrl: 'http://api2.com' }) + + const barrier = this.createBarrier() + const used: string[] = [] + + const run = async ({ apiTarget }) => { + used.push(apiTarget.baseUrl) + + await barrier.wait() + + return apiTarget.baseUrl + } + + const p1 = limiter.schedule(run) + const p2 = limiter.schedule(run) + + await Sleep.for(5).milliseconds().wait() + + barrier.release() + + const results = await Promise.all([p1, p2]) + + results.sort() + + assert.deepEqual(results, ['http://api1.com', 'http://api2.com']) + + used.sort() + assert.deepEqual(used, ['http://api1.com', 'http://api2.com']) + } + + @Test() + public async shouldBeAbleToTryWithTheSecondApiTargetIfTheFirstApiTargetIsAtFullCapacityInASequentialScenarioWithRoundRobinStrategy({ + assert + }: Context) { + const limiter = RateLimiter.build() + .store('memory', { windowMs: { second: 100 } }) + .key('request:api-key:/profile') + .apiTargetSelectionStrategy('round_robin') + .addRule({ type: 'second', limit: 1 }) + .addApiTarget({ id: 't1', baseUrl: 'http://api1.com' }) + .addApiTarget({ id: 't2', baseUrl: 'http://api2.com' }) + + const used: string[] = [] + + const tasks = Array.from({ length: 4 }, () => + limiter.schedule(({ apiTarget }) => { + used.push(apiTarget.baseUrl) + + return apiTarget.baseUrl + }) + ) + + const results = await Promise.all(tasks) + + assert.deepEqual(results, used) + assert.deepEqual(used, ['http://api1.com', 'http://api2.com', 'http://api1.com', 'http://api2.com']) + } + + @Test() + public async shouldBeAbleToTryWithTheSecondApiTargetIfTheFirstApiTargetIsAtFullCapacityInAConcurrentScenarioWithRoundRobinStrategy({ + assert + }: Context) { + const limiter = RateLimiter.build() + .maxConcurrent(2) + .store('memory', { windowMs: { second: 100 } }) + .key('request:api-key:/profile') + .apiTargetSelectionStrategy('round_robin') + .addRule({ type: 'second', limit: 1 }) + .addApiTarget({ id: 't1', baseUrl: 'http://api1.com' }) + .addApiTarget({ id: 't2', baseUrl: 'http://api2.com' }) + + const barrier = this.createBarrier() + const started: string[] = [] + + const p1 = limiter.schedule(async ({ apiTarget }) => { + started.push(apiTarget.baseUrl) + + await barrier.wait() + + return apiTarget.baseUrl + }) + const p2 = limiter.schedule(async ({ apiTarget }) => { + started.push(apiTarget.baseUrl) + + await barrier.wait() + + return apiTarget.baseUrl + }) + + for (let i = 0; i < 20 && started.length < 2; i++) { + await Sleep.for(1).milliseconds().wait() + } + + assert.deepEqual(started, ['http://api1.com', 'http://api2.com']) + + barrier.release() + + const result = await Promise.all([p1, p2]) + + result.sort() + + assert.deepEqual(result, ['http://api1.com', 'http://api2.com']) + } } From 55d9f7a477feda58512f28e794ee80b60982b7d9 Mon Sep 17 00:00:00 2001 From: jlenon7 Date: Mon, 8 Sep 2025 18:05:16 -0300 Subject: [PATCH 2/6] test(targets): add tests for API Targets --- package-lock.json | 8 +- package.json | 2 +- src/exceptions/MissingRuleException.ts | 16 +- src/ratelimiter/RateLimiterBuilder.ts | 128 +++++----- src/types/RateLimitApiTarget.ts | 23 +- src/types/RateLimitRetryCtx.ts | 7 +- src/types/RateLimitScheduleCtx.ts | 7 + src/types/RateLimiterOptions.ts | 6 +- tests/fixtures/config/cache.ts | 2 +- .../ratelimiter/RateLimiterBuilderTest.ts | 230 ++++++++++++++++-- 10 files changed, 304 insertions(+), 125 deletions(-) diff --git a/package-lock.json b/package-lock.json index 2e8bff9..ae5fdc9 100644 --- a/package-lock.json +++ b/package-lock.json @@ -10,7 +10,7 @@ "license": "MIT", "devDependencies": { "@athenna/cache": "^5.2.0", - "@athenna/common": "^5.18.0", + "@athenna/common": "^5.19.0", "@athenna/config": "^5.4.0", "@athenna/ioc": "^5.2.0", "@athenna/logger": "^5.10.0", @@ -46,9 +46,9 @@ } }, "node_modules/@athenna/common": { - "version": "5.18.0", - "resolved": "https://registry.npmjs.org/@athenna/common/-/common-5.18.0.tgz", - "integrity": "sha512-UW4A7b8zl42ZLltih6Sm231StqtUbmO3Yq9+hDjSglP8coyjbOlCJ97r/hGjcCpPGBBTET50LcI85aVkDGw7OA==", + "version": "5.19.0", + "resolved": "https://registry.npmjs.org/@athenna/common/-/common-5.19.0.tgz", + "integrity": "sha512-h1UrJFjl+0JS0lGRgnnujo/F4HXJg+/L7Ox7EJM0cPrsEOgg9kCFjGrsf0PEy6AkPwhbhlaYP2W5umw+6CsOew==", "dev": true, "license": "MIT", "dependencies": { diff --git a/package.json b/package.json index 8f182c8..f5cf986 100644 --- a/package.json +++ b/package.json @@ -51,7 +51,7 @@ }, "devDependencies": { "@athenna/cache": "^5.2.0", - "@athenna/common": "^5.18.0", + "@athenna/common": "^5.19.0", "@athenna/config": "^5.4.0", "@athenna/ioc": "^5.2.0", "@athenna/logger": "^5.10.0", diff --git a/src/exceptions/MissingRuleException.ts b/src/exceptions/MissingRuleException.ts index e7d676d..a1afbea 100644 --- a/src/exceptions/MissingRuleException.ts +++ b/src/exceptions/MissingRuleException.ts @@ -10,18 +10,10 @@ import { Exception } from '@athenna/common' export class MissingRuleException extends Exception { - public constructor(apiTargets?: string[]) { - let message = 'Missing rules value for rate limiter.' - let help = - 'This error happens when you forget to define rules for your RateLimiter instance.' - - if (apiTargets) { - message = `Missing rules value for your API Targets: ${apiTargets.join( - ', ' - )}` - help = - 'This error happens when you forget to define rules for you API Target or for your RateLimtier instance. You have two options, define a custom rule for all your API Targets or define a default rule in your RateLimiter to be used by default by API Targets.' - } + public constructor() { + const message = 'Missing rules value for rate limiter and API Targets.' + const help = + 'This error happens when you forget to define default rules for your RateLimiter instance and custom rules by API Target. You has two options, define a default rule in your RateLimiter that will be used by API Targets that does not have a rule or define a custom rule for all your API Targets.' super({ code: 'E_MISSING_RULE_ERROR', diff --git a/src/ratelimiter/RateLimiterBuilder.ts b/src/ratelimiter/RateLimiterBuilder.ts index 592fed7..212ad8b 100644 --- a/src/ratelimiter/RateLimiterBuilder.ts +++ b/src/ratelimiter/RateLimiterBuilder.ts @@ -19,7 +19,8 @@ import type { RateLimitRetryDecision } from '#src/types' -import { Macroable } from '@athenna/common' +import { Config } from '@athenna/config' +import { Json, String, Macroable } from '@athenna/common' import { RateLimitStore } from '#src/ratelimiter/RateLimitStore' import { MissingKeyException } from '#src/exceptions/MissingKeyException' import { MissingRuleException } from '#src/exceptions/MissingRuleException' @@ -30,8 +31,8 @@ export class RateLimiterBuilder extends Macroable { * Holds the options that will be used to build the rate limiter. */ private options: RateLimiterOptions = { - maxConcurrent: 1, jitterMs: 0, + maxConcurrent: 1, apiTargetSelectionStrategy: 'first_available' } @@ -70,23 +71,21 @@ export class RateLimiterBuilder extends Macroable { private nextWakeUpAt: number = 0 /** - * Cooldown per API Target for when an error happens. + * Create a custom id for an API Target by reading the metadata object. + * The object will always be sorted by keys. */ - private cooldownUntil = new Map() + private getApiTargetId(apiTarget: RateLimitApiTarget) { + return String.hash(JSON.stringify(Json.sort(apiTarget.metadata)), { + key: Config.get('app.key', 'ratelimiter') + }) + } /** - * Map the key inside the store. + * Create a custom key for an API Target to be used to map the + * API Target rules into the cache. */ private createApiTargetKey(apiTarget: RateLimitApiTarget) { - let host = '' - - try { - host = new URL(apiTarget.baseUrl).host - } catch { - host = apiTarget.baseUrl - } - - return `${this.options.key}:${host}:${apiTarget.id}` + return `${this.options.key}:${this.getApiTargetId(apiTarget)}` } /** @@ -155,6 +154,10 @@ export class RateLimiterBuilder extends Macroable { this.options.apiTargets = [] } + if (!apiTarget.id) { + apiTarget.id = this.getApiTargetId(apiTarget) + } + this.options.apiTargets.push(apiTarget) return this @@ -188,9 +191,13 @@ export class RateLimiterBuilder extends Macroable { return this } + /** + * Define the RateLmiter retry strategy. This is useful to control + * when and how we should proceed with the retry of API Targets. + */ public retryStrategy( fn: ( - ctx: RateLimitRetryCtx + ctx?: RateLimitRetryCtx ) => RateLimitRetryDecision | Promise ) { this.options.retryStrategy = fn @@ -261,7 +268,6 @@ export class RateLimiterBuilder extends Macroable { this.active = 0 this.nextWakeUpAt = 0 this.storeErrorCount = 0 - this.cooldownUntil.clear() if (this.timer) { clearTimeout(this.timer) @@ -297,9 +303,7 @@ export class RateLimiterBuilder extends Macroable { ) if (missingRuleApiTargets.length) { - throw new MissingRuleException( - missingRuleApiTargets.map(t => t.baseUrl) - ) + throw new MissingRuleException() } } @@ -377,14 +381,8 @@ export class RateLimiterBuilder extends Macroable { for (const i of this.createIdxBySelectionStrategy(pinnedApiTargetId)) { const apiTarget = this.options.apiTargets[i] - const cooldown = this.cooldownUntil.get(apiTarget.id) - - if (cooldown && cooldown > now) { - minWait = Math.min(minWait, cooldown - now) - continue - } - const key = this.createApiTargetKey(apiTarget) + const rules = apiTarget.rules?.length ? apiTarget.rules : this.options.rules @@ -402,9 +400,9 @@ export class RateLimiterBuilder extends Macroable { } break - } else { - minWait = Math.min(minWait, res.waitMs) } + + minWait = Math.min(minWait, res.waitMs) } catch (error) { this.storeErrorCount++ @@ -452,13 +450,13 @@ export class RateLimiterBuilder extends Macroable { item.run({ signal: item.signal, apiTarget: apiTargetChosen }) ) .then(result => { - this.release({ isToPump: true }) + this.releaseTask({ isToPump: true }) item.resolve(result) }) .catch(async error => { if (!this.options.retryStrategy) { - this.release({ isToPump: true }) + this.releaseTask({ isToPump: true }) item.reject(error) @@ -470,6 +468,7 @@ export class RateLimiterBuilder extends Macroable { const ctx: RateLimitRetryCtx = { key, error, + signal: item.signal, attempt: item.attempt, apiTarget: apiTargetChosen } @@ -478,7 +477,7 @@ export class RateLimiterBuilder extends Macroable { switch (decision.type) { case 'retry_same': - this.release({ isToPump: false }) + this.releaseTask({ isToPump: false }) item.attempt++ item.started = false @@ -495,7 +494,7 @@ export class RateLimiterBuilder extends Macroable { if (delay > 0) { this.nextWakeUpAt = Date.now() + delay - this.timer = setTimeout(() => this.pump(), delay) + this.timer = setTimeout(tryRun, delay) } else { this.pump() } @@ -503,14 +502,13 @@ export class RateLimiterBuilder extends Macroable { return case 'retry_other': { - if (decision.cooldownMs > 0) { - await this.options.store!.setCooldown( - key, - decision.cooldownMs - ) - } + const cooldown = Math.max(0, decision.cooldownMs || 0) - this.release({ isToPump: false }) + this.releaseTask({ isToPump: false }) + + if (cooldown > 0) { + await this.options.store!.setCooldown(key, cooldown) + } item.attempt++ item.started = false @@ -527,7 +525,7 @@ export class RateLimiterBuilder extends Macroable { if (delay > 0) { this.nextWakeUpAt = Date.now() + delay - this.timer = setTimeout(() => this.pump(), delay) + this.timer = setTimeout(tryRun, delay) } else { this.pump() } @@ -536,17 +534,15 @@ export class RateLimiterBuilder extends Macroable { } case 'cooldown': - await this.options.store!.setCooldown( - key, - Math.max(0, decision.cooldownMs) - ) + const ms = Math.max(0, decision.cooldownMs || 0) + + this.releaseTask({ isToPump: false }) + + await this.options.store!.setCooldown(key, ms) - if ( - decision.then === 'retry_same' || - decision.then === 'retry_other' - ) { - this.release({ isToPump: false }) + const then = decision.then + if (then === 'retry_same' || then === 'retry_other') { item.attempt++ item.started = false item.pinnedApiTargetId = @@ -565,16 +561,18 @@ export class RateLimiterBuilder extends Macroable { const delay = decision.cooldownMs this.nextWakeUpAt = Date.now() + delay - this.timer = setTimeout(() => this.pump(), delay) + this.timer = setTimeout(tryRun, delay) return } - this.release({ isToPump: true }) + this.releaseTask({ isToPump: true }) + item.reject(error) + break default: - this.release({ isToPump: true }) + this.releaseTask({ isToPump: true }) item.reject(error) } @@ -650,13 +648,13 @@ export class RateLimiterBuilder extends Macroable { Promise.resolve() .then(() => item.run({ signal: item.signal })) .then(result => { - this.release({ isToPump: true }) + this.releaseTask({ isToPump: true }) item.resolve(result) }) .catch(async error => { if (!this.options.retryStrategy) { - this.release({ isToPump: true }) + this.releaseTask({ isToPump: true }) item.reject(error) @@ -676,7 +674,7 @@ export class RateLimiterBuilder extends Macroable { case 'retry_other': const delay = Math.max(0, decision.delayMs ?? 0) - this.release({ isToPump: false }) + this.releaseTask({ isToPump: false }) item.attempt++ item.started = false @@ -690,7 +688,7 @@ export class RateLimiterBuilder extends Macroable { if (delay > 0) { this.nextWakeUpAt = Date.now() + delay - this.timer = setTimeout(() => this.pump(), delay) + this.timer = setTimeout(tryRun, delay) } else { this.pump() } @@ -709,7 +707,7 @@ export class RateLimiterBuilder extends Macroable { ) { const delay = decision.cooldownMs - this.release({ isToPump: false }) + this.releaseTask({ isToPump: false }) item.attempt++ item.started = false @@ -722,12 +720,12 @@ export class RateLimiterBuilder extends Macroable { this.queue.unshift(item) this.nextWakeUpAt = Date.now() + delay - this.timer = setTimeout(() => this.pump(), delay) + this.timer = setTimeout(tryRun, delay) return } - this.release({ isToPump: true }) + this.releaseTask({ isToPump: true }) item.reject(error) @@ -735,7 +733,7 @@ export class RateLimiterBuilder extends Macroable { } default: - this.release({ isToPump: true }) + this.releaseTask({ isToPump: true }) item.reject(error) } @@ -773,8 +771,6 @@ export class RateLimiterBuilder extends Macroable { indexes = this.createRoundRobinIdx() break case 'first_available': - indexes = this.createFirstAvailableIdx() - break default: indexes = this.createFirstAvailableIdx() } @@ -812,10 +808,16 @@ export class RateLimiterBuilder extends Macroable { /** * Release the rate limit task. */ - private release(options: { isToPump: boolean }) { + private releaseTask(options: { isToPump: boolean }) { this.active-- if (options.isToPump) { + if (this.timer) { + clearTimeout(this.timer) + + this.timer = null + } + this.pump() } } diff --git a/src/types/RateLimitApiTarget.ts b/src/types/RateLimitApiTarget.ts index 165d718..8a821bc 100644 --- a/src/types/RateLimitApiTarget.ts +++ b/src/types/RateLimitApiTarget.ts @@ -11,27 +11,24 @@ import type { RateLimitRule } from '#src/types' export type RateLimitApiTarget = { /** - * The rate limit target ID. Useful for logs and metrics - * and to create a unique key in your store only for this - * API Target. + * The rate limit target ID. By default this will be created by creating + * a hash from the API Target metadata object, but you can also define your + * own ID. */ - id: string + id?: string /** - * API target base URL that will be used to fetch the request. + * Define all the metadata for this API target to function. Metadata + * is required because we are going to create a hash from this object + * to store the rules inside the cache by ApiTarget. With this + * implementation you can create not only API rotations but also API + * Keys rotations at the same time. */ - baseUrl: string + metadata: Record /** * Custom rate limit rules for this API target. If not defined, * the default defined in RateLimiter will be used. */ rules?: RateLimitRule[] - - /** - * Define any kind of metadata for this API target. Metadata is - * useful to define informations such as API Keys to not only - * create API rotations but API Keys rotations at the same time. - */ - metadata?: Record } diff --git a/src/types/RateLimitRetryCtx.ts b/src/types/RateLimitRetryCtx.ts index 8597802..ec9311a 100644 --- a/src/types/RateLimitRetryCtx.ts +++ b/src/types/RateLimitRetryCtx.ts @@ -13,7 +13,12 @@ export type RateLimitRetryCtx = { /** * The error that has happened while trying to make the request. */ - error: unknown + error: Error + + /** + * The abort signal to abort the entire process when needed. + */ + signal?: AbortSignal /** * The cache key that was used to store the rate limit rules. diff --git a/src/types/RateLimitScheduleCtx.ts b/src/types/RateLimitScheduleCtx.ts index dd82939..dc3d51f 100644 --- a/src/types/RateLimitScheduleCtx.ts +++ b/src/types/RateLimitScheduleCtx.ts @@ -10,6 +10,13 @@ import type { RateLimitApiTarget } from '#src/types' export type RateLimitScheduleCtx = { + /** + * The abort signal to abort the entire process when needed. + */ signal?: AbortSignal + + /** + * The API Target that this retry is currently using. + */ apiTarget?: RateLimitApiTarget } diff --git a/src/types/RateLimiterOptions.ts b/src/types/RateLimiterOptions.ts index d3d88e0..11690ca 100644 --- a/src/types/RateLimiterOptions.ts +++ b/src/types/RateLimiterOptions.ts @@ -22,8 +22,8 @@ export type RateLimiterOptions = { /** * The logical key that will be used by store to save buckets. - * If targets are defined, it will be used as a prefix: - * `${key}:${target.baseUrl}`. + * If API Targers are defined, it will be used as a prefix from + * a hash created from API Targets metadata object: `${key}:${hash}`. */ key?: string @@ -36,7 +36,7 @@ export type RateLimiterOptions = { /** * The retry strategy for this rate limiter. This is useful to * give the power to the user when and how we should proceed with - * the retry of API Targets. + * the retry of failed executions. */ retryStrategy?: RateLimitRetryClosure diff --git a/tests/fixtures/config/cache.ts b/tests/fixtures/config/cache.ts index 1dbcef2..00a2c5a 100644 --- a/tests/fixtures/config/cache.ts +++ b/tests/fixtures/config/cache.ts @@ -1,5 +1,5 @@ /** - * @athenna/queue + * @athenna/cache * * (c) João Lenon * diff --git a/tests/unit/ratelimiter/RateLimiterBuilderTest.ts b/tests/unit/ratelimiter/RateLimiterBuilderTest.ts index 4c0bf60..73c6c0c 100644 --- a/tests/unit/ratelimiter/RateLimiterBuilderTest.ts +++ b/tests/unit/ratelimiter/RateLimiterBuilderTest.ts @@ -1,6 +1,15 @@ +/** + * @athenna/ratelimiter + * + * (c) João Lenon + * + * For the full copyright and license information, please view the LICENSE + * file that was distributed with this source code. + */ + import { RateLimiter } from '#src' -import { Path, Sleep, Uuid } from '@athenna/common' -import { CacheProvider } from '@athenna/cache' +import { Path, Sleep } from '@athenna/common' +import { Cache, CacheProvider } from '@athenna/cache' import { AfterEach, BeforeEach, Test, type Context } from '@athenna/test' import { MissingKeyException } from '#src/exceptions/MissingKeyException' import { MissingRuleException } from '#src/exceptions/MissingRuleException' @@ -16,6 +25,8 @@ export class RateLimiterBuilderTest { @AfterEach() public async afterEach() { + await Cache.store('memory').truncate() + Config.clear() ioc.reconstruct() } @@ -132,6 +143,26 @@ export class RateLimiterBuilderTest { assert.isAtLeast(Date.now() - dateStart, 400) } + @Test() + public async shouldBeAbleToBuildARateLimiterWithSettingMultipleRules({ assert }: Context) { + const limiter = RateLimiter.build() + .store('memory', { windowMs: { second: 100 } }) + .key('request:api-key:/profile') + .setRules([{ type: 'second', limit: 1 }]) + + const promises = [] + const dateStart = Date.now() + const numberOfRequests = 5 + + for (let i = 0; i < numberOfRequests; i++) { + promises.push(limiter.schedule(() => 'ok' + i)) + } + + await Promise.all(promises) + + assert.isAtLeast(Date.now() - dateStart, 400) + } + @Test() public async shouldBeAbleToHaveErrorsHappeningInsideTheRateLimiterHandler({ assert }: Context) { const limiter = RateLimiter.build() @@ -383,12 +414,14 @@ export class RateLimiterBuilderTest { limiter.schedule(() => new Promise(() => {})) limiter.schedule(() => 'ok') - await Sleep.for(0).milliseconds().wait() + await this.waitUntil(() => limiter.getAvailableInMs() > 0, 10, 200) + + const availableInMs = limiter.getAvailableInMs() - assert.isAtLeast(limiter.getAvailableInMs(), 60) - assert.isAtMost(limiter.getAvailableInMs(), 120) + assert.isAtLeast(availableInMs, 50) + assert.isAtMost(availableInMs, 120) - await this.waitUntil(() => limiter.getAvailableInMs() === 0, 5, 800) + await this.waitUntil(() => limiter.getAvailableInMs() === 0, 10, 300) assert.equal(limiter.getAvailableInMs(), 0) } @@ -414,7 +447,7 @@ export class RateLimiterBuilderTest { abortController.abort('testing') await Sleep.for(30).milliseconds().wait() - assert.equal(limiter.getQueuedCount(), 2) + assert.equal(limiter.getQueuedCount(), 1) barrier.release() @@ -541,7 +574,38 @@ export class RateLimiterBuilderTest { .store('memory', { windowMs: { second: 100 } }) .key('request:api-key:/profile') .addRule({ type: 'second', limit: 1 }) - .addApiTarget({ id: Uuid.generate(), baseUrl: 'http://api1.com' }) + .addApiTarget({ metadata: { baseUrl: 'http://api1.com' } }) + + const promises = [] + const dateStart = Date.now() + const numberOfRequests = 5 + + for (let i = 0; i < numberOfRequests; i++) { + promises.push( + limiter.schedule(({ apiTarget }) => { + assert.isDefined(apiTarget) + + return 'ok' + i + }) + ) + } + + await Promise.all(promises) + + assert.isAtLeast(Date.now() - dateStart, 400) + } + + @Test() + public async shouldBeAbleToBuildARateLimiterWithARuleOfOneRequestPerSecondSettingMultipleApiTargets({ + assert + }: Context) { + assert.plan(6) + + const limiter = RateLimiter.build() + .store('memory', { windowMs: { second: 100 } }) + .key('request:api-key:/profile') + .addRule({ type: 'second', limit: 1 }) + .setApiTargets([{ metadata: { baseUrl: 'http://api1.com' } }]) const promises = [] const dateStart = Date.now() @@ -568,7 +632,7 @@ export class RateLimiterBuilderTest { .store('memory', { windowMs: { second: 100 } }) .key('request:api-key:/profile') .addRule({ type: 'second', limit: 1 }) - .addApiTarget({ id: Uuid.generate(), baseUrl: 'http://api1.com' }) + .addApiTarget({ metadata: { baseUrl: 'http://api1.com' } }) await assert.rejects(() => { return limiter.schedule(() => { @@ -585,14 +649,14 @@ export class RateLimiterBuilderTest { .store('memory', { windowMs: { second: 100 } }) .key('request:api-key:/profile') .addRule({ type: 'second', limit: 1 }) - .addApiTarget({ id: 't1', baseUrl: 'http://api1.com' }) - .addApiTarget({ id: 't2', baseUrl: 'http://api2.com' }) + .addApiTarget({ metadata: { baseUrl: 'http://api1.com' } }) + .addApiTarget({ metadata: { baseUrl: 'http://api2.com' } }) - const first = await limiter.schedule(({ apiTarget }) => apiTarget.baseUrl) + const first = await limiter.schedule(({ apiTarget }) => apiTarget.metadata.baseUrl) assert.equal(first, 'http://api1.com') - const second = await limiter.schedule(({ apiTarget }) => apiTarget.baseUrl) + const second = await limiter.schedule(({ apiTarget }) => apiTarget.metadata.baseUrl) assert.equal(second, 'http://api2.com') } @@ -606,18 +670,18 @@ export class RateLimiterBuilderTest { .store('memory', { windowMs: { second: 100 } }) .key('request:api-key:/profile') .addRule({ type: 'second', limit: 1 }) - .addApiTarget({ id: 't1', baseUrl: 'http://api1.com' }) - .addApiTarget({ id: 't2', baseUrl: 'http://api2.com' }) + .addApiTarget({ metadata: { baseUrl: 'http://api1.com' } }) + .addApiTarget({ metadata: { baseUrl: 'http://api2.com' } }) const barrier = this.createBarrier() const used: string[] = [] const run = async ({ apiTarget }) => { - used.push(apiTarget.baseUrl) + used.push(apiTarget.metadata.baseUrl) await barrier.wait() - return apiTarget.baseUrl + return apiTarget.metadata.baseUrl } const p1 = limiter.schedule(run) @@ -646,16 +710,16 @@ export class RateLimiterBuilderTest { .key('request:api-key:/profile') .apiTargetSelectionStrategy('round_robin') .addRule({ type: 'second', limit: 1 }) - .addApiTarget({ id: 't1', baseUrl: 'http://api1.com' }) - .addApiTarget({ id: 't2', baseUrl: 'http://api2.com' }) + .addApiTarget({ id: 't1', metadata: { baseUrl: 'http://api1.com' } }) + .addApiTarget({ id: 't2', metadata: { baseUrl: 'http://api2.com' } }) const used: string[] = [] const tasks = Array.from({ length: 4 }, () => limiter.schedule(({ apiTarget }) => { - used.push(apiTarget.baseUrl) + used.push(apiTarget.metadata.baseUrl) - return apiTarget.baseUrl + return apiTarget.metadata.baseUrl }) ) @@ -665,6 +729,28 @@ export class RateLimiterBuilderTest { assert.deepEqual(used, ['http://api1.com', 'http://api2.com', 'http://api1.com', 'http://api2.com']) } + @Test() + public async shouldThrowMissingRuleExceptionIfRateLimiterRulesAndApiTargetRulesAreNotDefined({ assert }: Context) { + const limiter = RateLimiter.build() + .store('memory', { windowMs: { second: 100 } }) + .key('request:api-key:/profile') + .addApiTarget({ id: 't1', metadata: { baseUrl: 'http://api1.com' } }) + + await assert.rejects(() => limiter.schedule(() => {}), MissingRuleException) + } + + @Test() + public async shouldNotThrowMissingRuleExceptionIfRateLimiterRulesAreNotDefinedButApiTargetRulesAreDefined({ + assert + }: Context) { + const limiter = RateLimiter.build() + .store('memory', { windowMs: { second: 100 } }) + .key('request:api-key:/profile') + .addApiTarget({ id: 't1', rules: [{ type: 'second', limit: 1 }], metadata: { baseUrl: 'http://api1.com' } }) + + await assert.doesNotReject(() => limiter.schedule(() => {}), MissingRuleException) + } + @Test() public async shouldBeAbleToTryWithTheSecondApiTargetIfTheFirstApiTargetIsAtFullCapacityInAConcurrentScenarioWithRoundRobinStrategy({ assert @@ -675,25 +761,25 @@ export class RateLimiterBuilderTest { .key('request:api-key:/profile') .apiTargetSelectionStrategy('round_robin') .addRule({ type: 'second', limit: 1 }) - .addApiTarget({ id: 't1', baseUrl: 'http://api1.com' }) - .addApiTarget({ id: 't2', baseUrl: 'http://api2.com' }) + .addApiTarget({ id: 't1', metadata: { baseUrl: 'http://api1.com' } }) + .addApiTarget({ id: 't2', metadata: { baseUrl: 'http://api2.com' } }) const barrier = this.createBarrier() const started: string[] = [] const p1 = limiter.schedule(async ({ apiTarget }) => { - started.push(apiTarget.baseUrl) + started.push(apiTarget.metadata.baseUrl) await barrier.wait() - return apiTarget.baseUrl + return apiTarget.metadata.baseUrl }) const p2 = limiter.schedule(async ({ apiTarget }) => { - started.push(apiTarget.baseUrl) + started.push(apiTarget.metadata.baseUrl) await barrier.wait() - return apiTarget.baseUrl + return apiTarget.metadata.baseUrl }) for (let i = 0; i < 20 && started.length < 2; i++) { @@ -710,4 +796,94 @@ export class RateLimiterBuilderTest { assert.deepEqual(result, ['http://api1.com', 'http://api2.com']) } + + @Test() + public async shouldBeAbleToBuildARateLimiterDefiningRulesInApiTarget({ assert }: Context) { + assert.plan(6) + + const limiter = RateLimiter.build() + .store('memory', { windowMs: { second: 100 } }) + .key('request:api-key:/profile') + .addApiTarget({ rules: [{ type: 'second', limit: 1 }], metadata: { baseUrl: 'http://api0.com' } }) + + const promises = [] + const dateStart = Date.now() + const numberOfRequests = 5 + + for (let i = 0; i < numberOfRequests; i++) { + promises.push( + limiter.schedule(({ apiTarget }) => { + assert.isDefined(apiTarget) + + return 'ok' + i + }) + ) + } + + await Promise.all(promises) + + assert.isAtLeast(Date.now() - dateStart, 400) + } + + @Test() + public async shouldNotBeAbleToRetryRequestWithoutARetryStrategy({ assert }: Context) { + const limiter = RateLimiter.build() + .store('memory', { windowMs: { second: 100 } }) + .key('request:api-key:/profile') + .addApiTarget({ rules: [{ type: 'second', limit: 1 }], metadata: { baseUrl: 'http://api0.com' } }) + + await assert.rejects(() => + limiter.schedule(({ apiTarget }) => { + throw new Error(apiTarget.metadata.baseUrl) + }) + ) + } + + @Test() + public async shouldAlwaysFailTheRequestIfRetryStrategyDecideItShouldFail({ assert }: Context) { + const limiter = RateLimiter.build() + .store('memory', { windowMs: { second: 100 } }) + .key('request:api-key:/profile') + .addApiTarget({ rules: [{ type: 'second', limit: 1 }], metadata: { baseUrl: 'http://api0.com' } }) + .retryStrategy(ctx => { + if (ctx.error.message === 'fail') { + return { type: 'fail' } + } + }) + + await assert.rejects(() => + limiter.schedule(() => { + throw new Error('fail') + }) + ) + } + + @Test() + public async shouldAlwaysCooldownAndRetryTheRequestWithTheSameApiIfRetryStrategyDecideItShouldWaitToTryAgainWithSame({ + assert + }: Context) { + const limiter = RateLimiter.build() + .key('request:api-key:/profile') + .store('memory', { windowMs: { second: 100 } }) + .addApiTarget({ rules: [{ type: 'second', limit: 1 }], metadata: { baseUrl: 'http://api0.com' } }) + .addApiTarget({ rules: [{ type: 'second', limit: 1 }], metadata: { baseUrl: 'http://api1.com' } }) + .retryStrategy(() => { + return { type: 'cooldown', cooldownMs: 100, then: 'retry_same' } + }) + + const apiTargetUsed = [] + let isFirstRequest = true + + await limiter.schedule(({ apiTarget }) => { + apiTargetUsed.push(apiTarget.metadata.baseUrl) + + if (isFirstRequest) { + isFirstRequest = false + + throw new Error('fail') + } + }) + + assert.deepEqual(apiTargetUsed, ['http://api0.com', 'http://api0.com']) + } } From ca06ceacb2f4422e73adaaf2957cf1c961350322 Mon Sep 17 00:00:00 2001 From: jlenon7 Date: Mon, 8 Sep 2025 22:39:35 -0300 Subject: [PATCH 3/6] test(targets): add tests for retry strategy, retry_same, cooldown decisions --- src/ratelimiter/RateLimitStore.ts | 4 + src/ratelimiter/RateLimiterBuilder.ts | 896 +++++++++++------- .../ratelimiter/RateLimiterBuilderTest.ts | 39 +- 3 files changed, 570 insertions(+), 369 deletions(-) diff --git a/src/ratelimiter/RateLimitStore.ts b/src/ratelimiter/RateLimitStore.ts index d90e262..881bf86 100644 --- a/src/ratelimiter/RateLimitStore.ts +++ b/src/ratelimiter/RateLimitStore.ts @@ -28,6 +28,10 @@ export class RateLimitStore extends Macroable { this.options = options } + public async truncate() { + await Cache.store(this.options.store).truncate() + } + public async getOrInit(key: string, rules: RateLimitRule[]) { const cache = Cache.store(this.options.store) diff --git a/src/ratelimiter/RateLimiterBuilder.ts b/src/ratelimiter/RateLimiterBuilder.ts index 212ad8b..7758f32 100644 --- a/src/ratelimiter/RateLimiterBuilder.ts +++ b/src/ratelimiter/RateLimiterBuilder.ts @@ -20,8 +20,8 @@ import type { } from '#src/types' import { Config } from '@athenna/config' -import { Json, String, Macroable } from '@athenna/common' import { RateLimitStore } from '#src/ratelimiter/RateLimitStore' +import { Json, String, Macroable, Options } from '@athenna/common' import { MissingKeyException } from '#src/exceptions/MissingKeyException' import { MissingRuleException } from '#src/exceptions/MissingRuleException' import { MissingStoreException } from '#src/exceptions/MissingStoreException' @@ -70,26 +70,18 @@ export class RateLimiterBuilder extends Macroable { */ private nextWakeUpAt: number = 0 - /** - * Create a custom id for an API Target by reading the metadata object. - * The object will always be sorted by keys. - */ - private getApiTargetId(apiTarget: RateLimitApiTarget) { - return String.hash(JSON.stringify(Json.sort(apiTarget.metadata)), { - key: Config.get('app.key', 'ratelimiter') - }) - } - - /** - * Create a custom key for an API Target to be used to map the - * API Target rules into the cache. - */ - private createApiTargetKey(apiTarget: RateLimitApiTarget) { - return `${this.options.key}:${this.getApiTargetId(apiTarget)}` - } - /** * Logical key that will be used by store to save buckets. + * + * @example + * ```ts + * const limiter = RateLimiter.build() + * .store('memory') + * .addRule({ type: 'second', limit: 1 }) + * .key('request:/profile') + * + * await limiter.schedule(() => {...}) + * ``` */ public key(value: string) { this.options.key = value @@ -100,6 +92,16 @@ export class RateLimiterBuilder extends Macroable { /** * Define the store that will be responsible to save the * rate limit buckets. + * + * @example + * ```ts + * const limiter = RateLimiter.build() + * .key('request:/profile') + * .addRule({ type: 'second', limit: 1 }) + * .store('memory') + * + * await limiter.schedule(() => {...}) + * ``` */ public store( store: 'memory' | 'redis' | string, @@ -116,6 +118,17 @@ export class RateLimiterBuilder extends Macroable { /** * Set the max number of tasks that could run concurrently. + * + * @example + * ```ts + * const limiter = RateLimiter.build() + * .store('memory') + * .key('request:/profile') + * .addRule({ type: 'second', limit: 1 }) + * .maxConcurrent(10) + * + * await limiter.schedule(() => {...}) + * ``` */ public maxConcurrent(value: number) { this.options.maxConcurrent = Math.max(1, value ?? 1) @@ -126,6 +139,17 @@ export class RateLimiterBuilder extends Macroable { /** * Random jitter in milliseconds to avoid thundering herd in * distributed environments. + * + * @example + * ```ts + * const limiter = RateLimiter.build() + * .store('memory') + * .key('request:/profile') + * .addRule({ type: 'second', limit: 1 }) + * .randomJitter(1000) + * + * await limiter.schedule(() => {...}) + * ``` */ public jitterMs(value: number) { this.options.jitterMs = Math.max(0, value ?? 0) @@ -135,6 +159,16 @@ export class RateLimiterBuilder extends Macroable { /** * Add a new rate limit rule. + * + * @example + * ```ts + * const limiter = RateLimiter.build() + * .store('memory') + * .key('request:/profile') + * .addRule({ type: 'second', limit: 1 }) + * + * await limiter.schedule(() => {...}) + * ``` */ public addRule(rule: RateLimitRule) { if (!this.options.rules) { @@ -148,6 +182,17 @@ export class RateLimiterBuilder extends Macroable { /** * Add a new rate limit API target. + * + * @example + * ```ts + * const limiter = RateLimiter.build() + * .store('memory') + * .key('request:/profile') + * .addRule({ type: 'second', limit: 1 }) + * .addApiTarget({ metadata: { baseUrl: 'http://example.com' } }) + * + * await limiter.schedule(() => {...}) + * ``` */ public addApiTarget(apiTarget: RateLimitApiTarget) { if (!this.options.apiTargets) { @@ -165,6 +210,16 @@ export class RateLimiterBuilder extends Macroable { /** * Set multiple rate limit rules with one method call. + * + * @example + * ```ts + * const limiter = RateLimiter.build() + * .store('memory') + * .key('request:/profile') + * .setRules([{ type: 'second', limit: 1 }]) + * + * await limiter.schedule(() => {...}) + * ``` */ public setRules(rules: RateLimitRule[]) { rules.forEach(rule => this.addRule(rule)) @@ -174,6 +229,17 @@ export class RateLimiterBuilder extends Macroable { /** * Set multiple rate limit API targets with one method call. + * + * @example + * ```ts + * const limiter = RateLimiter.build() + * .store('memory') + * .key('request:/profile') + * .addRule({ type: 'second', limit: 1 }) + * .setApiTargets([{ metadata: { baseUrl: 'http://example.com' } }]) + * + * await limiter.schedule(() => {...}) + * ``` */ public setApiTargets(apiTargets: RateLimitApiTarget[]) { apiTargets.forEach(apiTarget => this.addApiTarget(apiTarget)) @@ -183,7 +249,19 @@ export class RateLimiterBuilder extends Macroable { /** * Define the API target selection strategy that will be used - * to select the next one when an API fails. + * to select the next one when an API Target fails. + * + * @example + * ```ts + * const limiter = RateLimiter.build() + * .store('memory') + * .key('request:/profile') + * .addRule({ type: 'second', limit: 1 }) + * .addApiTarget({ metadata: { baseUrl: 'http://example.com' } }) + * .apiTargetSelectionStrategy('round_robin') + * + * await limiter.schedule(() => {...}) + * ``` */ public apiTargetSelectionStrategy(value: 'first_available' | 'round_robin') { this.options.apiTargetSelectionStrategy = value @@ -193,7 +271,29 @@ export class RateLimiterBuilder extends Macroable { /** * Define the RateLmiter retry strategy. This is useful to control - * when and how we should proceed with the retry of API Targets. + * when and how we should proceed with the retry of tasks that failed + * to execute. + * + * @example + * ```ts + * const limiter = RateLimiter.build() + * .store('memory') + * .key('request:/profile') + * .addRule({ type: 'second', limit: 1 }) + * .retryStrategy(({ attempt }) => { + * const decision = { type: 'fail' } + * + * if (attempt === 3) { + * return decision + * } + * + * decision.type = 'retry_same' + * + * return decision + * }) + * + * await limiter.schedule(() => {...}) + * ``` */ public retryStrategy( fn: ( @@ -261,9 +361,20 @@ export class RateLimiterBuilder extends Macroable { } /** - * Drop all the tasks that are in the queue. + * Drop all the tasks that are in the queue and clear + * store. + * + * @example + * ```ts + * const limiter = RateLimiter.build() + * .store('memory') + * .key('request:/profile') + * .addRule({ type: 'second', limit: 1 }) + * + * await limiter.truncate() + * ``` */ - public truncate() { + public async truncate() { this.queue = [] this.active = 0 this.nextWakeUpAt = 0 @@ -274,12 +385,26 @@ export class RateLimiterBuilder extends Macroable { this.timer = null } + await this.options.store.truncate() + return this } /** * Schedule the execution of an async function respecting - * the rate limit rules. + * the rate limit rules and the API Targets. + * + * @example + * ```ts + * const limiter = RateLimiter.build() + * .store('memory') + * .key('request:/profile') + * .addRule({ type: 'second', limit: 1 }) + * + * const response = await limiter.schedule(() => { + * return fetch('http://example.com') + * }) + * ``` */ public schedule( closure: (ctx: RateLimitScheduleCtx) => T | Promise, @@ -347,478 +472,525 @@ export class RateLimiterBuilder extends Macroable { item.abortHandler = onAbort } - this.pump() + this.scheduleQueueItemRun() }) } /** - * Process the queue of tasks. + * Create a custom id for an API Target by reading the metadata object. + * The object will always be sorted by keys. */ - private pump() { - if (this.timer) { - return + private getApiTargetId(apiTarget: RateLimitApiTarget) { + return String.hash(JSON.stringify(Json.sort(apiTarget.metadata)), { + key: Config.get('app.key', 'ratelimiter') + }) + } + + /** + * Create a custom key for an API Target to be used to map the + * API Target rules into the cache. + */ + private createApiTargetKey(apiTarget: RateLimitApiTarget) { + return `${this.options.key}:${this.getApiTargetId(apiTarget)}` + } + + /** + * Get a random jitter or return 0 if user has not + * defined one. + */ + private randomJitter(): number { + if (!this.options.jitterMs) { + return 0 } - const tryRun = async () => { - this.timer = null + return Math.floor(Math.random() * this.options.jitterMs) + } - if (this.active >= this.options.maxConcurrent) { - return + /** + * Read the API Target selection strategy and defines which is + * going to be used. + */ + private createIdxBySelectionStrategy(pinnedApiTargetId?: string) { + if (pinnedApiTargetId) { + const i = this.options.apiTargets.findIndex( + a => a.id === pinnedApiTargetId + ) + + if (i >= 0) { + return [i] } + } - if (this.queue.length === 0) { - return + let indexes = [] + + switch (this.options.apiTargetSelectionStrategy) { + case 'round_robin': + indexes = this.createRoundRobinIdx() + break + case 'first_available': + default: + indexes = this.createFirstAvailableIdx() + } + + return indexes + } + + /** + * Create the indexes for when using round_robin selection strategy. + */ + private createRoundRobinIdx() { + return Array.from( + { length: this.options.apiTargets.length }, + (_, k) => (this.rrIndex + k) % this.options.apiTargets.length + ) + } + + /** + * Create the indexes for when using first_available selection strategy. + */ + private createFirstAvailableIdx() { + return Array.from({ length: this.options.apiTargets.length }, (_, k) => k) + } + + /** + * Release the rate limit task. + */ + private releaseTask(options: { isToScheduleTick: boolean }) { + this.active-- + + if (options.isToScheduleTick) { + if (this.timer) { + clearTimeout(this.timer) + + this.timer = null } - const now = Date.now() + this.scheduleQueueItemRun() + } + } - if (this.options.apiTargets?.length) { - let minWait = Number.POSITIVE_INFINITY - let apiTargetChosen: RateLimitApiTarget = null + /** + * Try process an item from the queue of tasks. + */ + private tryToRunQueueItem = async () => { + this.timer = null - const nextItem = this.queue[0] - const pinnedApiTargetId = nextItem?.pinnedApiTargetId + if (this.active >= this.options.maxConcurrent) { + return + } - for (const i of this.createIdxBySelectionStrategy(pinnedApiTargetId)) { - const apiTarget = this.options.apiTargets[i] - const key = this.createApiTargetKey(apiTarget) + if (this.queue.length === 0) { + return + } - const rules = apiTarget.rules?.length - ? apiTarget.rules - : this.options.rules + const now = Date.now() - try { - const res = await this.options.store.tryReserve(key, rules) + if (this.options.apiTargets?.length) { + let minWait = Number.POSITIVE_INFINITY + let apiTarget: RateLimitApiTarget = null - this.storeErrorCount = 0 + const nextItem = this.queue[0] + const pinnedApiTargetId = nextItem?.pinnedApiTargetId - if (res.allowed) { - apiTargetChosen = apiTarget + for (const i of this.createIdxBySelectionStrategy(pinnedApiTargetId)) { + const key = this.createApiTargetKey(this.options.apiTargets[i]) - if (this.options.apiTargetSelectionStrategy === 'round_robin') { - this.rrIndex = (i + 1) % this.options.apiTargets.length - } + const rules = this.options.apiTargets[i].rules?.length + ? this.options.apiTargets[i].rules + : this.options.rules - break - } + try { + const res = await this.options.store.tryReserve(key, rules) - minWait = Math.min(minWait, res.waitMs) - } catch (error) { - this.storeErrorCount++ + this.storeErrorCount = 0 - if (this.storeErrorCount > 10) { - while (this.queue.length) { - this.queue.shift().reject(error) - } + if (res.allowed) { + apiTarget = this.options.apiTargets[i] - throw error + if (this.options.apiTargetSelectionStrategy === 'round_robin') { + this.rrIndex = (i + 1) % this.options.apiTargets.length } - minWait = Math.min(minWait, 100) + break } - } - if (!apiTargetChosen) { - const delay = - (isFinite(minWait) ? minWait : 100) + this.randomJitter() + minWait = Math.min(minWait, res.waitMs) + } catch (error) { + this.storeErrorCount++ - this.nextWakeUpAt = now + delay - this.timer = setTimeout(tryRun, delay) + if (this.storeErrorCount > 10) { + while (this.queue.length) { + this.queue.shift().reject(error) + } - return + throw error + } + + minWait = Math.min(minWait, 100) } + } - const item = this.queue.shift() + if (!apiTarget) { + const delay = (isFinite(minWait) ? minWait : 100) + this.randomJitter() - if (item.signal?.aborted) { - item.reject(new DOMException('Aborted', 'AbortError')) - this.pump() - return - } + this.nextWakeUpAt = now + delay + this.scheduleQueueItemRun({ delay }) - item.started = true + return + } - if (item.signal && item.abortHandler) { - item.signal.removeEventListener('abort', item.abortHandler) - item.abortHandler = undefined - } + const item = this.queue.shift() - this.active++ + if (item.signal?.aborted) { + item.reject(new DOMException('Aborted', 'AbortError')) - Promise.resolve() - .then(() => - item.run({ signal: item.signal, apiTarget: apiTargetChosen }) - ) - .then(result => { - this.releaseTask({ isToPump: true }) + this.scheduleQueueItemRun() - item.resolve(result) - }) - .catch(async error => { - if (!this.options.retryStrategy) { - this.releaseTask({ isToPump: true }) + return + } - item.reject(error) + item.started = true - return - } + if (item.signal && item.abortHandler) { + item.signal.removeEventListener('abort', item.abortHandler) + item.abortHandler = undefined + } - const key = this.createApiTargetKey(apiTargetChosen) + this.active++ - const ctx: RateLimitRetryCtx = { - key, - error, - signal: item.signal, - attempt: item.attempt, - apiTarget: apiTargetChosen - } + Promise.resolve() + .then(() => item.run({ signal: item.signal, apiTarget })) + .then(result => { + this.releaseTask({ isToScheduleTick: true }) - const decision = await this.options.retryStrategy(ctx) + item.resolve(result) + }) + .catch(error => this.onFailInMultiMode({ error, item, apiTarget })) - switch (decision.type) { - case 'retry_same': - this.releaseTask({ isToPump: false }) + if (this.active < this.options.maxConcurrent) { + this.scheduleQueueItemRun() + } - item.attempt++ - item.started = false - item.pinnedApiTargetId = apiTargetChosen.id + return + } - if (item.signal?.aborted) { - item.reject(new DOMException('Aborted', 'AbortError')) - break - } + let waitMs = 0 + let allowed = false - this.queue.unshift(item) + try { + const res = await this.options.store.tryReserve( + this.options.key, + this.options.rules + ) - const delay = Math.max(0, decision.delayMs ?? 0) + this.storeErrorCount = 0 + + allowed = res.allowed + waitMs = res.waitMs + } catch (error) { + this.storeErrorCount++ + + /** + * If the store failed 10 times it means it is not working for some + * reason, in this case we can reject all the requests that are in + * the queue. + */ + if (this.storeErrorCount > 10) { + while (this.queue.length) { + this.queue.shift().reject(error) + } - if (delay > 0) { - this.nextWakeUpAt = Date.now() + delay - this.timer = setTimeout(tryRun, delay) - } else { - this.pump() - } + throw error + } - return + allowed = false + waitMs = 100 + } - case 'retry_other': { - const cooldown = Math.max(0, decision.cooldownMs || 0) + if (!allowed) { + const delay = waitMs + this.randomJitter() - this.releaseTask({ isToPump: false }) + this.nextWakeUpAt = now + delay + this.scheduleQueueItemRun({ delay }) - if (cooldown > 0) { - await this.options.store!.setCooldown(key, cooldown) - } + return + } - item.attempt++ - item.started = false - item.pinnedApiTargetId = undefined + const item = this.queue.shift() - if (item.signal?.aborted) { - item.reject(new DOMException('Aborted', 'AbortError')) - break - } + if (item.signal?.aborted) { + item.reject(new DOMException('Aborted', 'AbortError')) - this.queue.unshift(item) + this.scheduleQueueItemRun() - const delay = Math.max(0, decision.delayMs ?? 0) + return + } - if (delay > 0) { - this.nextWakeUpAt = Date.now() + delay - this.timer = setTimeout(tryRun, delay) - } else { - this.pump() - } + item.started = true - return - } + if (item.signal && item.abortHandler) { + item.signal.removeEventListener('abort', item.abortHandler) + item.abortHandler = undefined + } - case 'cooldown': - const ms = Math.max(0, decision.cooldownMs || 0) + this.active++ - this.releaseTask({ isToPump: false }) + Promise.resolve() + .then(() => item.run({ signal: item.signal })) + .then(result => { + this.releaseTask({ isToScheduleTick: true }) - await this.options.store!.setCooldown(key, ms) + item.resolve(result) + }) + .catch(error => this.onFailInSingleMode({ error, item })) - const then = decision.then + if (this.active < this.options.maxConcurrent) { + this.scheduleQueueItemRun() + } + } - if (then === 'retry_same' || then === 'retry_other') { - item.attempt++ - item.started = false - item.pinnedApiTargetId = - decision.then === 'retry_same' - ? apiTargetChosen.id - : undefined + /** + * Schedule to run another queue item. + */ + private scheduleQueueItemRun = (options?: { delay?: number }) => { + options = Options.create(options, { + delay: 0 + }) - if (item.signal?.aborted) { - item.reject(new DOMException('Aborted', 'AbortError')) + if (this.timer) { + return + } - break - } + const fire = async () => { + this.timer = null - this.queue.unshift(item) + await this.tryToRunQueueItem() + } - const delay = decision.cooldownMs + this.timer = setTimeout(fire, options.delay) + } - this.nextWakeUpAt = Date.now() + delay - this.timer = setTimeout(tryRun, delay) + /** + * Closure that deals with all the errors that happens when when running + * RateLimiter in single-mode. + */ + private onFailInSingleMode = async (options: { + error: Error + item: QueueItem + }) => { + if (!this.options.retryStrategy) { + this.releaseTask({ isToScheduleTick: true }) - return - } + options.item.reject(options.error) - this.releaseTask({ isToPump: true }) + return + } - item.reject(error) + const ctx: RateLimitRetryCtx = { + error: options.error, + key: this.options.key, + attempt: options.item.attempt + } - break - default: - this.releaseTask({ isToPump: true }) + const decision = await this.options.retryStrategy(ctx) - item.reject(error) - } - }) + switch (decision.type) { + case 'retry_same': + case 'retry_other': + const delay = Math.max(0, decision.delayMs ?? 0) - if (this.active < this.options.maxConcurrent) { - this.timer = setTimeout(tryRun, 0) - } + this.releaseTask({ isToScheduleTick: false }) - return - } + options.item.attempt++ + options.item.started = false - let waitMs = 0 - let allowed = false + if (options.item.signal?.aborted) { + options.item.reject(new DOMException('Aborted', 'AbortError')) - try { - const res = await this.options.store.tryReserve( - this.options.key, - this.options.rules - ) + break + } - this.storeErrorCount = 0 - allowed = res.allowed - waitMs = res.waitMs - } catch (error) { - this.storeErrorCount++ - - /** - * If the store failed 10 times it means it is not working for some - * reason, in this case we can reject all the requests that are in - * the queue. - */ - if (this.storeErrorCount > 10) { - while (this.queue.length) { - this.queue.shift().reject(error) - } + this.queue.unshift(options.item) - throw error + if (delay > 0) { + this.nextWakeUpAt = Date.now() + delay + this.scheduleQueueItemRun({ delay }) + } else { + this.scheduleQueueItemRun() } - allowed = false - waitMs = 100 - } + return - if (!allowed) { - const delay = waitMs + this.randomJitter() + case 'cooldown': { + await this.options.store!.setCooldown( + this.options.key, + Math.max(0, decision.cooldownMs) + ) - this.nextWakeUpAt = now + delay - this.timer = setTimeout(tryRun, delay) + if (decision.then === 'retry_same' || decision.then === 'retry_other') { + const delay = decision.cooldownMs - return - } + this.releaseTask({ isToScheduleTick: false }) - const item = this.queue.shift() + options.item.attempt++ + options.item.started = false - if (item.signal?.aborted) { - item.reject(new DOMException('Aborted', 'AbortError')) + if (options.item.signal?.aborted) { + options.item.reject(new DOMException('Aborted', 'AbortError')) - this.pump() + break + } - return - } + this.queue.unshift(options.item) + this.nextWakeUpAt = Date.now() + delay + this.scheduleQueueItemRun({ delay }) - item.started = true + return + } - if (item.signal && item.abortHandler) { - item.signal.removeEventListener('abort', item.abortHandler) - item.abortHandler = undefined - } + this.releaseTask({ isToScheduleTick: true }) - this.active++ + options.item.reject(options.error) - Promise.resolve() - .then(() => item.run({ signal: item.signal })) - .then(result => { - this.releaseTask({ isToPump: true }) + break + } - item.resolve(result) - }) - .catch(async error => { - if (!this.options.retryStrategy) { - this.releaseTask({ isToPump: true }) + default: + this.releaseTask({ isToScheduleTick: true }) - item.reject(error) + options.item.reject(options.error) + } + } - return - } + /** + * Closure that deals with all the errors that happens when when running + * RateLimiter in multi-mode. + */ + private onFailInMultiMode = async (options: { + error: Error + item: QueueItem + apiTarget: RateLimitApiTarget + }) => { + if (!this.options.retryStrategy) { + this.releaseTask({ isToScheduleTick: true }) - const ctx: RateLimitRetryCtx = { - error, - key: this.options.key, - attempt: item.attempt - } + options.item.reject(options.error) - const decision = await this.options.retryStrategy(ctx) + return + } - switch (decision.type) { - case 'retry_same': - case 'retry_other': - const delay = Math.max(0, decision.delayMs ?? 0) + const key = this.createApiTargetKey(options.apiTarget) - this.releaseTask({ isToPump: false }) + const ctx: RateLimitRetryCtx = { + key, + error: options.error, + signal: options.item.signal, + attempt: options.item.attempt, + apiTarget: options.apiTarget + } - item.attempt++ - item.started = false + const decision = await this.options.retryStrategy(ctx) - if (item.signal?.aborted) { - item.reject(new DOMException('Aborted', 'AbortError')) - break - } + switch (decision.type) { + case 'retry_same': + this.releaseTask({ isToScheduleTick: false }) - this.queue.unshift(item) + options.item.attempt++ + options.item.started = false + options.item.pinnedApiTargetId = options.apiTarget.id - if (delay > 0) { - this.nextWakeUpAt = Date.now() + delay - this.timer = setTimeout(tryRun, delay) - } else { - this.pump() - } + if (options.item.signal?.aborted) { + options.item.reject(new DOMException('Aborted', 'AbortError')) - return + break + } - case 'cooldown': { - await this.options.store!.setCooldown( - this.options.key, - Math.max(0, decision.cooldownMs) - ) + this.queue.unshift(options.item) - if ( - decision.then === 'retry_same' || - decision.then === 'retry_other' - ) { - const delay = decision.cooldownMs + const delay = Math.max(0, decision.delayMs ?? 0) - this.releaseTask({ isToPump: false }) + if (delay > 0) { + this.nextWakeUpAt = Date.now() + delay + this.scheduleQueueItemRun({ delay }) + } else { + this.scheduleQueueItemRun() + } - item.attempt++ - item.started = false + return - if (item.signal?.aborted) { - item.reject(new DOMException('Aborted', 'AbortError')) + case 'retry_other': { + const cooldown = Math.max(0, decision.cooldownMs || 0) - break - } + this.releaseTask({ isToScheduleTick: false }) - this.queue.unshift(item) - this.nextWakeUpAt = Date.now() + delay - this.timer = setTimeout(tryRun, delay) + if (cooldown > 0) { + await this.options.store!.setCooldown(key, cooldown) + } - return - } + options.item.attempt++ + options.item.started = false + options.item.pinnedApiTargetId = undefined - this.releaseTask({ isToPump: true }) + if (options.item.signal?.aborted) { + options.item.reject(new DOMException('Aborted', 'AbortError')) - item.reject(error) + break + } - break - } + this.queue.unshift(options.item) - default: - this.releaseTask({ isToPump: true }) + const delay = Math.max(0, decision.delayMs ?? 0) - item.reject(error) - } - }) + if (delay > 0) { + this.nextWakeUpAt = Date.now() + delay + this.scheduleQueueItemRun({ delay }) + } else { + this.scheduleQueueItemRun() + } - if (this.active < this.options.maxConcurrent) { - this.timer = setTimeout(tryRun, 0) + return } - } - this.timer = setTimeout(tryRun, 0) - } + case 'cooldown': + const ms = Math.max(0, decision.cooldownMs || 0) - /** - * Get a random jitter or return 0 if user has not - * defined one. - */ - private randomJitter(): number { - if (!this.options.jitterMs) { - return 0 - } + this.releaseTask({ isToScheduleTick: false }) - return Math.floor(Math.random() * this.options.jitterMs) - } + await this.options.store!.setCooldown(key, ms) - /** - * Read the API Target selection strategy and defines which is - * going to be used. - */ - private createIdxBySelectionStrategy(pinnedApiTargetId?: string) { - let indexes = [] + const then = decision.then - switch (this.options.apiTargetSelectionStrategy) { - case 'round_robin': - indexes = this.createRoundRobinIdx() - break - case 'first_available': - default: - indexes = this.createFirstAvailableIdx() - } + if (then === 'retry_same' || then === 'retry_other') { + options.item.attempt++ + options.item.started = false + options.item.pinnedApiTargetId = + decision.then === 'retry_same' ? options.apiTarget.id : undefined - if (pinnedApiTargetId) { - const i = this.options.apiTargets.findIndex( - a => a.id === pinnedApiTargetId - ) + if (options.item.signal?.aborted) { + options.item.reject(new DOMException('Aborted', 'AbortError')) - if (i >= 0) { - indexes = [i, ...indexes.filter(x => x !== i)] - } - } + break + } - return indexes - } + this.queue.unshift(options.item) - /** - * Create the indexes for when using round_robin selection strategy. - */ - private createRoundRobinIdx() { - return Array.from( - { length: this.options.apiTargets.length }, - (_, k) => (this.rrIndex + k) % this.options.apiTargets.length - ) - } + const delay = decision.cooldownMs - /** - * Create the indexes for when using first_available selection strategy. - */ - private createFirstAvailableIdx() { - return Array.from({ length: this.options.apiTargets.length }, (_, k) => k) - } + this.nextWakeUpAt = Date.now() + delay + this.scheduleQueueItemRun({ delay }) - /** - * Release the rate limit task. - */ - private releaseTask(options: { isToPump: boolean }) { - this.active-- + return + } - if (options.isToPump) { - if (this.timer) { - clearTimeout(this.timer) + this.releaseTask({ isToScheduleTick: true }) - this.timer = null - } + options.item.reject(options.error) + + break + default: + this.releaseTask({ isToScheduleTick: true }) - this.pump() + options.item.reject(options.error) } } } diff --git a/tests/unit/ratelimiter/RateLimiterBuilderTest.ts b/tests/unit/ratelimiter/RateLimiterBuilderTest.ts index 73c6c0c..e8336e3 100644 --- a/tests/unit/ratelimiter/RateLimiterBuilderTest.ts +++ b/tests/unit/ratelimiter/RateLimiterBuilderTest.ts @@ -309,7 +309,7 @@ export class RateLimiterBuilderTest { limiter.schedule(() => 'ok' + i) } - limiter.truncate() + await limiter.truncate() assert.equal(limiter.getActiveCount(), 0) assert.equal(limiter.getQueuedCount(), 0) @@ -330,8 +330,6 @@ export class RateLimiterBuilderTest { } assert.equal(limiter.getQueuedCount(), 5) - - limiter.truncate() } @Test() @@ -845,10 +843,8 @@ export class RateLimiterBuilderTest { .store('memory', { windowMs: { second: 100 } }) .key('request:api-key:/profile') .addApiTarget({ rules: [{ type: 'second', limit: 1 }], metadata: { baseUrl: 'http://api0.com' } }) - .retryStrategy(ctx => { - if (ctx.error.message === 'fail') { - return { type: 'fail' } - } + .retryStrategy(() => { + return { type: 'fail' } }) await assert.rejects(() => @@ -858,6 +854,35 @@ export class RateLimiterBuilderTest { ) } + @Test() + public async shouldAlwaysRetryTheRequestWithTheSameApiIfRetryStrategyDecideItShouldRetryWithTheSame({ + assert + }: Context) { + const limiter = RateLimiter.build() + .key('request:api-key:/profile') + .store('memory', { windowMs: { second: 100 } }) + .addApiTarget({ rules: [{ type: 'second', limit: 1 }], metadata: { baseUrl: 'http://api0.com' } }) + .addApiTarget({ rules: [{ type: 'second', limit: 1 }], metadata: { baseUrl: 'http://api1.com' } }) + .retryStrategy(() => { + return { type: 'retry_same' } + }) + + const apiTargetUsed = [] + let isFirstRequest = true + + await limiter.schedule(({ apiTarget }) => { + apiTargetUsed.push(apiTarget.metadata.baseUrl) + + if (isFirstRequest) { + isFirstRequest = false + + throw new Error('fail') + } + }) + + assert.deepEqual(apiTargetUsed, ['http://api0.com', 'http://api0.com']) + } + @Test() public async shouldAlwaysCooldownAndRetryTheRequestWithTheSameApiIfRetryStrategyDecideItShouldWaitToTryAgainWithSame({ assert From 397fb5b8dde53c77dbb63fceb23fa2666f3aff67 Mon Sep 17 00:00:00 2001 From: jlenon7 Date: Tue, 9 Sep 2025 11:40:33 -0300 Subject: [PATCH 4/6] test(targets): add tests for missing decisions --- src/ratelimiter/RateLimitStore.ts | 3 + src/ratelimiter/RateLimiterBuilder.ts | 151 +++++--------- src/types/QueueItem.ts | 1 + src/types/RateLimitRetryDecision.ts | 54 ++++- .../ratelimiter/RateLimiterBuilderTest.ts | 185 +++++++++++++++++- 5 files changed, 276 insertions(+), 118 deletions(-) diff --git a/src/ratelimiter/RateLimitStore.ts b/src/ratelimiter/RateLimitStore.ts index 881bf86..e62f7a8 100644 --- a/src/ratelimiter/RateLimitStore.ts +++ b/src/ratelimiter/RateLimitStore.ts @@ -32,6 +32,9 @@ export class RateLimitStore extends Macroable { await Cache.store(this.options.store).truncate() } + /** + * Get the rate limit buckets from the cache or initialize them. + */ public async getOrInit(key: string, rules: RateLimitRule[]) { const cache = Cache.store(this.options.store) diff --git a/src/ratelimiter/RateLimiterBuilder.ts b/src/ratelimiter/RateLimiterBuilder.ts index 7758f32..65bb271 100644 --- a/src/ratelimiter/RateLimiterBuilder.ts +++ b/src/ratelimiter/RateLimiterBuilder.ts @@ -19,6 +19,7 @@ import type { RateLimitRetryDecision } from '#src/types' +import { debug } from '#src/debug' import { Config } from '@athenna/config' import { RateLimitStore } from '#src/ratelimiter/RateLimitStore' import { Json, String, Macroable, Options } from '@athenna/common' @@ -510,10 +511,10 @@ export class RateLimiterBuilder extends Macroable { * Read the API Target selection strategy and defines which is * going to be used. */ - private createIdxBySelectionStrategy(pinnedApiTargetId?: string) { - if (pinnedApiTargetId) { + private createIdxBySelectionStrategy(item: QueueItem) { + if (item.pinnedApiTargetId) { const i = this.options.apiTargets.findIndex( - a => a.id === pinnedApiTargetId + a => a.id === item.pinnedApiTargetId ) if (i >= 0) { @@ -532,6 +533,14 @@ export class RateLimiterBuilder extends Macroable { indexes = this.createFirstAvailableIdx() } + if (item.avoidApiTargetId) { + const i = this.options.apiTargets.findIndex( + a => a.id === item.avoidApiTargetId + ) + + return indexes.filter(idx => idx !== i) + } + return indexes } @@ -590,9 +599,8 @@ export class RateLimiterBuilder extends Macroable { let apiTarget: RateLimitApiTarget = null const nextItem = this.queue[0] - const pinnedApiTargetId = nextItem?.pinnedApiTargetId - for (const i of this.createIdxBySelectionStrategy(pinnedApiTargetId)) { + for (const i of this.createIdxBySelectionStrategy(nextItem)) { const key = this.createApiTargetKey(this.options.apiTargets[i]) const rules = this.options.apiTargets[i].rules?.length @@ -793,16 +801,25 @@ export class RateLimiterBuilder extends Macroable { } const decision = await this.options.retryStrategy(ctx) + const cooldown = Math.max(0, decision.cooldownMs ?? 0) + + if (cooldown > 0) { + await this.options + .store!.setCooldown(this.options.key, cooldown) + .catch(() => { + debug('failed to set cooldown in cache for key %s', this.options.key) + }) + } switch (decision.type) { - case 'retry_same': case 'retry_other': - const delay = Math.max(0, decision.delayMs ?? 0) - + case 'retry_same': this.releaseTask({ isToScheduleTick: false }) options.item.attempt++ options.item.started = false + options.item.avoidApiTargetId = undefined + options.item.pinnedApiTargetId = undefined if (options.item.signal?.aborted) { options.item.reject(new DOMException('Aborted', 'AbortError')) @@ -812,48 +829,12 @@ export class RateLimiterBuilder extends Macroable { this.queue.unshift(options.item) - if (delay > 0) { - this.nextWakeUpAt = Date.now() + delay - this.scheduleQueueItemRun({ delay }) - } else { - this.scheduleQueueItemRun() - } - - return - - case 'cooldown': { - await this.options.store!.setCooldown( - this.options.key, - Math.max(0, decision.cooldownMs) - ) - - if (decision.then === 'retry_same' || decision.then === 'retry_other') { - const delay = decision.cooldownMs + const delay = cooldown - this.releaseTask({ isToScheduleTick: false }) - - options.item.attempt++ - options.item.started = false - - if (options.item.signal?.aborted) { - options.item.reject(new DOMException('Aborted', 'AbortError')) - - break - } - - this.queue.unshift(options.item) - this.nextWakeUpAt = Date.now() + delay - this.scheduleQueueItemRun({ delay }) - - return - } - - this.releaseTask({ isToScheduleTick: true }) - - options.item.reject(options.error) + this.nextWakeUpAt = Date.now() + cooldown + this.scheduleQueueItemRun({ delay }) - break - } + return default: this.releaseTask({ isToScheduleTick: true }) @@ -890,6 +871,13 @@ export class RateLimiterBuilder extends Macroable { } const decision = await this.options.retryStrategy(ctx) + const cooldown = Math.max(0, decision.cooldownMs ?? 0) + + if (cooldown > 0) { + await this.options.store!.setCooldown(key, cooldown).catch(() => { + debug('failed to set cooldown in cache for key %s', key) + }) + } switch (decision.type) { case 'retry_same': @@ -897,6 +885,7 @@ export class RateLimiterBuilder extends Macroable { options.item.attempt++ options.item.started = false + options.item.avoidApiTargetId = undefined options.item.pinnedApiTargetId = options.apiTarget.id if (options.item.signal?.aborted) { @@ -907,28 +896,19 @@ export class RateLimiterBuilder extends Macroable { this.queue.unshift(options.item) - const delay = Math.max(0, decision.delayMs ?? 0) + const delay = cooldown - if (delay > 0) { - this.nextWakeUpAt = Date.now() + delay - this.scheduleQueueItemRun({ delay }) - } else { - this.scheduleQueueItemRun() - } + this.nextWakeUpAt = Date.now() + delay + this.scheduleQueueItemRun({ delay }) return - case 'retry_other': { - const cooldown = Math.max(0, decision.cooldownMs || 0) - + case 'retry_other': this.releaseTask({ isToScheduleTick: false }) - if (cooldown > 0) { - await this.options.store!.setCooldown(key, cooldown) - } - options.item.attempt++ options.item.started = false + options.item.avoidApiTargetId = options.apiTarget.id options.item.pinnedApiTargetId = undefined if (options.item.signal?.aborted) { @@ -939,54 +919,11 @@ export class RateLimiterBuilder extends Macroable { this.queue.unshift(options.item) - const delay = Math.max(0, decision.delayMs ?? 0) - - if (delay > 0) { - this.nextWakeUpAt = Date.now() + delay - this.scheduleQueueItemRun({ delay }) - } else { - this.scheduleQueueItemRun() - } + this.nextWakeUpAt = Date.now() + this.scheduleQueueItemRun() return - } - - case 'cooldown': - const ms = Math.max(0, decision.cooldownMs || 0) - - this.releaseTask({ isToScheduleTick: false }) - - await this.options.store!.setCooldown(key, ms) - const then = decision.then - - if (then === 'retry_same' || then === 'retry_other') { - options.item.attempt++ - options.item.started = false - options.item.pinnedApiTargetId = - decision.then === 'retry_same' ? options.apiTarget.id : undefined - - if (options.item.signal?.aborted) { - options.item.reject(new DOMException('Aborted', 'AbortError')) - - break - } - - this.queue.unshift(options.item) - - const delay = decision.cooldownMs - - this.nextWakeUpAt = Date.now() + delay - this.scheduleQueueItemRun({ delay }) - - return - } - - this.releaseTask({ isToScheduleTick: true }) - - options.item.reject(options.error) - - break default: this.releaseTask({ isToScheduleTick: true }) diff --git a/src/types/QueueItem.ts b/src/types/QueueItem.ts index 475769d..2bb6489 100644 --- a/src/types/QueueItem.ts +++ b/src/types/QueueItem.ts @@ -17,5 +17,6 @@ export type QueueItem = { started: boolean signal?: AbortSignal attempt?: number + avoidApiTargetId?: string pinnedApiTargetId?: string } diff --git a/src/types/RateLimitRetryDecision.ts b/src/types/RateLimitRetryDecision.ts index c68d38b..c634305 100644 --- a/src/types/RateLimitRetryDecision.ts +++ b/src/types/RateLimitRetryDecision.ts @@ -8,11 +8,53 @@ */ export type RateLimitRetryDecision = - | { type: 'fail' } - | { type: 'retry_same'; delayMs?: number } - | { type: 'retry_other'; delayMs?: number; cooldownMs?: number } | { - type: 'cooldown' - cooldownMs: number - then?: 'fail' | 'retry_same' | 'retry_other' + /** + * Decide that RateLimtiter should fail the request. Use this when + * you have already retried as much as possible and you want RateLimiter + * to throw the exception. + */ + type: 'fail' + + /** + * Define for how long time your API Target will be blocked from usage. + * This is a global state that will be respected by your store when + * defining if it's allowed to run with that API Target or not. + */ + cooldownMs?: number + } + | { + /** + * Decide that your next try should be with the same API Target. + * Returning `retry_same` will basically avoid your RateLimiter + * from using any other API Target until you decide something else. + * + * This decision works when using single API Targets or multiple. + */ + type: 'retry_same' + + /** + * Define for how long time your API Target will be blocked from usage. + * This is a global state that will be respected by your store when + * defining if it's allowed to run with that API Target or not. + */ + cooldownMs?: number + } + | { + /** + * Decide that your next try should be with another API Target. + * Returning `retry_other` will basically avoid your RateLimiter + * from using the last API Target until you decide something else. + * + * This decision only takes effect when using multiple API Targets, + * If using none or only one, it will use `retry_same` by default. + */ + type: 'retry_other' + + /** + * Define for how long time your API Target will be blocked from usage. + * This is a global state that will be respected by your store when + * defining if it's allowed to run with that API Target or not. + */ + cooldownMs?: number } diff --git a/tests/unit/ratelimiter/RateLimiterBuilderTest.ts b/tests/unit/ratelimiter/RateLimiterBuilderTest.ts index e8336e3..ea51057 100644 --- a/tests/unit/ratelimiter/RateLimiterBuilderTest.ts +++ b/tests/unit/ratelimiter/RateLimiterBuilderTest.ts @@ -828,7 +828,7 @@ export class RateLimiterBuilderTest { const limiter = RateLimiter.build() .store('memory', { windowMs: { second: 100 } }) .key('request:api-key:/profile') - .addApiTarget({ rules: [{ type: 'second', limit: 1 }], metadata: { baseUrl: 'http://api0.com' } }) + .addRule({ type: 'second', limit: 1 }) await assert.rejects(() => limiter.schedule(({ apiTarget }) => { @@ -842,7 +842,7 @@ export class RateLimiterBuilderTest { const limiter = RateLimiter.build() .store('memory', { windowMs: { second: 100 } }) .key('request:api-key:/profile') - .addApiTarget({ rules: [{ type: 'second', limit: 1 }], metadata: { baseUrl: 'http://api0.com' } }) + .addRule({ type: 'second', limit: 1 }) .retryStrategy(() => { return { type: 'fail' } }) @@ -855,7 +855,95 @@ export class RateLimiterBuilderTest { } @Test() - public async shouldAlwaysRetryTheRequestWithTheSameApiIfRetryStrategyDecideItShouldRetryWithTheSame({ + public async shouldAlwaysRetryTheRequestIfRetryStrategyDecideItShouldRetry({ assert }: Context) { + const limiter = RateLimiter.build() + .key('request:api-key:/profile') + .store('memory', { windowMs: { second: 100 } }) + .addRule({ type: 'second', limit: 1 }) + .retryStrategy(({ attempt }) => { + if (attempt === 1) { + return { type: 'retry_same' } + } + + return { type: 'fail' } + }) + + let retriedForCount = 0 + let isFirstRequest = true + + await limiter.schedule(() => { + retriedForCount++ + + if (isFirstRequest) { + isFirstRequest = false + + throw new Error('fail') + } + }) + + assert.deepEqual(retriedForCount, 2) + } + + @Test() + public async shouldAlwaysCooldownAndFailTheRequestIfRetryStrategyDecideItShouldCooldownAndFail({ assert }: Context) { + const limiter = RateLimiter.build() + .key('request:api-key:/profile') + .store('memory', { windowMs: { second: 100 } }) + .addRule({ type: 'second', limit: 1 }) + .retryStrategy(({ attempt }) => { + if (attempt === 1) { + return { type: 'retry_same', cooldownMs: 100 } + } + + return { type: 'fail', cooldownMs: 100 } + }) + + let retriedForCount = 0 + + await assert.rejects(() => + limiter.schedule(() => { + retriedForCount++ + throw new Error('fail') + }) + ) + + assert.deepEqual(retriedForCount, 2) + } + + @Test() + public async shouldAlwaysCooldownAndRetryTheRequestIfRetryStrategyDecideItShouldCooldownAndRetry({ + assert + }: Context) { + const limiter = RateLimiter.build() + .key('request:api-key:/profile') + .store('memory', { windowMs: { second: 100 } }) + .addRule({ type: 'second', limit: 1 }) + .retryStrategy(({ attempt }) => { + if (attempt === 1) { + return { type: 'retry_same', cooldownMs: 100 } + } + + return { type: 'fail' } + }) + + let retriedForCount = 0 + let isFirstRequest = true + + await limiter.schedule(() => { + retriedForCount++ + + if (isFirstRequest) { + isFirstRequest = false + + throw new Error('fail') + } + }) + + assert.deepEqual(retriedForCount, 2) + } + + @Test() + public async shouldAlwaysRetryTheRequestWithTheSameApiTargetIfRetryStrategyDecideItShouldRetryWithTheSame({ assert }: Context) { const limiter = RateLimiter.build() @@ -884,7 +972,94 @@ export class RateLimiterBuilderTest { } @Test() - public async shouldAlwaysCooldownAndRetryTheRequestWithTheSameApiIfRetryStrategyDecideItShouldWaitToTryAgainWithSame({ + public async shouldAlwaysRetryTheRequestWithTheOtherApiTargetIfRetryStrategyDecideItShouldRetryWithOther({ + assert + }: Context) { + const limiter = RateLimiter.build() + .key('request:api-key:/profile') + .store('memory', { windowMs: { second: 100 } }) + .addApiTarget({ rules: [{ type: 'second', limit: 1 }], metadata: { baseUrl: 'http://api0.com' } }) + .addApiTarget({ rules: [{ type: 'second', limit: 1 }], metadata: { baseUrl: 'http://api1.com' } }) + .retryStrategy(() => { + return { type: 'retry_other' } + }) + + const apiTargetUsed = [] + let isFirstRequest = true + + await limiter.schedule(({ apiTarget }) => { + apiTargetUsed.push(apiTarget.metadata.baseUrl) + + if (isFirstRequest) { + isFirstRequest = false + + throw new Error('fail') + } + }) + + assert.deepEqual(apiTargetUsed, ['http://api0.com', 'http://api1.com']) + } + + @Test() + public async shouldAlwaysCooldownAndFailTheRequestWithApiTargetIfRetryStrategyDecideItShouldCooldownAnFail({ + assert + }: Context) { + const limiter = RateLimiter.build() + .key('request:api-key:/profile') + .store('memory', { windowMs: { second: 100 } }) + .addApiTarget({ rules: [{ type: 'second', limit: 1 }], metadata: { baseUrl: 'http://api0.com' } }) + .addApiTarget({ rules: [{ type: 'second', limit: 1 }], metadata: { baseUrl: 'http://api1.com' } }) + .retryStrategy(({ attempt }) => { + if (attempt === 1) { + return { type: 'retry_other', cooldownMs: 100 } + } + + return { type: 'fail', cooldownMs: 100 } + }) + + let retriedForCount = 0 + + await assert.rejects(() => + limiter.schedule(() => { + retriedForCount++ + throw new Error('fail') + }) + ) + + assert.deepEqual(retriedForCount, 2) + } + + @Test() + public async shouldAlwaysCooldownAndRetryTheRequestWithOtherApiTargetIfRetryStrategyDecideItShouldCooldownAndTryWithOther({ + assert + }: Context) { + const limiter = RateLimiter.build() + .key('request:api-key:/profile') + .store('memory', { windowMs: { second: 100 } }) + .addApiTarget({ rules: [{ type: 'second', limit: 1 }], metadata: { baseUrl: 'http://api0.com' } }) + .addApiTarget({ rules: [{ type: 'second', limit: 1 }], metadata: { baseUrl: 'http://api1.com' } }) + .retryStrategy(() => { + return { type: 'retry_other', cooldownMs: 100 } + }) + + const apiTargetUsed = [] + let isFirstRequest = true + + await limiter.schedule(({ apiTarget }) => { + apiTargetUsed.push(apiTarget.metadata.baseUrl) + + if (isFirstRequest) { + isFirstRequest = false + + throw new Error('fail') + } + }) + + assert.deepEqual(apiTargetUsed, ['http://api0.com', 'http://api1.com']) + } + + @Test() + public async shouldAlwaysCooldownAndRetryTheRequestWithTheSameApiTargetIfRetryStrategyDecideItShouldCooldownAndTryWithSame({ assert }: Context) { const limiter = RateLimiter.build() @@ -893,7 +1068,7 @@ export class RateLimiterBuilderTest { .addApiTarget({ rules: [{ type: 'second', limit: 1 }], metadata: { baseUrl: 'http://api0.com' } }) .addApiTarget({ rules: [{ type: 'second', limit: 1 }], metadata: { baseUrl: 'http://api1.com' } }) .retryStrategy(() => { - return { type: 'cooldown', cooldownMs: 100, then: 'retry_same' } + return { type: 'retry_same', cooldownMs: 100 } }) const apiTargetUsed = [] From 66c7762775debd1113928017f561c5301123a317 Mon Sep 17 00:00:00 2001 From: jlenon7 Date: Tue, 9 Sep 2025 14:46:09 -0300 Subject: [PATCH 5/6] refactor(targets): rename ApiTarget to Target --- src/exceptions/MissingRuleException.ts | 4 +- src/ratelimiter/RateLimiterBuilder.ts | 128 +++++---- src/types/QueueItem.ts | 4 +- src/types/RateLimitRetryCtx.ts | 6 +- src/types/RateLimitRetryDecision.ts | 24 +- src/types/RateLimitScheduleCtx.ts | 6 +- ...teLimitApiTarget.ts => RateLimitTarget.ts} | 10 +- src/types/RateLimiterOptions.ts | 10 +- src/types/index.ts | 2 +- tests/fixtures/config/cache.ts | 2 +- .../ratelimiter/RateLimiterBuilderTest.ts | 262 +++++++++--------- 11 files changed, 235 insertions(+), 223 deletions(-) rename src/types/{RateLimitApiTarget.ts => RateLimitTarget.ts} (68%) diff --git a/src/exceptions/MissingRuleException.ts b/src/exceptions/MissingRuleException.ts index a1afbea..3b4ce44 100644 --- a/src/exceptions/MissingRuleException.ts +++ b/src/exceptions/MissingRuleException.ts @@ -11,9 +11,9 @@ import { Exception } from '@athenna/common' export class MissingRuleException extends Exception { public constructor() { - const message = 'Missing rules value for rate limiter and API Targets.' + const message = 'Missing rules value for rate limiter and targets.' const help = - 'This error happens when you forget to define default rules for your RateLimiter instance and custom rules by API Target. You has two options, define a default rule in your RateLimiter that will be used by API Targets that does not have a rule or define a custom rule for all your API Targets.' + 'This error happens when you forget to define default rules for your RateLimiter instance and custom rules by target. You has two options, define a default rule in your RateLimiter that will be used by targets that does not have a rule or define a custom rule for all your targets.' super({ code: 'E_MISSING_RULE_ERROR', diff --git a/src/ratelimiter/RateLimiterBuilder.ts b/src/ratelimiter/RateLimiterBuilder.ts index 65bb271..d41333b 100644 --- a/src/ratelimiter/RateLimiterBuilder.ts +++ b/src/ratelimiter/RateLimiterBuilder.ts @@ -11,9 +11,9 @@ import type { QueueItem, RateLimitRule, ScheduleOptions, + RateLimitTarget, RateLimitRetryCtx, RateLimiterOptions, - RateLimitApiTarget, RateLimitScheduleCtx, RateLimitStoreOptions, RateLimitRetryDecision @@ -34,7 +34,7 @@ export class RateLimiterBuilder extends Macroable { private options: RateLimiterOptions = { jitterMs: 0, maxConcurrent: 1, - apiTargetSelectionStrategy: 'first_available' + targetSelectionStrategy: 'first_available' } /** @@ -182,7 +182,7 @@ export class RateLimiterBuilder extends Macroable { } /** - * Add a new rate limit API target. + * Add a new rate limit target. * * @example * ```ts @@ -190,21 +190,21 @@ export class RateLimiterBuilder extends Macroable { * .store('memory') * .key('request:/profile') * .addRule({ type: 'second', limit: 1 }) - * .addApiTarget({ metadata: { baseUrl: 'http://example.com' } }) + * .addTarget({ metadata: { baseUrl: 'http://example.com' } }) * * await limiter.schedule(() => {...}) * ``` */ - public addApiTarget(apiTarget: RateLimitApiTarget) { - if (!this.options.apiTargets) { - this.options.apiTargets = [] + public addTarget(target: RateLimitTarget) { + if (!this.options.targets) { + this.options.targets = [] } - if (!apiTarget.id) { - apiTarget.id = this.getApiTargetId(apiTarget) + if (!target.id) { + target.id = this.getTargetId(target) } - this.options.apiTargets.push(apiTarget) + this.options.targets.push(target) return this } @@ -229,7 +229,7 @@ export class RateLimiterBuilder extends Macroable { } /** - * Set multiple rate limit API targets with one method call. + * Set multiple rate limit targets with one method call. * * @example * ```ts @@ -237,20 +237,20 @@ export class RateLimiterBuilder extends Macroable { * .store('memory') * .key('request:/profile') * .addRule({ type: 'second', limit: 1 }) - * .setApiTargets([{ metadata: { baseUrl: 'http://example.com' } }]) + * .setTargets([{ metadata: { baseUrl: 'http://example.com' } }]) * * await limiter.schedule(() => {...}) * ``` */ - public setApiTargets(apiTargets: RateLimitApiTarget[]) { - apiTargets.forEach(apiTarget => this.addApiTarget(apiTarget)) + public setTargets(targets: RateLimitTarget[]) { + targets.forEach(target => this.addTarget(target)) return this } /** - * Define the API target selection strategy that will be used - * to select the next one when an API Target fails. + * Define the target selection strategy that will be used + * to select the next one when an target fails. * * @example * ```ts @@ -258,20 +258,20 @@ export class RateLimiterBuilder extends Macroable { * .store('memory') * .key('request:/profile') * .addRule({ type: 'second', limit: 1 }) - * .addApiTarget({ metadata: { baseUrl: 'http://example.com' } }) - * .apiTargetSelectionStrategy('round_robin') + * .addTarget({ metadata: { baseUrl: 'http://example.com' } }) + * .targetSelectionStrategy('round_robin') * * await limiter.schedule(() => {...}) * ``` */ - public apiTargetSelectionStrategy(value: 'first_available' | 'round_robin') { - this.options.apiTargetSelectionStrategy = value + public targetSelectionStrategy(value: 'first_available' | 'round_robin') { + this.options.targetSelectionStrategy = value return this } /** - * Define the RateLmiter retry strategy. This is useful to control + * Define the RateLimiter retry strategy. This is useful to control * when and how we should proceed with the retry of tasks that failed * to execute. * @@ -393,7 +393,7 @@ export class RateLimiterBuilder extends Macroable { /** * Schedule the execution of an async function respecting - * the rate limit rules and the API Targets. + * the rate limit rules and the targets. * * @example * ```ts @@ -420,15 +420,15 @@ export class RateLimiterBuilder extends Macroable { } if (!this.options.rules?.length) { - if (!this.options.apiTargets?.length) { + if (!this.options.targets?.length) { throw new MissingRuleException() } - const missingRuleApiTargets = this.options.apiTargets.filter( - apiTarget => !apiTarget.rules + const missingRuleTargets = this.options.targets.filter( + target => !target.rules ) - if (missingRuleApiTargets.length) { + if (missingRuleTargets.length) { throw new MissingRuleException() } } @@ -478,21 +478,21 @@ export class RateLimiterBuilder extends Macroable { } /** - * Create a custom id for an API Target by reading the metadata object. + * Create a custom id for an target by reading the metadata object. * The object will always be sorted by keys. */ - private getApiTargetId(apiTarget: RateLimitApiTarget) { - return String.hash(JSON.stringify(Json.sort(apiTarget.metadata)), { + public getTargetId(target: RateLimitTarget) { + return String.hash(JSON.stringify(Json.sort(target.metadata)), { key: Config.get('app.key', 'ratelimiter') }) } /** - * Create a custom key for an API Target to be used to map the - * API Target rules into the cache. + * Create a custom key for an target to be used to map the + * target rules into the cache. */ - private createApiTargetKey(apiTarget: RateLimitApiTarget) { - return `${this.options.key}:${this.getApiTargetId(apiTarget)}` + public createTargetKey(target: RateLimitTarget) { + return `${this.options.key}:${this.getTargetId(target)}` } /** @@ -508,13 +508,13 @@ export class RateLimiterBuilder extends Macroable { } /** - * Read the API Target selection strategy and defines which is + * Read the target selection strategy and defines which is * going to be used. */ private createIdxBySelectionStrategy(item: QueueItem) { - if (item.pinnedApiTargetId) { - const i = this.options.apiTargets.findIndex( - a => a.id === item.pinnedApiTargetId + if (item.pinnedTargetId) { + const i = this.options.targets.findIndex( + a => a.id === item.pinnedTargetId ) if (i >= 0) { @@ -524,7 +524,7 @@ export class RateLimiterBuilder extends Macroable { let indexes = [] - switch (this.options.apiTargetSelectionStrategy) { + switch (this.options.targetSelectionStrategy) { case 'round_robin': indexes = this.createRoundRobinIdx() break @@ -533,10 +533,8 @@ export class RateLimiterBuilder extends Macroable { indexes = this.createFirstAvailableIdx() } - if (item.avoidApiTargetId) { - const i = this.options.apiTargets.findIndex( - a => a.id === item.avoidApiTargetId - ) + if (item.avoidTargetId) { + const i = this.options.targets.findIndex(a => a.id === item.avoidTargetId) return indexes.filter(idx => idx !== i) } @@ -549,8 +547,8 @@ export class RateLimiterBuilder extends Macroable { */ private createRoundRobinIdx() { return Array.from( - { length: this.options.apiTargets.length }, - (_, k) => (this.rrIndex + k) % this.options.apiTargets.length + { length: this.options.targets.length }, + (_, k) => (this.rrIndex + k) % this.options.targets.length ) } @@ -558,7 +556,7 @@ export class RateLimiterBuilder extends Macroable { * Create the indexes for when using first_available selection strategy. */ private createFirstAvailableIdx() { - return Array.from({ length: this.options.apiTargets.length }, (_, k) => k) + return Array.from({ length: this.options.targets.length }, (_, k) => k) } /** @@ -594,17 +592,17 @@ export class RateLimiterBuilder extends Macroable { const now = Date.now() - if (this.options.apiTargets?.length) { + if (this.options.targets?.length) { let minWait = Number.POSITIVE_INFINITY - let apiTarget: RateLimitApiTarget = null + let target: RateLimitTarget = null const nextItem = this.queue[0] for (const i of this.createIdxBySelectionStrategy(nextItem)) { - const key = this.createApiTargetKey(this.options.apiTargets[i]) + const key = this.createTargetKey(this.options.targets[i]) - const rules = this.options.apiTargets[i].rules?.length - ? this.options.apiTargets[i].rules + const rules = this.options.targets[i].rules?.length + ? this.options.targets[i].rules : this.options.rules try { @@ -613,10 +611,10 @@ export class RateLimiterBuilder extends Macroable { this.storeErrorCount = 0 if (res.allowed) { - apiTarget = this.options.apiTargets[i] + target = this.options.targets[i] - if (this.options.apiTargetSelectionStrategy === 'round_robin') { - this.rrIndex = (i + 1) % this.options.apiTargets.length + if (this.options.targetSelectionStrategy === 'round_robin') { + this.rrIndex = (i + 1) % this.options.targets.length } break @@ -638,7 +636,7 @@ export class RateLimiterBuilder extends Macroable { } } - if (!apiTarget) { + if (!target) { const delay = (isFinite(minWait) ? minWait : 100) + this.randomJitter() this.nextWakeUpAt = now + delay @@ -667,13 +665,13 @@ export class RateLimiterBuilder extends Macroable { this.active++ Promise.resolve() - .then(() => item.run({ signal: item.signal, apiTarget })) + .then(() => item.run({ signal: item.signal, target })) .then(result => { this.releaseTask({ isToScheduleTick: true }) item.resolve(result) }) - .catch(error => this.onFailInMultiMode({ error, item, apiTarget })) + .catch(error => this.onFailInMultiMode({ error, item, target })) if (this.active < this.options.maxConcurrent) { this.scheduleQueueItemRun() @@ -818,8 +816,8 @@ export class RateLimiterBuilder extends Macroable { options.item.attempt++ options.item.started = false - options.item.avoidApiTargetId = undefined - options.item.pinnedApiTargetId = undefined + options.item.avoidTargetId = undefined + options.item.pinnedTargetId = undefined if (options.item.signal?.aborted) { options.item.reject(new DOMException('Aborted', 'AbortError')) @@ -850,7 +848,7 @@ export class RateLimiterBuilder extends Macroable { private onFailInMultiMode = async (options: { error: Error item: QueueItem - apiTarget: RateLimitApiTarget + target: RateLimitTarget }) => { if (!this.options.retryStrategy) { this.releaseTask({ isToScheduleTick: true }) @@ -860,14 +858,14 @@ export class RateLimiterBuilder extends Macroable { return } - const key = this.createApiTargetKey(options.apiTarget) + const key = this.createTargetKey(options.target) const ctx: RateLimitRetryCtx = { key, error: options.error, signal: options.item.signal, attempt: options.item.attempt, - apiTarget: options.apiTarget + target: options.target } const decision = await this.options.retryStrategy(ctx) @@ -885,8 +883,8 @@ export class RateLimiterBuilder extends Macroable { options.item.attempt++ options.item.started = false - options.item.avoidApiTargetId = undefined - options.item.pinnedApiTargetId = options.apiTarget.id + options.item.avoidTargetId = undefined + options.item.pinnedTargetId = options.target.id if (options.item.signal?.aborted) { options.item.reject(new DOMException('Aborted', 'AbortError')) @@ -908,8 +906,8 @@ export class RateLimiterBuilder extends Macroable { options.item.attempt++ options.item.started = false - options.item.avoidApiTargetId = options.apiTarget.id - options.item.pinnedApiTargetId = undefined + options.item.avoidTargetId = options.target.id + options.item.pinnedTargetId = undefined if (options.item.signal?.aborted) { options.item.reject(new DOMException('Aborted', 'AbortError')) diff --git a/src/types/QueueItem.ts b/src/types/QueueItem.ts index 2bb6489..e8bf1d9 100644 --- a/src/types/QueueItem.ts +++ b/src/types/QueueItem.ts @@ -17,6 +17,6 @@ export type QueueItem = { started: boolean signal?: AbortSignal attempt?: number - avoidApiTargetId?: string - pinnedApiTargetId?: string + avoidTargetId?: string + pinnedTargetId?: string } diff --git a/src/types/RateLimitRetryCtx.ts b/src/types/RateLimitRetryCtx.ts index ec9311a..7e091ec 100644 --- a/src/types/RateLimitRetryCtx.ts +++ b/src/types/RateLimitRetryCtx.ts @@ -7,7 +7,7 @@ * file that was distributed with this source code. */ -import type { RateLimitApiTarget } from '#src/types' +import type { RateLimitTarget } from '#src/types' export type RateLimitRetryCtx = { /** @@ -31,7 +31,7 @@ export type RateLimitRetryCtx = { attempt: number /** - * The API Target that this retry is currently using. + * The target that this retry is currently using. */ - apiTarget?: RateLimitApiTarget + target?: RateLimitTarget } diff --git a/src/types/RateLimitRetryDecision.ts b/src/types/RateLimitRetryDecision.ts index c634305..a527642 100644 --- a/src/types/RateLimitRetryDecision.ts +++ b/src/types/RateLimitRetryDecision.ts @@ -17,44 +17,44 @@ export type RateLimitRetryDecision = type: 'fail' /** - * Define for how long time your API Target will be blocked from usage. + * Define for how long time your target will be blocked from usage. * This is a global state that will be respected by your store when - * defining if it's allowed to run with that API Target or not. + * defining if it's allowed to run with that target or not. */ cooldownMs?: number } | { /** - * Decide that your next try should be with the same API Target. + * Decide that your next try should be with the same target. * Returning `retry_same` will basically avoid your RateLimiter - * from using any other API Target until you decide something else. + * from using any other target until you decide something else. * - * This decision works when using single API Targets or multiple. + * This decision works when using single targets or multiple. */ type: 'retry_same' /** - * Define for how long time your API Target will be blocked from usage. + * Define for how long time your target will be blocked from usage. * This is a global state that will be respected by your store when - * defining if it's allowed to run with that API Target or not. + * defining if it's allowed to run with that target or not. */ cooldownMs?: number } | { /** - * Decide that your next try should be with another API Target. + * Decide that your next try should be with another target. * Returning `retry_other` will basically avoid your RateLimiter - * from using the last API Target until you decide something else. + * from using the last target until you decide something else. * - * This decision only takes effect when using multiple API Targets, + * This decision only takes effect when using multiple targets, * If using none or only one, it will use `retry_same` by default. */ type: 'retry_other' /** - * Define for how long time your API Target will be blocked from usage. + * Define for how long time your target will be blocked from usage. * This is a global state that will be respected by your store when - * defining if it's allowed to run with that API Target or not. + * defining if it's allowed to run with that target or not. */ cooldownMs?: number } diff --git a/src/types/RateLimitScheduleCtx.ts b/src/types/RateLimitScheduleCtx.ts index dc3d51f..8af8ef9 100644 --- a/src/types/RateLimitScheduleCtx.ts +++ b/src/types/RateLimitScheduleCtx.ts @@ -7,7 +7,7 @@ * file that was distributed with this source code. */ -import type { RateLimitApiTarget } from '#src/types' +import type { RateLimitTarget } from '#src/types' export type RateLimitScheduleCtx = { /** @@ -16,7 +16,7 @@ export type RateLimitScheduleCtx = { signal?: AbortSignal /** - * The API Target that this retry is currently using. + * The target that this retry is currently using. */ - apiTarget?: RateLimitApiTarget + target?: RateLimitTarget } diff --git a/src/types/RateLimitApiTarget.ts b/src/types/RateLimitTarget.ts similarity index 68% rename from src/types/RateLimitApiTarget.ts rename to src/types/RateLimitTarget.ts index 8a821bc..a540775 100644 --- a/src/types/RateLimitApiTarget.ts +++ b/src/types/RateLimitTarget.ts @@ -9,25 +9,25 @@ import type { RateLimitRule } from '#src/types' -export type RateLimitApiTarget = { +export type RateLimitTarget = { /** * The rate limit target ID. By default this will be created by creating - * a hash from the API Target metadata object, but you can also define your + * a hash from the target metadata object, but you can also define your * own ID. */ id?: string /** - * Define all the metadata for this API target to function. Metadata + * Define all the metadata for this target to function. Metadata * is required because we are going to create a hash from this object - * to store the rules inside the cache by ApiTarget. With this + * to store the rules inside the cache by Target. With this * implementation you can create not only API rotations but also API * Keys rotations at the same time. */ metadata: Record /** - * Custom rate limit rules for this API target. If not defined, + * Custom rate limit rules for this target. If not defined, * the default defined in RateLimiter will be used. */ rules?: RateLimitRule[] diff --git a/src/types/RateLimiterOptions.ts b/src/types/RateLimiterOptions.ts index 11690ca..239f5ea 100644 --- a/src/types/RateLimiterOptions.ts +++ b/src/types/RateLimiterOptions.ts @@ -9,7 +9,7 @@ import type { RateLimitRule, - RateLimitApiTarget, + RateLimitTarget, RateLimitRetryClosure } from '#src/types' import type { RateLimitStore } from '#src/ratelimiter/RateLimitStore' @@ -23,7 +23,7 @@ export type RateLimiterOptions = { /** * The logical key that will be used by store to save buckets. * If API Targers are defined, it will be used as a prefix from - * a hash created from API Targets metadata object: `${key}:${hash}`. + * a hash created from targets metadata object: `${key}:${hash}`. */ key?: string @@ -31,7 +31,7 @@ export type RateLimiterOptions = { * The api targets that will be used to create API rotations when * some of them fails. */ - apiTargets?: RateLimitApiTarget[] + targets?: RateLimitTarget[] /** * The retry strategy for this rate limiter. This is useful to @@ -60,10 +60,10 @@ export type RateLimiterOptions = { jitterMs?: number /** - * Define the selection strategy that will be used to select which API target + * Define the selection strategy that will be used to select which target * will be used next when some of them fails. * * @default 'first_available' */ - apiTargetSelectionStrategy: 'first_available' | 'round_robin' + targetSelectionStrategy: 'first_available' | 'round_robin' } diff --git a/src/types/index.ts b/src/types/index.ts index 221879f..cbd312a 100644 --- a/src/types/index.ts +++ b/src/types/index.ts @@ -11,8 +11,8 @@ export * from '#src/types/Reserve' export * from '#src/types/QueueItem' export * from '#src/types/RateLimitRule' export * from '#src/types/ScheduleOptions' +export * from '#src/types/RateLimitTarget' export * from '#src/types/RateLimitRetryCtx' -export * from '#src/types/RateLimitApiTarget' export * from '#src/types/RateLimiterOptions' export * from '#src/types/RateLimitScheduleCtx' export * from '#src/types/RateLimitRetryClosure' diff --git a/tests/fixtures/config/cache.ts b/tests/fixtures/config/cache.ts index 00a2c5a..e260491 100644 --- a/tests/fixtures/config/cache.ts +++ b/tests/fixtures/config/cache.ts @@ -1,5 +1,5 @@ /** - * @athenna/cache + * @athenna/ratelimiter * * (c) João Lenon * diff --git a/tests/unit/ratelimiter/RateLimiterBuilderTest.ts b/tests/unit/ratelimiter/RateLimiterBuilderTest.ts index ea51057..52ef855 100644 --- a/tests/unit/ratelimiter/RateLimiterBuilderTest.ts +++ b/tests/unit/ratelimiter/RateLimiterBuilderTest.ts @@ -7,7 +7,7 @@ * file that was distributed with this source code. */ -import { RateLimiter } from '#src' +import { RateLimiter, RateLimitStore } from '#src' import { Path, Sleep } from '@athenna/common' import { Cache, CacheProvider } from '@athenna/cache' import { AfterEach, BeforeEach, Test, type Context } from '@athenna/test' @@ -565,14 +565,14 @@ export class RateLimiterBuilderTest { } @Test() - public async shouldBeAbleToBuildARateLimiterWithARuleOfOneRequestPerSecondWithAnApiTarget({ assert }: Context) { + public async shouldBeAbleToBuildARateLimiterWithARuleOfOneRequestPerSecondWithAnTarget({ assert }: Context) { assert.plan(6) const limiter = RateLimiter.build() .store('memory', { windowMs: { second: 100 } }) .key('request:api-key:/profile') .addRule({ type: 'second', limit: 1 }) - .addApiTarget({ metadata: { baseUrl: 'http://api1.com' } }) + .addTarget({ metadata: { baseUrl: 'http://api1.com' } }) const promises = [] const dateStart = Date.now() @@ -580,8 +580,8 @@ export class RateLimiterBuilderTest { for (let i = 0; i < numberOfRequests; i++) { promises.push( - limiter.schedule(({ apiTarget }) => { - assert.isDefined(apiTarget) + limiter.schedule(({ target }) => { + assert.isDefined(target) return 'ok' + i }) @@ -594,7 +594,7 @@ export class RateLimiterBuilderTest { } @Test() - public async shouldBeAbleToBuildARateLimiterWithARuleOfOneRequestPerSecondSettingMultipleApiTargets({ + public async shouldBeAbleToBuildARateLimiterWithARuleOfOneRequestPerSecondSettingMultipleTargets({ assert }: Context) { assert.plan(6) @@ -603,7 +603,7 @@ export class RateLimiterBuilderTest { .store('memory', { windowMs: { second: 100 } }) .key('request:api-key:/profile') .addRule({ type: 'second', limit: 1 }) - .setApiTargets([{ metadata: { baseUrl: 'http://api1.com' } }]) + .setTargets([{ metadata: { baseUrl: 'http://api1.com' } }]) const promises = [] const dateStart = Date.now() @@ -611,8 +611,8 @@ export class RateLimiterBuilderTest { for (let i = 0; i < numberOfRequests; i++) { promises.push( - limiter.schedule(({ apiTarget }) => { - assert.isDefined(apiTarget) + limiter.schedule(({ target }) => { + assert.isDefined(target) return 'ok' + i }) @@ -625,12 +625,12 @@ export class RateLimiterBuilderTest { } @Test() - public async shouldBeAbleToHaveErrorsHappeningInsideTheRateLimiterHandlerEvenWithAnApiTargetSet({ assert }: Context) { + public async shouldBeAbleToHaveErrorsHappeningInsideTheRateLimiterHandlerEvenWithAnTargetSet({ assert }: Context) { const limiter = RateLimiter.build() .store('memory', { windowMs: { second: 100 } }) .key('request:api-key:/profile') .addRule({ type: 'second', limit: 1 }) - .addApiTarget({ metadata: { baseUrl: 'http://api1.com' } }) + .addTarget({ metadata: { baseUrl: 'http://api1.com' } }) await assert.rejects(() => { return limiter.schedule(() => { @@ -640,52 +640,68 @@ export class RateLimiterBuilderTest { } @Test() - public async shouldBeAbleToTryWithTheSecondApiTargetIfTheFirstApiTargetIsAtFullCapacityInASequentialScenario({ + public async shouldBeAbleToTryWithTheSecondTargetIfTheFirstTargetIsAtFullCapacityInASequentialScenario({ assert }: Context) { + const api1 = { metadata: { baseUrl: 'http://api1.com' } } + const api2 = { metadata: { baseUrl: 'http://api2.com' } } + const limiter = RateLimiter.build() .store('memory', { windowMs: { second: 100 } }) .key('request:api-key:/profile') .addRule({ type: 'second', limit: 1 }) - .addApiTarget({ metadata: { baseUrl: 'http://api1.com' } }) - .addApiTarget({ metadata: { baseUrl: 'http://api2.com' } }) + .addTarget(api1) + .addTarget(api2) - const first = await limiter.schedule(({ apiTarget }) => apiTarget.metadata.baseUrl) + const store = new RateLimitStore({ store: 'memory', windowMs: { second: 100 } }) - assert.equal(first, 'http://api1.com') + await store.setCooldown(limiter.createTargetKey(api1), 1000) - const second = await limiter.schedule(({ apiTarget }) => apiTarget.metadata.baseUrl) + const result = await limiter.schedule(({ target }) => { + return target.metadata.baseUrl + }) - assert.equal(second, 'http://api2.com') + assert.deepEqual(result, 'http://api2.com') } @Test() - public async shouldBeAbleToTryWithTheSecondApiTargetIfTheFirstApiTargetIsAtFullCapacityInAConcurrentScenario({ + public async shouldBeAbleToTryWithTheSecondTargetIfTheFirstTargetIsAtFullCapacityInAConcurrentScenario({ assert }: Context) { + const api1 = { metadata: { baseUrl: 'http://api1.com' } } + const api2 = { metadata: { baseUrl: 'http://api2.com' } } + const limiter = RateLimiter.build() .maxConcurrent(2) .store('memory', { windowMs: { second: 100 } }) .key('request:api-key:/profile') .addRule({ type: 'second', limit: 1 }) - .addApiTarget({ metadata: { baseUrl: 'http://api1.com' } }) - .addApiTarget({ metadata: { baseUrl: 'http://api2.com' } }) + .addTarget(api1) + .addTarget(api2) + + const store = new RateLimitStore({ store: 'memory', windowMs: { second: 100 } }) + + await store.setCooldown(limiter.createTargetKey(api1), 1000) - const barrier = this.createBarrier() const used: string[] = [] + const barrier = this.createBarrier() - const run = async ({ apiTarget }) => { - used.push(apiTarget.metadata.baseUrl) + const run = async ({ target }) => { + used.push(target.metadata.baseUrl) await barrier.wait() - return apiTarget.metadata.baseUrl + return target.metadata.baseUrl } const p1 = limiter.schedule(run) const p2 = limiter.schedule(run) - await Sleep.for(5).milliseconds().wait() + await Sleep.for(105).milliseconds().wait() + + used.sort() + + assert.deepEqual(used, ['http://api2.com', 'http://api2.com']) barrier.release() @@ -693,116 +709,114 @@ export class RateLimiterBuilderTest { results.sort() - assert.deepEqual(results, ['http://api1.com', 'http://api2.com']) - - used.sort() - assert.deepEqual(used, ['http://api1.com', 'http://api2.com']) + assert.deepEqual(results, ['http://api2.com', 'http://api2.com']) } @Test() - public async shouldBeAbleToTryWithTheSecondApiTargetIfTheFirstApiTargetIsAtFullCapacityInASequentialScenarioWithRoundRobinStrategy({ + public async shouldBeAbleToTryWithTheSecondTargetIfTheFirstTargetIsAtFullCapacityInASequentialScenarioWithRoundRobinStrategy({ assert }: Context) { + const api1 = { id: 'api1', metadata: { baseUrl: 'http://api1.com' } } + const api2 = { id: 'api2', metadata: { baseUrl: 'http://api2.com' } } + const limiter = RateLimiter.build() .store('memory', { windowMs: { second: 100 } }) .key('request:api-key:/profile') - .apiTargetSelectionStrategy('round_robin') + .targetSelectionStrategy('round_robin') .addRule({ type: 'second', limit: 1 }) - .addApiTarget({ id: 't1', metadata: { baseUrl: 'http://api1.com' } }) - .addApiTarget({ id: 't2', metadata: { baseUrl: 'http://api2.com' } }) - - const used: string[] = [] - - const tasks = Array.from({ length: 4 }, () => - limiter.schedule(({ apiTarget }) => { - used.push(apiTarget.metadata.baseUrl) + .addTarget(api1) + .addTarget(api2) - return apiTarget.metadata.baseUrl - }) - ) + const store = new RateLimitStore({ store: 'memory', windowMs: { second: 100 } }) - const results = await Promise.all(tasks) + await store.setCooldown(limiter.createTargetKey(api1), 1000) - assert.deepEqual(results, used) - assert.deepEqual(used, ['http://api1.com', 'http://api2.com', 'http://api1.com', 'http://api2.com']) - } - - @Test() - public async shouldThrowMissingRuleExceptionIfRateLimiterRulesAndApiTargetRulesAreNotDefined({ assert }: Context) { - const limiter = RateLimiter.build() - .store('memory', { windowMs: { second: 100 } }) - .key('request:api-key:/profile') - .addApiTarget({ id: 't1', metadata: { baseUrl: 'http://api1.com' } }) + const result = await limiter.schedule(({ target }) => { + return target.metadata.baseUrl + }) - await assert.rejects(() => limiter.schedule(() => {}), MissingRuleException) + assert.deepEqual(result, 'http://api2.com') } @Test() - public async shouldNotThrowMissingRuleExceptionIfRateLimiterRulesAreNotDefinedButApiTargetRulesAreDefined({ + public async shouldBeAbleToTryWithTheSecondTargetIfTheFirstTargetIsAtFullCapacityInAConcurrentScenarioWithRoundRobinStrategy({ assert }: Context) { - const limiter = RateLimiter.build() - .store('memory', { windowMs: { second: 100 } }) - .key('request:api-key:/profile') - .addApiTarget({ id: 't1', rules: [{ type: 'second', limit: 1 }], metadata: { baseUrl: 'http://api1.com' } }) + const api1 = { id: 'api1', metadata: { baseUrl: 'http://api1.com' } } + const api2 = { id: 'api2', metadata: { baseUrl: 'http://api2.com' } } - await assert.doesNotReject(() => limiter.schedule(() => {}), MissingRuleException) - } - - @Test() - public async shouldBeAbleToTryWithTheSecondApiTargetIfTheFirstApiTargetIsAtFullCapacityInAConcurrentScenarioWithRoundRobinStrategy({ - assert - }: Context) { const limiter = RateLimiter.build() .maxConcurrent(2) .store('memory', { windowMs: { second: 100 } }) .key('request:api-key:/profile') - .apiTargetSelectionStrategy('round_robin') + .targetSelectionStrategy('round_robin') .addRule({ type: 'second', limit: 1 }) - .addApiTarget({ id: 't1', metadata: { baseUrl: 'http://api1.com' } }) - .addApiTarget({ id: 't2', metadata: { baseUrl: 'http://api2.com' } }) + .addTarget(api1) + .addTarget(api2) + + const store = new RateLimitStore({ store: 'memory', windowMs: { second: 100 } }) + await store.setCooldown(limiter.createTargetKey(api1), 1000) + + const used: string[] = [] const barrier = this.createBarrier() - const started: string[] = [] - const p1 = limiter.schedule(async ({ apiTarget }) => { - started.push(apiTarget.metadata.baseUrl) + const run = async ({ target }) => { + used.push(target.metadata.baseUrl) await barrier.wait() - return apiTarget.metadata.baseUrl - }) - const p2 = limiter.schedule(async ({ apiTarget }) => { - started.push(apiTarget.metadata.baseUrl) + return target.metadata.baseUrl + } - await barrier.wait() + const p1 = limiter.schedule(run) + const p2 = limiter.schedule(run) - return apiTarget.metadata.baseUrl - }) + await Sleep.for(105).milliseconds().wait() - for (let i = 0; i < 20 && started.length < 2; i++) { - await Sleep.for(1).milliseconds().wait() - } + used.sort() - assert.deepEqual(started, ['http://api1.com', 'http://api2.com']) + assert.deepEqual(used, ['http://api2.com', 'http://api2.com']) barrier.release() - const result = await Promise.all([p1, p2]) + const results = await Promise.all([p1, p2]) + + results.sort() + + assert.deepEqual(results, ['http://api2.com', 'http://api2.com']) + } - result.sort() + @Test() + public async shouldThrowMissingRuleExceptionIfRateLimiterRulesAndTargetRulesAreNotDefined({ assert }: Context) { + const limiter = RateLimiter.build() + .store('memory', { windowMs: { second: 100 } }) + .key('request:api-key:/profile') + .addTarget({ id: 't1', metadata: { baseUrl: 'http://api1.com' } }) + + await assert.rejects(() => limiter.schedule(() => {}), MissingRuleException) + } - assert.deepEqual(result, ['http://api1.com', 'http://api2.com']) + @Test() + public async shouldNotThrowMissingRuleExceptionIfRateLimiterRulesAreNotDefinedButTargetRulesAreDefined({ + assert + }: Context) { + const limiter = RateLimiter.build() + .store('memory', { windowMs: { second: 100 } }) + .key('request:api-key:/profile') + .addTarget({ id: 't1', rules: [{ type: 'second', limit: 1 }], metadata: { baseUrl: 'http://api1.com' } }) + + await assert.doesNotReject(() => limiter.schedule(() => {}), MissingRuleException) } @Test() - public async shouldBeAbleToBuildARateLimiterDefiningRulesInApiTarget({ assert }: Context) { + public async shouldBeAbleToBuildARateLimiterDefiningRulesInTarget({ assert }: Context) { assert.plan(6) const limiter = RateLimiter.build() .store('memory', { windowMs: { second: 100 } }) .key('request:api-key:/profile') - .addApiTarget({ rules: [{ type: 'second', limit: 1 }], metadata: { baseUrl: 'http://api0.com' } }) + .addTarget({ rules: [{ type: 'second', limit: 1 }], metadata: { baseUrl: 'http://api0.com' } }) const promises = [] const dateStart = Date.now() @@ -810,8 +824,8 @@ export class RateLimiterBuilderTest { for (let i = 0; i < numberOfRequests; i++) { promises.push( - limiter.schedule(({ apiTarget }) => { - assert.isDefined(apiTarget) + limiter.schedule(({ target }) => { + assert.isDefined(target) return 'ok' + i }) @@ -831,8 +845,8 @@ export class RateLimiterBuilderTest { .addRule({ type: 'second', limit: 1 }) await assert.rejects(() => - limiter.schedule(({ apiTarget }) => { - throw new Error(apiTarget.metadata.baseUrl) + limiter.schedule(({ target }) => { + throw new Error(target.metadata.baseUrl) }) ) } @@ -943,23 +957,23 @@ export class RateLimiterBuilderTest { } @Test() - public async shouldAlwaysRetryTheRequestWithTheSameApiTargetIfRetryStrategyDecideItShouldRetryWithTheSame({ + public async shouldAlwaysRetryTheRequestWithTheSameTargetIfRetryStrategyDecideItShouldRetryWithTheSame({ assert }: Context) { const limiter = RateLimiter.build() .key('request:api-key:/profile') .store('memory', { windowMs: { second: 100 } }) - .addApiTarget({ rules: [{ type: 'second', limit: 1 }], metadata: { baseUrl: 'http://api0.com' } }) - .addApiTarget({ rules: [{ type: 'second', limit: 1 }], metadata: { baseUrl: 'http://api1.com' } }) + .addTarget({ rules: [{ type: 'second', limit: 1 }], metadata: { baseUrl: 'http://api0.com' } }) + .addTarget({ rules: [{ type: 'second', limit: 1 }], metadata: { baseUrl: 'http://api1.com' } }) .retryStrategy(() => { return { type: 'retry_same' } }) - const apiTargetUsed = [] + const targetUsed = [] let isFirstRequest = true - await limiter.schedule(({ apiTarget }) => { - apiTargetUsed.push(apiTarget.metadata.baseUrl) + await limiter.schedule(({ target }) => { + targetUsed.push(target.metadata.baseUrl) if (isFirstRequest) { isFirstRequest = false @@ -968,27 +982,27 @@ export class RateLimiterBuilderTest { } }) - assert.deepEqual(apiTargetUsed, ['http://api0.com', 'http://api0.com']) + assert.deepEqual(targetUsed, ['http://api0.com', 'http://api0.com']) } @Test() - public async shouldAlwaysRetryTheRequestWithTheOtherApiTargetIfRetryStrategyDecideItShouldRetryWithOther({ + public async shouldAlwaysRetryTheRequestWithTheOtherTargetIfRetryStrategyDecideItShouldRetryWithOther({ assert }: Context) { const limiter = RateLimiter.build() .key('request:api-key:/profile') .store('memory', { windowMs: { second: 100 } }) - .addApiTarget({ rules: [{ type: 'second', limit: 1 }], metadata: { baseUrl: 'http://api0.com' } }) - .addApiTarget({ rules: [{ type: 'second', limit: 1 }], metadata: { baseUrl: 'http://api1.com' } }) + .addTarget({ rules: [{ type: 'second', limit: 1 }], metadata: { baseUrl: 'http://api0.com' } }) + .addTarget({ rules: [{ type: 'second', limit: 1 }], metadata: { baseUrl: 'http://api1.com' } }) .retryStrategy(() => { return { type: 'retry_other' } }) - const apiTargetUsed = [] + const targetUsed = [] let isFirstRequest = true - await limiter.schedule(({ apiTarget }) => { - apiTargetUsed.push(apiTarget.metadata.baseUrl) + await limiter.schedule(({ target }) => { + targetUsed.push(target.metadata.baseUrl) if (isFirstRequest) { isFirstRequest = false @@ -997,18 +1011,18 @@ export class RateLimiterBuilderTest { } }) - assert.deepEqual(apiTargetUsed, ['http://api0.com', 'http://api1.com']) + assert.deepEqual(targetUsed, ['http://api0.com', 'http://api1.com']) } @Test() - public async shouldAlwaysCooldownAndFailTheRequestWithApiTargetIfRetryStrategyDecideItShouldCooldownAnFail({ + public async shouldAlwaysCooldownAndFailTheRequestWithTargetIfRetryStrategyDecideItShouldCooldownAnFail({ assert }: Context) { const limiter = RateLimiter.build() .key('request:api-key:/profile') .store('memory', { windowMs: { second: 100 } }) - .addApiTarget({ rules: [{ type: 'second', limit: 1 }], metadata: { baseUrl: 'http://api0.com' } }) - .addApiTarget({ rules: [{ type: 'second', limit: 1 }], metadata: { baseUrl: 'http://api1.com' } }) + .addTarget({ rules: [{ type: 'second', limit: 1 }], metadata: { baseUrl: 'http://api0.com' } }) + .addTarget({ rules: [{ type: 'second', limit: 1 }], metadata: { baseUrl: 'http://api1.com' } }) .retryStrategy(({ attempt }) => { if (attempt === 1) { return { type: 'retry_other', cooldownMs: 100 } @@ -1030,23 +1044,23 @@ export class RateLimiterBuilderTest { } @Test() - public async shouldAlwaysCooldownAndRetryTheRequestWithOtherApiTargetIfRetryStrategyDecideItShouldCooldownAndTryWithOther({ + public async shouldAlwaysCooldownAndRetryTheRequestWithOtherTargetIfRetryStrategyDecideItShouldCooldownAndTryWithOther({ assert }: Context) { const limiter = RateLimiter.build() .key('request:api-key:/profile') .store('memory', { windowMs: { second: 100 } }) - .addApiTarget({ rules: [{ type: 'second', limit: 1 }], metadata: { baseUrl: 'http://api0.com' } }) - .addApiTarget({ rules: [{ type: 'second', limit: 1 }], metadata: { baseUrl: 'http://api1.com' } }) + .addTarget({ rules: [{ type: 'second', limit: 1 }], metadata: { baseUrl: 'http://api0.com' } }) + .addTarget({ rules: [{ type: 'second', limit: 1 }], metadata: { baseUrl: 'http://api1.com' } }) .retryStrategy(() => { return { type: 'retry_other', cooldownMs: 100 } }) - const apiTargetUsed = [] + const targetUsed = [] let isFirstRequest = true - await limiter.schedule(({ apiTarget }) => { - apiTargetUsed.push(apiTarget.metadata.baseUrl) + await limiter.schedule(({ target }) => { + targetUsed.push(target.metadata.baseUrl) if (isFirstRequest) { isFirstRequest = false @@ -1055,27 +1069,27 @@ export class RateLimiterBuilderTest { } }) - assert.deepEqual(apiTargetUsed, ['http://api0.com', 'http://api1.com']) + assert.deepEqual(targetUsed, ['http://api0.com', 'http://api1.com']) } @Test() - public async shouldAlwaysCooldownAndRetryTheRequestWithTheSameApiTargetIfRetryStrategyDecideItShouldCooldownAndTryWithSame({ + public async shouldAlwaysCooldownAndRetryTheRequestWithTheSameTargetIfRetryStrategyDecideItShouldCooldownAndTryWithSame({ assert }: Context) { const limiter = RateLimiter.build() .key('request:api-key:/profile') .store('memory', { windowMs: { second: 100 } }) - .addApiTarget({ rules: [{ type: 'second', limit: 1 }], metadata: { baseUrl: 'http://api0.com' } }) - .addApiTarget({ rules: [{ type: 'second', limit: 1 }], metadata: { baseUrl: 'http://api1.com' } }) + .addTarget({ rules: [{ type: 'second', limit: 1 }], metadata: { baseUrl: 'http://api0.com' } }) + .addTarget({ rules: [{ type: 'second', limit: 1 }], metadata: { baseUrl: 'http://api1.com' } }) .retryStrategy(() => { return { type: 'retry_same', cooldownMs: 100 } }) - const apiTargetUsed = [] + const targetUsed = [] let isFirstRequest = true - await limiter.schedule(({ apiTarget }) => { - apiTargetUsed.push(apiTarget.metadata.baseUrl) + await limiter.schedule(({ target }) => { + targetUsed.push(target.metadata.baseUrl) if (isFirstRequest) { isFirstRequest = false @@ -1084,6 +1098,6 @@ export class RateLimiterBuilderTest { } }) - assert.deepEqual(apiTargetUsed, ['http://api0.com', 'http://api0.com']) + assert.deepEqual(targetUsed, ['http://api0.com', 'http://api0.com']) } } From c486cae7c86f41d3e93f7d78fd595f815a0b3d9e Mon Sep 17 00:00:00 2001 From: jlenon7 Date: Tue, 9 Sep 2025 15:02:47 -0300 Subject: [PATCH 6/6] refactor(targets): rename cooldownMs to currentTargetCooldownMs --- src/ratelimiter/RateLimiterBuilder.ts | 4 ++-- src/types/RateLimitRetryDecision.ts | 6 ++--- .../ratelimiter/RateLimiterBuilderTest.ts | 24 +++++++++---------- 3 files changed, 17 insertions(+), 17 deletions(-) diff --git a/src/ratelimiter/RateLimiterBuilder.ts b/src/ratelimiter/RateLimiterBuilder.ts index d41333b..d791b86 100644 --- a/src/ratelimiter/RateLimiterBuilder.ts +++ b/src/ratelimiter/RateLimiterBuilder.ts @@ -799,7 +799,7 @@ export class RateLimiterBuilder extends Macroable { } const decision = await this.options.retryStrategy(ctx) - const cooldown = Math.max(0, decision.cooldownMs ?? 0) + const cooldown = Math.max(0, decision.currentTargetCooldownMs ?? 0) if (cooldown > 0) { await this.options @@ -869,7 +869,7 @@ export class RateLimiterBuilder extends Macroable { } const decision = await this.options.retryStrategy(ctx) - const cooldown = Math.max(0, decision.cooldownMs ?? 0) + const cooldown = Math.max(0, decision.currentTargetCooldownMs ?? 0) if (cooldown > 0) { await this.options.store!.setCooldown(key, cooldown).catch(() => { diff --git a/src/types/RateLimitRetryDecision.ts b/src/types/RateLimitRetryDecision.ts index a527642..8e21a7d 100644 --- a/src/types/RateLimitRetryDecision.ts +++ b/src/types/RateLimitRetryDecision.ts @@ -21,7 +21,7 @@ export type RateLimitRetryDecision = * This is a global state that will be respected by your store when * defining if it's allowed to run with that target or not. */ - cooldownMs?: number + currentTargetCooldownMs?: number } | { /** @@ -38,7 +38,7 @@ export type RateLimitRetryDecision = * This is a global state that will be respected by your store when * defining if it's allowed to run with that target or not. */ - cooldownMs?: number + currentTargetCooldownMs?: number } | { /** @@ -56,5 +56,5 @@ export type RateLimitRetryDecision = * This is a global state that will be respected by your store when * defining if it's allowed to run with that target or not. */ - cooldownMs?: number + currentTargetCooldownMs?: number } diff --git a/tests/unit/ratelimiter/RateLimiterBuilderTest.ts b/tests/unit/ratelimiter/RateLimiterBuilderTest.ts index 52ef855..803bc4f 100644 --- a/tests/unit/ratelimiter/RateLimiterBuilderTest.ts +++ b/tests/unit/ratelimiter/RateLimiterBuilderTest.ts @@ -289,11 +289,11 @@ export class RateLimiterBuilderTest { /** * Third one only after releasing the 300ms window (rule minute) */ - assert.isAtLeast(starts[2] - starts[1], 300) - assert.isAtLeast(starts[3] - starts[2], 200) - assert.isAtLeast(starts[4] - starts[3], 300) - assert.isAtLeast(starts[5] - starts[4], 200) - assert.isAtLeast(total, 800) + assert.isAtLeast(starts[2] - starts[1], 299) + assert.isAtLeast(starts[3] - starts[2], 199) + assert.isAtLeast(starts[4] - starts[3], 299) + assert.isAtLeast(starts[5] - starts[4], 199) + assert.isAtLeast(total, 799) } @Test() @@ -906,10 +906,10 @@ export class RateLimiterBuilderTest { .addRule({ type: 'second', limit: 1 }) .retryStrategy(({ attempt }) => { if (attempt === 1) { - return { type: 'retry_same', cooldownMs: 100 } + return { type: 'retry_same', currentTargetCooldownMs: 100 } } - return { type: 'fail', cooldownMs: 100 } + return { type: 'fail', currentTargetCooldownMs: 100 } }) let retriedForCount = 0 @@ -934,7 +934,7 @@ export class RateLimiterBuilderTest { .addRule({ type: 'second', limit: 1 }) .retryStrategy(({ attempt }) => { if (attempt === 1) { - return { type: 'retry_same', cooldownMs: 100 } + return { type: 'retry_same', currentTargetCooldownMs: 100 } } return { type: 'fail' } @@ -1025,10 +1025,10 @@ export class RateLimiterBuilderTest { .addTarget({ rules: [{ type: 'second', limit: 1 }], metadata: { baseUrl: 'http://api1.com' } }) .retryStrategy(({ attempt }) => { if (attempt === 1) { - return { type: 'retry_other', cooldownMs: 100 } + return { type: 'retry_other', currentTargetCooldownMs: 100 } } - return { type: 'fail', cooldownMs: 100 } + return { type: 'fail', currentTargetCooldownMs: 100 } }) let retriedForCount = 0 @@ -1053,7 +1053,7 @@ export class RateLimiterBuilderTest { .addTarget({ rules: [{ type: 'second', limit: 1 }], metadata: { baseUrl: 'http://api0.com' } }) .addTarget({ rules: [{ type: 'second', limit: 1 }], metadata: { baseUrl: 'http://api1.com' } }) .retryStrategy(() => { - return { type: 'retry_other', cooldownMs: 100 } + return { type: 'retry_other', currentTargetCooldownMs: 100 } }) const targetUsed = [] @@ -1082,7 +1082,7 @@ export class RateLimiterBuilderTest { .addTarget({ rules: [{ type: 'second', limit: 1 }], metadata: { baseUrl: 'http://api0.com' } }) .addTarget({ rules: [{ type: 'second', limit: 1 }], metadata: { baseUrl: 'http://api1.com' } }) .retryStrategy(() => { - return { type: 'retry_same', cooldownMs: 100 } + return { type: 'retry_same', currentTargetCooldownMs: 100 } }) const targetUsed = []