Ensure RateLimiter utilizes the token-bucket algorithm (#2097)

Co-authored-by: Tim <hello@timsmart.co>

Author: IMax153

.changeset/silly-yaks-allow.md

@@ -0,0 +1,42 @@
+---
+"effect": patch
+---
+
+Updates the `RateLimiter.make` constructor to take an object of `RateLimiter.Options`, which allows for specifying the rate-limiting algorithm to utilize:
+
+You can choose from either the `token-bucket` or the `fixed-window` algorithms for rate-limiting.
+
+```ts
+export declare namespace RateLimiter {
+  export interface Options {
+    /**
+     * The maximum number of requests that should be allowed.
+     */
+    readonly limit: number
+    /**
+     * The interval to utilize for rate-limiting requests. The semantics of the
+     * specified `interval` vary depending on the chosen `algorithm`:
+     *
+     * `token-bucket`: The maximum number of requests will be spread out over
+     * the provided interval if no tokens are available.
+     *
+     * For example, for a `RateLimiter` using the `token-bucket` algorithm with
+     * a `limit` of `10` and an `interval` of `1 seconds`, `1` request can be
+     * made every `100 millis`.
+     *
+     * `fixed-window`: The maximum number of requests will be reset during each
+     * interval. For example, for a `RateLimiter` using the `fixed-window`
+     * algorithm with a `limit` of `10` and an `interval` of `1 seconds`, a
+     * maximum of `10` requests can be made each second.
+     */
+    readonly interval: DurationInput
+    /**
+     * The algorithm to utilize for rate-limiting requests.
+     *
+     * Defaults to `token-bucket`.
+     */
+    readonly algorithm?: "fixed-window" | "token-bucket"
+  }
+}
+```
+

.changeset/tender-apples-grin.md

@@ -0,0 +1,5 @@
+---
+"effect": patch
+---
+
+return the resulting available permits from Semaphore.release

packages/effect/src/Effect.ts

@@ -4673,10 +4673,14 @@ export interface Permit {
  * @since 2.0.0
  */
 export interface Semaphore {
+  /** when the given amount of permits are available, run the effect and release the permits when finished */
   withPermits(permits: number): <A, E, R>(self: Effect<A, E, R>) => Effect<A, E, R>
+  /** take the given amount of permits, suspending if they are not yet available */
   take(permits: number): Effect<number>
-  release(permits: number): Effect<void>
-  releaseAll: Effect<void>
+  /** release the given amount of permits, and return the resulting available permits */
+  release(permits: number): Effect<number>
+  /** release all the taken permits, and return the resulting available permits */
+  releaseAll: Effect<number>
 }
 
 /**

packages/effect/src/RateLimiter.ts

@@ -1,6 +1,5 @@
 /**
- * Limits the number of calls to a resource to a maximum amount in some interval
- * using the token bucket algorithm.
+ * Limits the number of calls to a resource to a maximum amount in some interval.
  *
  * @since 2.0.0
  */
@@ -10,6 +9,11 @@ import * as internal from "./internal/rateLimiter.js"
 import type { Scope } from "./Scope.js"
 
 /**
+ * Limits the number of calls to a resource to a maximum amount in some interval.
+ *
+ * Note that only the moment of starting the effect is rate limited: the number
+ * of concurrent executions is not bounded.
+ *
  * @since 2.0.0
  * @category models
  */
@@ -18,10 +22,47 @@ export interface RateLimiter {
 }
 
 /**
- * Constructs a new `RateLimiter` with the specified limit and window.
- *
- * Limits the number of calls to a resource to a maximum amount in some interval
- * using the token bucket algorithm.
+ * @since 2.0.0
+ */
+export declare namespace RateLimiter {
+  /**
+   * @since 2.0.0
+   * @category models
+   */
+  export interface Options {
+    /**
+     * The maximum number of requests that should be allowed.
+     */
+    readonly limit: number
+    /**
+     * The interval to utilize for rate-limiting requests. The semantics of the
+     * specified `interval` vary depending on the chosen `algorithm`:
+     *
+     * `token-bucket`: The maximum number of requests will be spread out over
+     * the provided interval if no tokens are available.
+     *
+     * For example, for a `RateLimiter` using the `token-bucket` algorithm with
+     * a `limit` of `10` and an `interval` of `1 seconds`, `1` request can be
+     * made every `100 millis`.
+     *
+     * `fixed-window`: The maximum number of requests will be reset during each
+     * interval. For example, for a `RateLimiter` using the `fixed-window`
+     * algorithm with a `limit` of `10` and an `interval` of `1 seconds`, a
+     * maximum of `10` requests can be made each second.
+     */
+    readonly interval: DurationInput
+    /**
+     * The algorithm to utilize for rate-limiting requests.
+     *
+     * Defaults to `token-bucket`.
+     */
+    readonly algorithm?: "fixed-window" | "token-bucket"
+  }
+}
+
+/**
+ * Constructs a new `RateLimiter` which will utilize the specified algorithm
+ * to limit requests (defaults to `token-bucket`).
  *
  * Notes
  * - Only the moment of starting the effect is rate limited. The number of concurrent executions is not bounded.
@@ -34,8 +75,8 @@ export interface RateLimiter {
  *
  * const program = Effect.scoped(
  *   Effect.gen(function* ($) {
- *     const perMinuteRL = yield* $(RateLimiter.make(30, "1 minutes"))
- *     const perSecondRL = yield* $(RateLimiter.make(2, "1 seconds"))
+ *     const perMinuteRL = yield* $(RateLimiter.make({ limit: 30, interval: "1 minutes" }))
+ *     const perSecondRL = yield* $(RateLimiter.make({ limit: 2, interval: "1 seconds" }))
  *
  *     // This rate limiter respects both the 30 calls per minute
  *     // and the 2 calls per second constraints.
@@ -52,11 +93,7 @@ export interface RateLimiter {
  * @since 2.0.0
  * @category constructors
  */
-export const make: (limit: number, window: DurationInput) => Effect<
-  RateLimiter,
-  never,
-  Scope
-> = internal.make
+export const make: (options: RateLimiter.Options) => Effect<RateLimiter, never, Scope> = internal.make
 
 /**
  * Alters the per-effect cost of the rate-limiter.
@@ -72,7 +109,7 @@ export const make: (limit: number, window: DurationInput) => Effect<
  * const program = Effect.scoped(
  *   Effect.gen(function* ($) {
  *     // Create a rate limiter that has an hourly limit of 1000 credits
- *     const rateLimiter = yield* $(RateLimiter.make(1000, "1 hours"));
+ *     const rateLimiter = yield* $(RateLimiter.make({ limit: 1000, interval: "1 hours" }));
  *     // Query API costs 1 credit per call ( 1 is the default cost )
  *     const queryAPIRL = compose(rateLimiter, RateLimiter.withCost(1));
  *     // Mutation API costs 5 credits per call

packages/effect/src/index.ts

@@ -545,8 +545,7 @@ export * as Queue from "./Queue.js"
 export * as Random from "./Random.js"
 
 /**
- * Limits the number of calls to a resource to a maximum amount in some interval
- * using the token bucket algorithm.
+ * Limits the number of calls to a resource to a maximum amount in some interval.
  *
  * @since 2.0.0
  */

packages/effect/src/internal/effect/circular.ts

@@ -65,22 +65,24 @@ class Semaphore {
       return Either.right(core.succeed(n))
     })
 
-  readonly updateTaken = (f: (n: number) => number): Effect.Effect<void> =>
+  readonly updateTaken = (f: (n: number) => number): Effect.Effect<number> =>
     core.withFiberRuntime((fiber) => {
       this.taken = f(this.taken)
-      fiber.getFiberRef(currentScheduler).scheduleTask(() => {
-        const iter = this.waiters.values()
-        let item = iter.next()
-        while (item.done === false && item.value() === true) {
-          item = iter.next()
-        }
-      }, fiber.getFiberRef(core.currentSchedulingPriority))
-      return core.unit
+      if (this.waiters.size > 0) {
+        fiber.getFiberRef(currentScheduler).scheduleTask(() => {
+          const iter = this.waiters.values()
+          let item = iter.next()
+          while (item.done === false && item.value() === true) {
+            item = iter.next()
+          }
+        }, fiber.getFiberRef(core.currentSchedulingPriority))
+      }
+      return core.succeed(this.free)
     })
 
-  readonly release = (n: number): Effect.Effect<void> => this.updateTaken((taken) => taken - n)
+  readonly release = (n: number): Effect.Effect<number> => this.updateTaken((taken) => taken - n)
 
-  readonly releaseAll: Effect.Effect<void> = this.updateTaken((_) => 0)
+  readonly releaseAll: Effect.Effect<number> = this.updateTaken((_) => 0)
 
   readonly withPermits = (n: number) => <R, E, A>(self: Effect.Effect<R, E, A>) =>
     core.uninterruptibleMask((restore) =>

packages/effect/src/internal/rateLimiter.ts

@@ -1,44 +1,92 @@
 import type { DurationInput } from "../Duration.js"
+import * as Duration from "../Duration.js"
 import * as Effect from "../Effect.js"
 import * as FiberRef from "../FiberRef.js"
 import { globalValue } from "../GlobalValue.js"
-import type { RateLimiter } from "../RateLimiter.js"
-import type { Scope } from "../Scope.js"
-import * as SynchronizedRef from "../SynchronizedRef.js"
+import type * as RateLimiter from "../RateLimiter.js"
+import type * as Scope from "../Scope.js"
 
 /** @internal */
-const currentCost = globalValue(
-  Symbol.for("effect/RateLimiter/currentCost"),
-  () => FiberRef.unsafeMake(1)
-)
+export const make = ({
+  algorithm = "token-bucket",
+  interval,
+  limit
+}: RateLimiter.RateLimiter.Options): Effect.Effect<
+  RateLimiter.RateLimiter,
+  never,
+  Scope.Scope
+> => {
+  switch (algorithm) {
+    case "fixed-window": {
+      return fixedWindow(limit, interval)
+    }
+    case "token-bucket": {
+      return tokenBucket(limit, interval)
+    }
+  }
+}
 
-/** @internal */
-export const make = (limit: number, window: DurationInput): Effect.Effect<
-  RateLimiter,
+const tokenBucket = (limit: number, window: DurationInput): Effect.Effect<
+  RateLimiter.RateLimiter,
   never,
-  Scope
+  Scope.Scope
 > =>
   Effect.gen(function*(_) {
-    const scope = yield* _(Effect.scope)
+    const millisPerToken = Math.ceil(Duration.toMillis(window) / limit)
     const semaphore = yield* _(Effect.makeSemaphore(limit))
-    const ref = yield* _(SynchronizedRef.make(false))
-    const reset = SynchronizedRef.updateEffect(
-      ref,
-      (running) =>
-        running ? Effect.succeed(true) : Effect.sleep(window).pipe(
-          Effect.zipRight(SynchronizedRef.set(ref, false)),
-          Effect.zipRight(semaphore.releaseAll),
-          Effect.forkIn(scope),
-          Effect.interruptible,
-          Effect.as(true)
-        )
+    const latch = yield* _(Effect.makeSemaphore(0))
+    const refill: Effect.Effect<void> = Effect.sleep(millisPerToken).pipe(
+      Effect.zipRight(latch.releaseAll),
+      Effect.zipRight(semaphore.release(1)),
+      Effect.flatMap((free) => free === limit ? Effect.unit : refill)
     )
+    yield* _(
+      latch.take(1),
+      Effect.zipRight(refill),
+      Effect.forever,
+      Effect.forkScoped,
+      Effect.interruptible
+    )
+    const take = Effect.uninterruptibleMask((restore) =>
+      Effect.flatMap(
+        FiberRef.get(currentCost),
+        (cost) => Effect.zipRight(restore(semaphore.take(cost)), latch.release(1))
+      )
+    )
+    return (effect) => Effect.zipRight(take, effect)
+  })
 
-    const cost = FiberRef.get(currentCost).pipe(Effect.flatMap(semaphore.take))
-    const take = Effect.zipRight(cost, reset)
-
+const fixedWindow = (limit: number, window: DurationInput): Effect.Effect<
+  RateLimiter.RateLimiter,
+  never,
+  Scope.Scope
+> =>
+  Effect.gen(function*(_) {
+    const semaphore = yield* _(Effect.makeSemaphore(limit))
+    const latch = yield* _(Effect.makeSemaphore(0))
+    yield* _(
+      latch.take(1),
+      Effect.zipRight(Effect.sleep(window)),
+      Effect.zipRight(latch.releaseAll),
+      Effect.zipRight(semaphore.releaseAll),
+      Effect.forever,
+      Effect.forkScoped,
+      Effect.interruptible
+    )
+    const take = Effect.uninterruptibleMask((restore) =>
+      Effect.flatMap(
+        FiberRef.get(currentCost),
+        (cost) => Effect.zipRight(restore(semaphore.take(cost)), latch.release(1))
+      )
+    )
     return (effect) => Effect.zipRight(take, effect)
   })
 
+/** @internal */
+const currentCost = globalValue(
+  Symbol.for("effect/RateLimiter/currentCost"),
+  () => FiberRef.unsafeMake(1)
+)
+
 /** @internal */
 export const withCost = (cost: number) => Effect.locally(currentCost, cost)