diff --git a/package-lock.json b/package-lock.json index 62d277e..f66355c 100644 --- a/package-lock.json +++ b/package-lock.json @@ -1,12 +1,12 @@ { "name": "@athenna/ratelimiter", - "version": "5.5.0", + "version": "5.6.0", "lockfileVersion": 3, "requires": true, "packages": { "": { "name": "@athenna/ratelimiter", - "version": "5.5.0", + "version": "5.6.0", "license": "MIT", "devDependencies": { "@athenna/cache": "^5.2.0", diff --git a/package.json b/package.json index b4d2628..88d1a1d 100644 --- a/package.json +++ b/package.json @@ -1,6 +1,6 @@ { "name": "@athenna/ratelimiter", - "version": "5.5.0", + "version": "5.6.0", "description": "Respect the rate limit rules of API's you need to consume.", "license": "MIT", "author": "João Lenon ", diff --git a/src/index.ts b/src/index.ts index 6e63b54..8d5b218 100644 --- a/src/index.ts +++ b/src/index.ts @@ -11,4 +11,5 @@ export * from '#src/types' export * from '#src/ratelimiter/RateLimiter' export * from '#src/ratelimiter/RateLimitStore' +export * from '#src/ratelimiter/RateLimitTarget' export * from '#src/ratelimiter/RateLimiterBuilder' diff --git a/src/ratelimiter/RateLimitTarget.ts b/src/ratelimiter/RateLimitTarget.ts new file mode 100644 index 0000000..6f7b4ea --- /dev/null +++ b/src/ratelimiter/RateLimitTarget.ts @@ -0,0 +1,115 @@ +/** + * @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, + RateLimitRawTarget, + RateLimiterOptions +} from '#src/types' + +import { Config } from '@athenna/config' +import { Json, String, Macroable } from '@athenna/common' + +export class RateLimitTarget extends Macroable { + /** + * The rate limit target ID. By default this will be created by creating + * a hash from the target metadata object, but you can also define your + * own ID. + */ + public id?: string + + /** + * 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 Target. With this + * implementation you can create not only API rotations but also API + * Keys rotations at the same time. + */ + public metadata: Record + + /** + * The options that were used to create the rate limiter. + */ + public options: RateLimiterOptions + + /** + * Custom rate limit rules for this target. If not defined, + * the default defined in RateLimiter will be used. + */ + public rules?: RateLimitRule[] + + public constructor(raw: RateLimitRawTarget, options: RateLimiterOptions) { + super() + + this.options = options + this.metadata = raw.metadata || {} + this.rules = raw.rules?.length ? raw.rules : options.rules + + this.id = + raw.id || + String.hash(JSON.stringify(Json.sort(this.metadata)), { + key: Config.get('app.key', 'ratelimiter') + }) + } + + /** + * Create a custom key for an target to be used to map the + * target rules into the cache. + */ + public getKey() { + if (this.id === '__implicit__') { + return this.options.key + } + + return `${this.options.key}:${this.id}` + } + + /** + * Get the current number of remaining requests for a specific rule type. + * This exposes the internal rate limiter state for comparison with API headers. + */ + public async getRemaining(type: RateLimitRule['type']) { + return this.options.store!.getRemaining(this.getKey(), type, this.rules) + } + + /** + * Get the timestamp when the rate limit will reset for a specific rule type. + * This exposes the internal rate limiter state for comparison with API headers. + */ + public async getResetAt(type: RateLimitRule['type']) { + return this.options.store!.getResetAt(this.getKey(), type, this.rules) + } + + /** + * Manually update the remaining request count for a specific rule type. + * This allows syncing the internal rate limiter state with external API + * rate limit headers. + */ + public async updateRemaining(remaining: number, type: RateLimitRule['type']) { + await this.options.store!.setRemaining( + this.getKey(), + type, + remaining, + this.rules + ) + } + + /** + * Manually update the reset time for a specific rule type based on API headers. + * This shifts all timestamps in the bucket to align with the API's reset schedule. + */ + public async updateResetAt(seconds: number, type: RateLimitRule['type']) { + await this.options.store!.setResetAt( + this.getKey(), + type, + seconds, + this.rules + ) + } +} diff --git a/src/ratelimiter/RateLimiterBuilder.ts b/src/ratelimiter/RateLimiterBuilder.ts index 9dd4710..af42744 100644 --- a/src/ratelimiter/RateLimiterBuilder.ts +++ b/src/ratelimiter/RateLimiterBuilder.ts @@ -11,19 +11,19 @@ import type { QueueItem, RateLimitRule, ScheduleOptions, - RateLimitTarget, RateLimitRetryCtx, RateLimiterOptions, + RateLimitRawTarget, + RateLimitPendingCtx, RateLimitScheduleCtx, RateLimitStoreOptions, - RateLimitRetryDecision, - RateLimitPendingCtx + RateLimitRetryDecision } from '#src/types' import { debug } from '#src/debug' -import { Config } from '@athenna/config' +import { Macroable, Options } from '@athenna/common' import { RateLimitStore } from '#src/ratelimiter/RateLimitStore' -import { Json, String, Macroable, Options } from '@athenna/common' +import { RateLimitTarget } from '#src/ratelimiter/RateLimitTarget' import { MissingKeyException } from '#src/exceptions/MissingKeyException' import { MissingRuleException } from '#src/exceptions/MissingRuleException' import { MissingStoreException } from '#src/exceptions/MissingStoreException' @@ -226,16 +226,12 @@ export class RateLimiterBuilder extends Macroable { * await limiter.schedule(() => {...}) * ``` */ - public addTarget(target: RateLimitTarget) { + public addTarget(raw: RateLimitRawTarget) { if (!this.options.targets) { this.options.targets = [] } - if (!target.id) { - target.id = this.getTargetId(target) - } - - this.options.targets.push(target) + this.options.targets.push(new RateLimitTarget(raw, this.options)) return this } @@ -273,8 +269,8 @@ export class RateLimiterBuilder extends Macroable { * await limiter.schedule(() => {...}) * ``` */ - public setTargets(targets: RateLimitTarget[]) { - targets.forEach(target => this.addTarget(target)) + public setTargets(raws: RateLimitRawTarget[]) { + raws.forEach(raw => this.addTarget(raw)) return this } @@ -509,21 +505,21 @@ export class RateLimiterBuilder extends Macroable { } /** - * Create a custom id for an target by reading the metadata object. - * The object will always be sorted by keys. - */ - 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 target to be used to map the - * target rules into the cache. + * Get the target instance from a raw target. This method doesn't register the + * target in the limiter, it only returns the target instance. + * + * @example + * ```ts + * const limiter = RateLimiter.build() + * .store('memory') + * .key('request:/profile') + * + * const api1 = { metadata: { baseUrl: 'http://api1.com' } } + * const key = limiter.getTarget(api1).getKey() + * ``` */ - public createTargetKey(target: RateLimitTarget) { - return `${this.options.key}:${this.getTargetId(target)}` + public getTarget(raw: RateLimitRawTarget) { + return new RateLimitTarget(raw, this.options) } /** @@ -632,19 +628,19 @@ export class RateLimiterBuilder extends Macroable { const nextItem = this.queue[0] for (const i of this.createIdxBySelectionStrategy(nextItem)) { - const key = this.createTargetKey(this.options.targets[i]) - - const rules = this.options.targets[i].rules?.length - ? this.options.targets[i].rules - : this.options.rules + const possibleTarget = this.options.targets[i] + const key = possibleTarget.getKey() try { - const res = await this.options.store.tryReserve(key, rules) + const res = await this.options.store.tryReserve( + key, + possibleTarget.rules + ) this.storeErrorCount = 0 if (res.allowed) { - target = this.options.targets[i] + target = possibleTarget if (this.options.targetSelectionStrategy === 'round_robin') { this.rrIndex = (i + 1) % this.options.targets.length @@ -698,43 +694,8 @@ export class RateLimiterBuilder extends Macroable { this.active++ - const targetKey = this.createTargetKey(target) - const rules = target.rules?.length ? target.rules : this.options.rules - - const enrichedTarget = { - ...target, - getRemaining: async (ruleType: RateLimitRule['type']) => { - return this.options.store!.getRemaining(targetKey, ruleType, rules) - }, - getResetAt: async (ruleType: RateLimitRule['type']) => { - return this.options.store!.getResetAt(targetKey, ruleType, rules) - }, - updateRemaining: async ( - remaining: number, - ruleType: RateLimitRule['type'] - ) => { - await this.options.store!.setRemaining( - targetKey, - ruleType, - remaining, - rules - ) - }, - updateResetAt: async ( - secondsUntilReset: number, - ruleType: RateLimitRule['type'] - ) => { - await this.options.store!.setResetAt( - targetKey, - ruleType, - secondsUntilReset, - rules - ) - } - } - Promise.resolve() - .then(() => item.run({ signal: item.signal, target: enrichedTarget })) + .then(() => item.run({ signal: item.signal, target })) .then(result => { this.releaseTask({ isToScheduleTick: true }) @@ -811,49 +772,10 @@ export class RateLimiterBuilder extends Macroable { this.active++ - const implicitTarget = { - id: '__implicit__', - metadata: {}, - getRemaining: async (ruleType: RateLimitRule['type']) => { - return this.options.store!.getRemaining( - this.options.key, - ruleType, - this.options.rules - ) - }, - getResetAt: async (ruleType: RateLimitRule['type']) => { - return this.options.store!.getResetAt( - this.options.key, - ruleType, - this.options.rules - ) - }, - updateRemaining: async ( - remaining: number, - ruleType: RateLimitRule['type'] - ) => { - await this.options.store!.setRemaining( - this.options.key, - ruleType, - remaining, - this.options.rules - ) - }, - updateResetAt: async ( - secondsUntilReset: number, - ruleType: RateLimitRule['type'] - ) => { - await this.options.store!.setResetAt( - this.options.key, - ruleType, - secondsUntilReset, - this.options.rules - ) - } - } + const target = new RateLimitTarget({ id: '__implicit__' }, this.options) Promise.resolve() - .then(() => item.run({ signal: item.signal, target: implicitTarget })) + .then(() => item.run({ signal: item.signal, target })) .then(result => { this.releaseTask({ isToScheduleTick: true }) @@ -976,7 +898,7 @@ export class RateLimiterBuilder extends Macroable { return } - const key = this.createTargetKey(options.target) + const key = options.target.getKey() const ctx: RateLimitRetryCtx = { key, diff --git a/src/types/RateLimitTarget.ts b/src/types/RateLimitRawTarget.ts similarity index 93% rename from src/types/RateLimitTarget.ts rename to src/types/RateLimitRawTarget.ts index a540775..99f1f22 100644 --- a/src/types/RateLimitTarget.ts +++ b/src/types/RateLimitRawTarget.ts @@ -9,7 +9,7 @@ import type { RateLimitRule } from '#src/types' -export type RateLimitTarget = { +export type RateLimitRawTarget = { /** * The rate limit target ID. By default this will be created by creating * a hash from the target metadata object, but you can also define your @@ -24,7 +24,7 @@ export type RateLimitTarget = { * implementation you can create not only API rotations but also API * Keys rotations at the same time. */ - metadata: Record + metadata?: Record /** * Custom rate limit rules for this target. If not defined, diff --git a/src/types/RateLimitScheduleCtx.ts b/src/types/RateLimitScheduleCtx.ts index a048e2c..f546378 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 { RateLimitTarget, RateLimitRule } from '#src/types' +import type { RateLimitTarget } from '#src/ratelimiter/RateLimitTarget' export type RateLimitScheduleCtx = { /** @@ -19,48 +19,5 @@ export type RateLimitScheduleCtx = { * The target that this schedule is currently using. This will always * be present, even in single mode where an implicit target is created. */ - target: RateLimitTarget & { - /** - * Get the current number of remaining requests for a specific rule type. - * This exposes the internal rate limiter state for comparison with API headers. - * - * @param ruleType The type of rule to query ('second', 'minute', etc.) - * @returns The number of remaining requests - */ - getRemaining: (ruleType: RateLimitRule['type']) => Promise - - /** - * Get the timestamp when the rate limit will reset for a specific rule type. - * This exposes the internal rate limiter state for comparison with API headers. - * - * @param ruleType The type of rule to query ('second', 'minute', etc.) - * @returns Unix timestamp in milliseconds when the oldest request expires - */ - getResetAt: (ruleType: RateLimitRule['type']) => Promise - - /** - * Manually update the remaining request count for a specific rule type. - * This allows syncing the internal rate limiter state with external API - * rate limit headers. - * - * @param remaining The number of remaining requests - * @param ruleType The type of rule to update ('second', 'minute', etc.) - */ - updateRemaining: ( - remaining: number, - ruleType: RateLimitRule['type'] - ) => Promise - - /** - * Manually update the reset time for a specific rule type based on API headers. - * This shifts all timestamps in the bucket to align with the API's reset schedule. - * - * @param secondsUntilReset Number of seconds until the rate limit resets (from API header) - * @param ruleType The type of rule to update ('second', 'minute', etc.) - */ - updateResetAt: ( - secondsUntilReset: number, - ruleType: RateLimitRule['type'] - ) => Promise - } + target: RateLimitTarget } diff --git a/src/types/index.ts b/src/types/index.ts index 0899b7c..ea9bf56 100644 --- a/src/types/index.ts +++ b/src/types/index.ts @@ -11,7 +11,7 @@ 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/RateLimitRawTarget' export * from '#src/types/RateLimitRetryCtx' export * from '#src/types/RateLimiterOptions' export * from '#src/types/RateLimitPendingCtx' diff --git a/tests/unit/ratelimiter/RateLimiterBuilderTest.ts b/tests/unit/ratelimiter/RateLimiterBuilderTest.ts index 2ff460e..c40f36d 100644 --- a/tests/unit/ratelimiter/RateLimiterBuilderTest.ts +++ b/tests/unit/ratelimiter/RateLimiterBuilderTest.ts @@ -655,7 +655,7 @@ export class RateLimiterBuilderTest { const store = new RateLimitStore({ store: 'memory', windowMs: { second: 100 } }) - await store.setCooldown(limiter.createTargetKey(api1), 5000) + await store.setCooldown(limiter.getTarget(api1).getKey(), 5000) await Sleep.for(10).milliseconds().wait() const result = await limiter.schedule(({ target }) => { @@ -682,7 +682,7 @@ export class RateLimiterBuilderTest { const store = new RateLimitStore({ store: 'memory', windowMs: { second: 100 } }) - await store.setCooldown(limiter.createTargetKey(api1), 5000) + await store.setCooldown(limiter.getTarget(api1).getKey(), 5000) await Sleep.for(50).milliseconds().wait() const used: string[] = [] @@ -731,7 +731,7 @@ export class RateLimiterBuilderTest { const store = new RateLimitStore({ store: 'memory', windowMs: { second: 100 } }) - await store.setCooldown(limiter.createTargetKey(api1), 5000) + await store.setCooldown(limiter.getTarget(api1).getKey(), 5000) await Sleep.for(10).milliseconds().wait() const result = await limiter.schedule(({ target }) => { @@ -759,7 +759,7 @@ export class RateLimiterBuilderTest { const store = new RateLimitStore({ store: 'memory', windowMs: { second: 100 } }) - await store.setCooldown(limiter.createTargetKey(api1), 5000) + await store.setCooldown(limiter.getTarget(api1).getKey(), 5000) await Sleep.for(50).milliseconds().wait() const used: string[] = []