diff --git a/package-lock.json b/package-lock.json index a3ace32..4ac4a28 100644 --- a/package-lock.json +++ b/package-lock.json @@ -1,12 +1,12 @@ { "name": "@athenna/ratelimiter", - "version": "5.11.0", + "version": "5.13.0", "lockfileVersion": 3, "requires": true, "packages": { "": { "name": "@athenna/ratelimiter", - "version": "5.11.0", + "version": "5.13.0", "license": "MIT", "devDependencies": { "@athenna/cache": "^5.2.0", diff --git a/package.json b/package.json index 49b132a..225afed 100644 --- a/package.json +++ b/package.json @@ -1,6 +1,6 @@ { "name": "@athenna/ratelimiter", - "version": "5.11.0", + "version": "5.13.0", "description": "Respect the rate limit rules of API's you need to consume.", "license": "MIT", "author": "João Lenon ", diff --git a/src/ratelimiter/RateLimitStore.ts b/src/ratelimiter/RateLimitStore.ts index f2eab61..3344b42 100644 --- a/src/ratelimiter/RateLimitStore.ts +++ b/src/ratelimiter/RateLimitStore.ts @@ -38,15 +38,27 @@ export class RateLimitStore extends Macroable { public async getOrInit(key: string, rules: RateLimitRule[]) { const cache = Cache.store(this.options.store) - let buckets = await cache.get(key) + const buckets = await cache.get(key) if (!buckets) { - buckets = JSON.stringify(rules.map(() => [])) + const initialized = JSON.stringify(rules.map(() => [])) - await cache.set(key, buckets) + await cache.set(key, initialized) + + return JSON.parse(initialized) as number[][] + } + + const parsed = JSON.parse(buckets) as number[][] + + if (parsed.length !== rules.length) { + const reconciled = rules.map((_, i) => parsed[i] ?? []) + + await cache.set(key, JSON.stringify(reconciled)) + + return reconciled } - return JSON.parse(buckets) as number[][] + return parsed } /** diff --git a/src/ratelimiter/RateLimiterBuilder.ts b/src/ratelimiter/RateLimiterBuilder.ts index fdb9c16..6fe71e2 100644 --- a/src/ratelimiter/RateLimiterBuilder.ts +++ b/src/ratelimiter/RateLimiterBuilder.ts @@ -64,8 +64,9 @@ export class RateLimiterBuilder extends Macroable { /** * Index for when using round_robin selection strategy. + * Starts as null to indicate it has not been randomly initialized yet. */ - private rrIndex = 0 + private rrIndex: number | null = null /** * Holds the setTimeout id to be able to disable it @@ -627,6 +628,10 @@ export class RateLimiterBuilder extends Macroable { const nextItem = this.queue[0] + if (this.options.targetSelectionStrategy === 'round_robin' && this.rrIndex === null) { + this.rrIndex = Math.floor(Math.random() * this.options.targets.length) + } + for (const i of this.createIdxBySelectionStrategy(nextItem)) { const possibleTarget = this.options.targets[i] const key = possibleTarget.getKey() diff --git a/tests/unit/ratelimiter/RateLimiterBuilderTest.ts b/tests/unit/ratelimiter/RateLimiterBuilderTest.ts index 4ecfcd9..0d3eeb6 100644 --- a/tests/unit/ratelimiter/RateLimiterBuilderTest.ts +++ b/tests/unit/ratelimiter/RateLimiterBuilderTest.ts @@ -791,6 +791,103 @@ export class RateLimiterBuilderTest { assert.deepEqual(results, ['http://api2.com', 'http://api2.com']) } + @Test() + public async shouldRotateTargetsSequentiallyWhenUsingRoundRobinStrategy({ assert }: Context) { + const targets = [ + { id: 'api1', metadata: { baseUrl: 'http://api1.com' } }, + { id: 'api2', metadata: { baseUrl: 'http://api2.com' } }, + { id: 'api3', metadata: { baseUrl: 'http://api3.com' } } + ] + + const limiter = RateLimiter.build() + .store('memory', { windowMs: { second: 1000 } }) + .key('request:rr-rotation') + .targetSelectionStrategy('round_robin') + .addRule({ type: 'second', limit: 100 }) + .setTargets(targets) + + const results: string[] = [] + + for (let i = 0; i < 6; i++) { + const result = await limiter.schedule(({ target }) => target.metadata.baseUrl) + results.push(result) + } + + const urls = ['http://api1.com', 'http://api2.com', 'http://api3.com'] + const startIdx = urls.indexOf(results[0]) + const expected = Array.from({ length: 6 }, (_, i) => urls[(startIdx + i) % 3]) + + assert.deepEqual(results, expected) + } + + @Test() + public async shouldRotateTargetsConcurrentlyWhenUsingRoundRobinStrategy({ assert }: Context) { + const api1 = { id: 'api1', metadata: { baseUrl: 'http://api1.com' } } + const api2 = { id: 'api2', metadata: { baseUrl: 'http://api2.com' } } + + const limiter = RateLimiter.build() + .maxConcurrent(2) + .store('memory', { windowMs: { second: 100 } }) + .key('request:rr-rotation-concurrent') + .targetSelectionStrategy('round_robin') + .addRule({ type: 'second', limit: 10 }) + .addTarget(api1) + .addTarget(api2) + + const used: string[] = [] + const barrier = this.createBarrier() + + const run = async ({ target }) => { + used.push(target.metadata.baseUrl) + await barrier.wait() + return target.metadata.baseUrl + } + + const p1 = limiter.schedule(run) + const p2 = limiter.schedule(run) + + await this.waitUntil(() => used.length === 2, 10, 1000) + + used.sort() + + assert.deepEqual(used, ['http://api1.com', 'http://api2.com']) + + barrier.release() + + const results = await Promise.all([p1, p2]) + + results.sort() + + assert.deepEqual(results, ['http://api1.com', 'http://api2.com']) + } + + @Test() + public async shouldDistributeStartingTargetAcrossMultipleInstancesWhenUsingRoundRobinStrategy({ assert }: Context) { + const api1 = { id: 'api1', metadata: { baseUrl: 'http://api1.com' } } + const api2 = { id: 'api2', metadata: { baseUrl: 'http://api2.com' } } + + const firstTargets: string[] = [] + + for (let run = 0; run < 20; run++) { + const limiter = RateLimiter.build() + .store('memory', { windowMs: { second: 1000 } }) + .key(`request:rr-dist-${run}`) + .targetSelectionStrategy('round_robin') + .addRule({ type: 'second', limit: 100 }) + .addTarget(api1) + .addTarget(api2) + + const result = await limiter.schedule(({ target }) => target.metadata.baseUrl) + firstTargets.push(result) + } + + const api1Count = firstTargets.filter(t => t === 'http://api1.com').length + const api2Count = firstTargets.filter(t => t === 'http://api2.com').length + + assert.isAbove(api1Count, 0) + assert.isAbove(api2Count, 0) + } + @Test() public async shouldThrowMissingRuleExceptionIfRateLimiterRulesAndTargetRulesAreNotDefined({ assert }: Context) { const limiter = RateLimiter.build()