@upstash/ratelimit 0.4.3 → 0.4.5-canary.0

package/src/single.ts

@@ -0,0 +1,487 @@
+import type { Duration } from "./duration";
+import { ms } from "./duration";
+import type { Algorithm, RegionContext } from "./types";
+import type { Redis } from "./types";
+
+import { Ratelimit } from "./ratelimit";
+export type RegionRatelimitConfig = {
+  /**
+   * Instance of `@upstash/redis`
+   * @see https://github.com/upstash/upstash-redis#quick-start
+   */
+  redis: Redis;
+  /**
+   * The ratelimiter function to use.
+   *
+   * Choose one of the predefined ones or implement your own.
+   * Available algorithms are exposed via static methods:
+   * - Ratelimiter.fixedWindow
+   * - Ratelimiter.slidingWindow
+   * - Ratelimiter.tokenBucket
+   */
+  limiter: Algorithm<RegionContext>;
+  /**
+   * All keys in redis are prefixed with this.
+   *
+   * @default `@upstash/ratelimit`
+   */
+  prefix?: string;
+
+  /**
+   * If enabled, the ratelimiter will keep a global cache of identifiers, that have
+   * exhausted their ratelimit. In serverless environments this is only possible if
+   * you create the ratelimiter instance outside of your handler function. While the
+   * function is still hot, the ratelimiter can block requests without having to
+   * request data from redis, thus saving time and money.
+   *
+   * Whenever an identifier has exceeded its limit, the ratelimiter will add it to an
+   * internal list together with its reset timestamp. If the same identifier makes a
+   * new request before it is reset, we can immediately reject it.
+   *
+   * Set to `false` to disable.
+   *
+   * If left undefined, a map is created automatically, but it can only work
+   * if the map or the ratelimit instance is created outside your serverless function handler.
+   */
+  ephemeralCache?: Map<string, number> | false;
+
+  /**
+   * If set, the ratelimiter will allow requests to pass after this many milliseconds.
+   *
+   * Use this if you want to allow requests in case of network problems
+   */
+  timeout?: number;
+
+  /**
+   * If enabled, the ratelimiter will store analytics data in redis, which you can check out at
+   * https://console.upstash.com/ratelimit
+   *
+   * @default true
+   */
+  analytics?: boolean;
+};
+
+/**
+ * Ratelimiter using serverless redis from https://upstash.com/
+ *
+ * @example
+ * ```ts
+ * const { limit } = new Ratelimit({
+ *    redis: Redis.fromEnv(),
+ *    limiter: Ratelimit.slidingWindow(
+ *      "30 m", // interval of 30 minutes
+ *      10,     // Allow 10 requests per window of 30 minutes
+ *    )
+ * })
+ *
+ * ```
+ */
+export class RegionRatelimit extends Ratelimit<RegionContext> {
+  /**
+   * Create a new Ratelimit instance by providing a `@upstash/redis` instance and the algorithm of your choice.
+   */
+
+  constructor(config: RegionRatelimitConfig) {
+    super({
+      prefix: config.prefix,
+      limiter: config.limiter,
+      timeout: config.timeout,
+      analytics: config.analytics,
+      ctx: {
+        redis: config.redis,
+      },
+      ephemeralCache: config.ephemeralCache,
+    });
+  }
+
+  /**
+   * Each request inside a fixed time increases a counter.
+   * Once the counter reaches the maximum allowed number, all further requests are
+   * rejected.
+   *
+   * **Pro:**
+   *
+   * - Newer requests are not starved by old ones.
+   * - Low storage cost.
+   *
+   * **Con:**
+   *
+   * A burst of requests near the boundary of a window can result in a very
+   * high request rate because two windows will be filled with requests quickly.
+   *
+   * @param tokens - How many requests a user can make in each time window.
+   * @param window - A fixed timeframe
+   */
+  static fixedWindow(
+    /**
+     * How many requests are allowed per window.
+     */
+    tokens: number,
+    /**
+     * The duration in which `tokens` requests are allowed.
+     */
+    window: Duration,
+  ): Algorithm<RegionContext> {
+    const windowDuration = ms(window);
+
+    const script = `
+    local key     = KEYS[1]
+    local window  = ARGV[1]
+    
+    local r = redis.call("INCR", key)
+    if r == 1 then 
+    -- The first time this key is set, the value will be 1.
+    -- So we only need the expire command once
+    redis.call("PEXPIRE", key, window)
+    end
+    
+    return r
+    `;
+
+    return async function (ctx: RegionContext, identifier: string) {
+      const bucket = Math.floor(Date.now() / windowDuration);
+      const key = [identifier, bucket].join(":");
+
+      if (ctx.cache) {
+        const { blocked, reset } = ctx.cache.isBlocked(identifier);
+        if (blocked) {
+          return {
+            success: false,
+            limit: tokens,
+            remaining: 0,
+            reset: reset,
+            pending: Promise.resolve(),
+          };
+        }
+      }
+      const usedTokensAfterUpdate = (await ctx.redis.eval(
+        script,
+        [key],
+        [windowDuration],
+      )) as number;
+
+      const success = usedTokensAfterUpdate <= tokens;
+      const reset = (bucket + 1) * windowDuration;
+      if (ctx.cache && !success) {
+        ctx.cache.blockUntil(identifier, reset);
+      }
+
+      return {
+        success,
+        limit: tokens,
+        remaining: Math.max(0, tokens - usedTokensAfterUpdate),
+        reset,
+        pending: Promise.resolve(),
+      };
+    };
+  }
+
+  /**
+   * Combined approach of `slidingLogs` and `fixedWindow` with lower storage
+   * costs than `slidingLogs` and improved boundary behavior by calculating a
+   * weighted score between two windows.
+   *
+   * **Pro:**
+   *
+   * Good performance allows this to scale to very high loads.
+   *
+   * **Con:**
+   *
+   * Nothing major.
+   *
+   * @param tokens - How many requests a user can make in each time window.
+   * @param window - The duration in which the user can max X requests.
+   */
+  static slidingWindow(
+    /**
+     * How many requests are allowed per window.
+     */
+    tokens: number,
+    /**
+     * The duration in which `tokens` requests are allowed.
+     */
+    window: Duration,
+  ): Algorithm<RegionContext> {
+    const script = `
+      local currentKey  = KEYS[1]           -- identifier including prefixes
+      local previousKey = KEYS[2]           -- key of the previous bucket
+      local tokens      = tonumber(ARGV[1]) -- tokens per window
+      local now         = ARGV[2]           -- current timestamp in milliseconds
+      local window      = ARGV[3]           -- interval in milliseconds
+
+      local requestsInCurrentWindow = redis.call("GET", currentKey)
+      if requestsInCurrentWindow == false then
+        requestsInCurrentWindow = 0
+      end
+
+      local requestsInPreviousWindow = redis.call("GET", previousKey)
+      if requestsInPreviousWindow == false then
+        requestsInPreviousWindow = 0
+      end
+      local percentageInCurrent = ( now % window ) / window
+      -- weighted requests to consider from the previous window
+      requestsInPreviousWindow = math.floor(( 1 - percentageInCurrent ) * requestsInPreviousWindow)
+      if requestsInPreviousWindow + requestsInCurrentWindow >= tokens then
+        return -1
+      end
+
+      local newValue = redis.call("INCR", currentKey)
+      if newValue == 1 then 
+        -- The first time this key is set, the value will be 1.
+        -- So we only need the expire command once
+        redis.call("PEXPIRE", currentKey, window * 2 + 1000) -- Enough time to overlap with a new window + 1 second
+      end
+      return tokens - ( newValue + requestsInPreviousWindow )
+      `;
+    const windowSize = ms(window);
+    return async function (ctx: RegionContext, identifier: string) {
+      const now = Date.now();
+
+      const currentWindow = Math.floor(now / windowSize);
+      const currentKey = [identifier, currentWindow].join(":");
+      const previousWindow = currentWindow - 1;
+      const previousKey = [identifier, previousWindow].join(":");
+
+      if (ctx.cache) {
+        const { blocked, reset } = ctx.cache.isBlocked(identifier);
+        if (blocked) {
+          return {
+            success: false,
+            limit: tokens,
+            remaining: 0,
+            reset: reset,
+            pending: Promise.resolve(),
+          };
+        }
+      }
+
+      const remaining = (await ctx.redis.eval(
+        script,
+        [currentKey, previousKey],
+        [tokens, now, windowSize],
+      )) as number;
+
+      const success = remaining >= 0;
+      const reset = (currentWindow + 1) * windowSize;
+      if (ctx.cache && !success) {
+        ctx.cache.blockUntil(identifier, reset);
+      }
+      return {
+        success,
+        limit: tokens,
+        remaining: Math.max(0, remaining),
+        reset,
+        pending: Promise.resolve(),
+      };
+    };
+  }
+
+  /**
+   * You have a bucket filled with `{maxTokens}` tokens that refills constantly
+   * at `{refillRate}` per `{interval}`.
+   * Every request will remove one token from the bucket and if there is no
+   * token to take, the request is rejected.
+   *
+   * **Pro:**
+   *
+   * - Bursts of requests are smoothed out and you can process them at a constant
+   * rate.
+   * - Allows to set a higher initial burst limit by setting `maxTokens` higher
+   * than `refillRate`
+   */
+  static tokenBucket(
+    /**
+     * How many tokens are refilled per `interval`
+     *
+     * An interval of `10s` and refillRate of 5 will cause a new token to be added every 2 seconds.
+     */
+    refillRate: number,
+    /**
+     * The interval for the `refillRate`
+     */
+    interval: Duration,
+    /**
+     * Maximum number of tokens.
+     * A newly created bucket starts with this many tokens.
+     * Useful to allow higher burst limits.
+     */
+    maxTokens: number,
+  ): Algorithm<RegionContext> {
+    const script = `
+        local key         = KEYS[1]           -- identifier including prefixes
+        local maxTokens   = tonumber(ARGV[1]) -- maximum number of tokens
+        local interval    = tonumber(ARGV[2]) -- size of the window in milliseconds
+        local refillRate  = tonumber(ARGV[3]) -- how many tokens are refilled after each interval
+        local now         = tonumber(ARGV[4]) -- current timestamp in milliseconds
+        local remaining   = 0
+        
+        local bucket = redis.call("HMGET", key, "updatedAt", "tokens")
+        
+        if bucket[1] == false then
+          -- The bucket does not exist yet, so we create it and add a ttl.
+          remaining = maxTokens - 1
+          
+          redis.call("HMSET", key, "updatedAt", now, "tokens", remaining)
+          redis.call("PEXPIRE", key, interval)
+
+          return {remaining, now + interval}
+        end
+
+        -- The bucket does exist
+    
+        local updatedAt = tonumber(bucket[1])
+        local tokens = tonumber(bucket[2])
+
+        if now >= updatedAt + interval then
+          if tokens <= 0 then 
+            -- No more tokens were left before the refill.
+            remaining = math.min(maxTokens, refillRate) - 1
+          else
+            remaining = math.min(maxTokens, tokens + refillRate) - 1
+          end
+        redis.call("HMSET", key, "updatedAt", now, "tokens", remaining)
+        return {remaining, now + interval}
+      end
+      
+      remaining = tokens - 1
+      redis.call("HSET", key, "tokens", remaining)
+      return {remaining, updatedAt + interval}
+
+    
+       `;
+
+    const intervalDuration = ms(interval);
+    return async function (ctx: RegionContext, identifier: string) {
+      if (ctx.cache) {
+        const { blocked, reset } = ctx.cache.isBlocked(identifier);
+        if (blocked) {
+          return {
+            success: false,
+            limit: maxTokens,
+            remaining: 0,
+            reset: reset,
+            pending: Promise.resolve(),
+          };
+        }
+      }
+
+      const now = Date.now();
+      const key = [identifier, Math.floor(now / intervalDuration)].join(":");
+
+      const [remaining, reset] = (await ctx.redis.eval(
+        script,
+        [key],
+        [maxTokens, intervalDuration, refillRate, now],
+      )) as [number, number];
+
+      const success = remaining > 0;
+      if (ctx.cache && !success) {
+        ctx.cache.blockUntil(identifier, reset);
+      }
+
+      return {
+        success,
+        limit: maxTokens,
+        remaining,
+        reset,
+        pending: Promise.resolve(),
+      };
+    };
+  }
+
+  /**
+   * cachedFixedWindow first uses the local cache to decide if a request may pass and then updates
+   * it asynchronously.
+   * This is experimental and not yet recommended for production use.
+   *
+   * @experimental
+   *
+   * Each request inside a fixed time increases a counter.
+   * Once the counter reaches the maximum allowed number, all further requests are
+   * rejected.
+   *
+   * **Pro:**
+   *
+   * - Newer requests are not starved by old ones.
+   * - Low storage cost.
+   *
+   * **Con:**
+   *
+   * A burst of requests near the boundary of a window can result in a very
+   * high request rate because two windows will be filled with requests quickly.
+   *
+   * @param tokens - How many requests a user can make in each time window.
+   * @param window - A fixed timeframe
+   */
+  static cachedFixedWindow(
+    /**
+     * How many requests are allowed per window.
+     */
+    tokens: number,
+    /**
+     * The duration in which `tokens` requests are allowed.
+     */
+    window: Duration,
+  ): Algorithm<RegionContext> {
+    const windowDuration = ms(window);
+
+    const script = `
+      local key     = KEYS[1]
+      local window  = ARGV[1]
+      
+      local r = redis.call("INCR", key)
+      if r == 1 then 
+      -- The first time this key is set, the value will be 1.
+      -- So we only need the expire command once
+      redis.call("PEXPIRE", key, window)
+      end
+      
+      return r
+      `;
+
+    return async function (ctx: RegionContext, identifier: string) {
+      if (!ctx.cache) {
+        throw new Error("This algorithm requires a cache");
+      }
+      const bucket = Math.floor(Date.now() / windowDuration);
+      const key = [identifier, bucket].join(":");
+      const reset = (bucket + 1) * windowDuration;
+
+      const hit = typeof ctx.cache.get(key) === "number";
+      if (hit) {
+        const cachedTokensAfterUpdate = ctx.cache.incr(key);
+        const success = cachedTokensAfterUpdate < tokens;
+
+        const pending = success
+          ? ctx.redis.eval(script, [key], [windowDuration]).then((t) => {
+              ctx.cache!.set(key, t as number);
+            })
+          : Promise.resolve();
+
+        return {
+          success,
+          limit: tokens,
+          remaining: tokens - cachedTokensAfterUpdate,
+          reset: reset,
+          pending,
+        };
+      }
+
+      const usedTokensAfterUpdate = (await ctx.redis.eval(
+        script,
+        [key],
+        [windowDuration],
+      )) as number;
+      ctx.cache.set(key, usedTokensAfterUpdate);
+      const remaining = tokens - usedTokensAfterUpdate;
+
+      return {
+        success: remaining >= 0,
+        limit: tokens,
+        remaining,
+        reset: reset,
+        pending: Promise.resolve(),
+      };
+    };
+  }
+}

