Merge pull request #3 from AthennaIO/feat/retry-and-api-targets

feat: implement retry strategy and targets

Author: jlenon7

package-lock.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 <lenon@athenna.io>",
@@ -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",
@@ -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",

package.json

@@ -11,10 +11,14 @@ import { Exception } from '@athenna/common'
 
 export class MissingRuleException extends Exception {
   public constructor() {
+    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 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',
-      help: 'This errors happens when you forget to define rules for your RateLimiter instance.',
-      message: 'Missing rules value for rate limiter.'
+      help,
+      message
     })
   }
 }

src/exceptions/MissingRuleException.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'

src/index.ts

@@ -0,0 +1,155 @@
+/**
+ * @athenna/ratelimiter
+ *
+ * (c) João Lenon <lenon@athenna.io>
+ *
+ * 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 truncate() {
+    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)
+
+    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
+  }
+}