feat: update remaining and reset for targets

package.json

@@ -1,6 +1,6 @@
 {
   "name": "@athenna/ratelimiter",
-  "version": "5.4.0",
+  "version": "5.5.0",
   "description": "Respect the rate limit rules of API's you need to consume.",
   "license": "MIT",
   "author": "João Lenon <lenon@athenna.io>",

src/ratelimiter/RateLimitStore.ts

@@ -152,4 +152,195 @@ export class RateLimitStore extends Macroable {
 
     return reserve
   }
+
+  /**
+   * 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 setRemaining(
+    key: string,
+    ruleType: RateLimitRule['type'],
+    remaining: number,
+    rules: RateLimitRule[]
+  ) {
+    debug(
+      'setting remaining to %d for rule type %s in %s store for key %s',
+      remaining,
+      ruleType,
+      this.options.store,
+      key
+    )
+
+    const now = Date.now()
+    const cache = Cache.store(this.options.store)
+    const buckets = await this.getOrInit(key, rules)
+    const ruleIndex = rules.findIndex(rule => rule.type === ruleType)
+
+    if (ruleIndex === -1) {
+      debug('rule type %s not found for key %s', ruleType, key)
+      return
+    }
+
+    const rule = rules[ruleIndex]
+    const bucket = buckets[ruleIndex]
+    const used = Math.max(0, rule.limit - remaining)
+
+    bucket.length = 0
+
+    for (let i = 0; i < used; i++) {
+      bucket.push(now)
+    }
+
+    await cache.set(key, JSON.stringify(buckets))
+
+    debug(
+      'updated bucket for rule type %s: %d used, %d remaining',
+      ruleType,
+      used,
+      remaining
+    )
+  }
+
+  /**
+   * 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(
+    key: string,
+    ruleType: RateLimitRule['type'],
+    rules: RateLimitRule[]
+  ) {
+    debug(
+      'getting remaining for rule type %s in %s store for key %s',
+      ruleType,
+      this.options.store,
+      key
+    )
+
+    const now = Date.now()
+    const buckets = await this.getOrInit(key, rules)
+    const ruleIndex = rules.findIndex(rule => rule.type === ruleType)
+
+    if (ruleIndex === -1) {
+      debug('rule type %s not found for key %s', ruleType, key)
+      return 0
+    }
+
+    const rule = rules[ruleIndex]
+    const bucket = buckets[ruleIndex]
+    const window = this.options.windowMs[rule.type]
+
+    while (bucket.length && bucket[0] <= now - window) {
+      bucket.shift()
+    }
+
+    const remaining = Math.max(0, rule.limit - bucket.length)
+
+    debug('remaining for rule type %s: %d', ruleType, remaining)
+
+    return remaining
+  }
+
+  /**
+   * 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.
+   * Returns the Unix timestamp in milliseconds when the oldest request expires.
+   */
+  public async getResetAt(
+    key: string,
+    ruleType: RateLimitRule['type'],
+    rules: RateLimitRule[]
+  ) {
+    debug(
+      'getting reset time for rule type %s in %s store for key %s',
+      ruleType,
+      this.options.store,
+      key
+    )
+
+    const now = Date.now()
+    const buckets = await this.getOrInit(key, rules)
+    const ruleIndex = rules.findIndex(rule => rule.type === ruleType)
+
+    if (ruleIndex === -1) {
+      debug('rule type %s not found for key %s', ruleType, key)
+      return now
+    }
+
+    const rule = rules[ruleIndex]
+    const bucket = buckets[ruleIndex]
+    const window = this.options.windowMs[rule.type]
+
+    while (bucket.length && bucket[0] <= now - window) {
+      bucket.shift()
+    }
+
+    if (bucket.length === 0) {
+      debug('bucket empty, resets now')
+      return now
+    }
+
+    const earliestTimestamp = bucket[0]
+    const resetAt = earliestTimestamp + window
+
+    debug('reset time for rule type %s: %d', ruleType, resetAt)
+
+    return resetAt
+  }
+
+  /**
+   * 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 setResetAt(
+    key: string,
+    ruleType: RateLimitRule['type'],
+    secondsUntilReset: number,
+    rules: RateLimitRule[]
+  ) {
+    debug(
+      'setting reset time to %d seconds for rule type %s in %s store for key %s',
+      secondsUntilReset,
+      ruleType,
+      this.options.store,
+      key
+    )
+
+    const now = Date.now()
+    const cache = Cache.store(this.options.store)
+    const buckets = await this.getOrInit(key, rules)
+    const ruleIndex = rules.findIndex(rule => rule.type === ruleType)
+
+    if (ruleIndex === -1) {
+      debug('rule type %s not found for key %s', ruleType, key)
+      return
+    }
+
+    const rule = rules[ruleIndex]
+    const bucket = buckets[ruleIndex]
+    const window = this.options.windowMs[rule.type]
+
+    while (bucket.length && bucket[0] <= now - window) {
+      bucket.shift()
+    }
+
+    if (bucket.length === 0) {
+      debug('bucket empty, nothing to shift')
+      return
+    }
+
+    const targetResetAt = now + secondsUntilReset * 1000
+    const earliestTimestamp = bucket[0]
+    const currentResetAt = earliestTimestamp + window
+    const timeDiff = targetResetAt - currentResetAt
+
+    for (let i = 0; i < bucket.length; i++) {
+      bucket[i] = bucket[i] + timeDiff
+    }
+
+    await cache.set(key, JSON.stringify(buckets))
+
+    debug('shifted timestamps by %d ms for rule type %s', timeDiff, ruleType)
+  }
 }

src/ratelimiter/RateLimiterBuilder.ts

@@ -698,8 +698,43 @@ 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 }))
+        .then(() => item.run({ signal: item.signal, target: enrichedTarget }))
         .then(result => {
           this.releaseTask({ isToScheduleTick: true })
 
@@ -776,8 +811,49 @@ 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
+        )
+      }
+    }
+
     Promise.resolve()
-      .then(() => item.run({ signal: item.signal }))
+      .then(() => item.run({ signal: item.signal, target: implicitTarget }))
       .then(result => {
         this.releaseTask({ isToScheduleTick: true })
 
@@ -907,7 +983,8 @@ export class RateLimiterBuilder extends Macroable {
       error: options.error,
       signal: options.item.signal,
       attempt: options.item.attempt,
-      target: options.target
+      target: options.target,
+      targets: this.options.targets
     }
 
     const decision = await this.options.retryStrategy(ctx)

src/types/RateLimitRetryCtx.ts

@@ -34,4 +34,9 @@ export type RateLimitRetryCtx = {
    * The target that this retry is currently using.
    */
   target?: RateLimitTarget
+
+  /**
+   * All the targets that are available to use for this retry.
+   */
+  targets?: RateLimitTarget[]
 }