package/src/test_utils.ts

@@ -0,0 +1,65 @@
+import crypto from "node:crypto";
+import { Ratelimit } from "./ratelimit";
+import { Context } from "./types";
+
+type Metrics = {
+  requests: number;
+  success: number;
+  rejected: number;
+};
+export class TestHarness<TContext extends Context> {
+  /**
+   * Used as prefix for redis keys
+   */
+  public readonly id: string;
+
+  private readonly ratelimit: Ratelimit<TContext>;
+  public metrics: Metrics;
+
+  public latencies: Record<string, { start: number; end: number }> = {};
+
+  constructor(ratelimit: Ratelimit<TContext>) {
+    this.ratelimit = ratelimit;
+    this.id = crypto.randomUUID();
+    this.metrics = {
+      requests: 0,
+      success: 0,
+      rejected: 0,
+    };
+  }
+
+  /**
+   * @param rps - req per second
+   * @param duration - duration in seconds
+   */
+  public async attack(rps: number, duration: number): Promise<void> {
+    const promises: Promise<{ success: boolean; pending: Promise<unknown> }>[] = [];
+
+    for (let i = 0; i < duration; i++) {
+      for (let r = 0; r < rps; r++) {
+        this.metrics.requests++;
+        const id = crypto.randomUUID();
+        this.latencies[id] = { start: Date.now(), end: -1 };
+        promises.push(
+          this.ratelimit.limit(this.id).then((res) => {
+            this.latencies[id].end = Date.now();
+            return res;
+          }),
+        );
+        await new Promise((r) => setTimeout(r, 1000 / rps));
+      }
+    }
+
+    await Promise.all(
+      promises.map(async (p) => {
+        const { success, pending } = await p;
+        await pending;
+        if (success) {
+          this.metrics.success++;
+        } else {
+          this.metrics.rejected++;
+        }
+      }),
+    );
+  }
+}

package/src/tools/seed.ts

@@ -0,0 +1,37 @@
+import { Redis } from "@upstash/redis";
+import { Analytics } from "../analytics";
+
+const redis = Redis.fromEnv();
+
+const identifier = new Set<string>();
+
+function getId(): string {
+  if (identifier.size === 0 || Math.random() > 0.95) {
+    const newIp = new Array(4)
+      .fill(0)
+      .map((_) => Math.floor(Math.random() * 256))
+      .join(".");
+    identifier.add(newIp);
+  }
+  return [...identifier][Math.floor(Math.random() * identifier.size)];
+}
+
+const a = new Analytics({ redis });
+
+async function main() {
+  const now = Date.now();
+  for (let i = 0; i < 1000; i++) {
+    console.log(i);
+    await Promise.all(
+      new Array(100).fill(0).map((_) =>
+        a.record({
+          time: now - Math.round(Math.random() * 7 * 24 * 60 * 60 * 1000),
+          identifier: getId(),
+          success: Math.random() > 0.2,
+        }),
+      ),
+    );
+  }
+}
+
+main();

package/src/types.ts

@@ -0,0 +1,78 @@
+/**
+ * EphemeralCache is used to block certain identifiers right away in case they have already exceeded the ratelimit.
+ */
+export interface EphemeralCache {
+  isBlocked: (identifier: string) => { blocked: boolean; reset: number };
+  blockUntil: (identifier: string, reset: number) => void;
+
+  set: (key: string, value: number) => void;
+  get: (key: string) => number | null;
+
+  incr: (key: string) => number;
+}
+
+export type RegionContext = { redis: Redis; cache?: EphemeralCache };
+export type MultiRegionContext = { redis: Redis[]; cache?: EphemeralCache };
+
+export type Context = RegionContext | MultiRegionContext;
+export type RatelimitResponse = {
+  /**
+   * Whether the request may pass(true) or exceeded the limit(false)
+   */
+  success: boolean;
+  /**
+   * Maximum number of requests allowed within a window.
+   */
+  limit: number;
+  /**
+   * How many requests the user has left within the current window.
+   */
+  remaining: number;
+  /**
+   * Unix timestamp in milliseconds when the limits are reset.
+   */
+  reset: number;
+
+  /**
+   * For the MultiRegion setup we do some synchronizing in the background, after returning the current limit.
+   * In most case you can simply ignore this.
+   *
+   * On Vercel Edge or Cloudflare workers, you need to explicitly handle the pending Promise like this:
+   *
+   * **Vercel Edge:**
+   * https://nextjs.org/docs/api-reference/next/server#nextfetchevent
+   *
+   * ```ts
+   * const { pending } = await ratelimit.limit("id")
+   * event.waitUntil(pending)
+   * ```
+   *
+   * **Cloudflare Worker:**
+   * https://developers.cloudflare.com/workers/runtime-apis/fetch-event/#syntax-module-worker
+   *
+   * ```ts
+   * const { pending } = await ratelimit.limit("id")
+   * context.waitUntil(pending)
+   * ```
+   */
+  pending: Promise<unknown>;
+};
+
+export type Algorithm<TContext> = (
+  ctx: TContext,
+  identifier: string,
+  opts?: {
+    cache?: EphemeralCache;
+  },
+) => Promise<RatelimitResponse>;
+
+/**
+ * This is all we need from the redis sdk.
+ */
+export interface Redis {
+  sadd: <TData>(key: string, ...members: TData[]) => Promise<number>;
+
+  eval: <TArgs extends unknown[], TData = unknown>(
+    ...args: [script: string, keys: string[], args: TArgs]
+  ) => Promise<TData>;
+}

package/src/version.ts

@@ -0,0 +1 @@
+export const VERSION = "0.4.5-canary.0";