@upstash/ratelimit 0.4.5-canary.0 → 1.0.0

package/src/multi.ts

@@ -1,365 +0,0 @@
-import { Cache } from "./cache";
-import type { Duration } from "./duration";
-import { ms } from "./duration";
-import { Ratelimit } from "./ratelimit";
-import type { Algorithm, MultiRegionContext } from "./types";
-import type { Redis } from "./types";
-
-function randomId(): string {
-  let result = "";
-  const characters = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789";
-  const charactersLength = characters.length;
-  for (let i = 0; i < 16; i++) {
-    result += characters.charAt(Math.floor(Math.random() * charactersLength));
-  }
-  return result;
-}
-
-export type MultiRegionRatelimitConfig = {
-  /**
-   * Instances 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:
-   * - MultiRegionRatelimit.fixedWindow
-   */
-  limiter: Algorithm<MultiRegionContext>;
-  /**
-   * 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 MultiRegionRatelimit({
- *    redis: Redis.fromEnv(),
- *    limiter: MultiRegionRatelimit.fixedWindow(
- *      10,     // Allow 10 requests per window of 30 minutes
- *      "30 m", // interval of 30 minutes
- *    )
- * })
- *
- * ```
- */
-export class MultiRegionRatelimit extends Ratelimit<MultiRegionContext> {
-  /**
-   * Create a new Ratelimit instance by providing a `@upstash/redis` instance and the algorithn of your choice.
-   */
-  constructor(config: MultiRegionRatelimitConfig) {
-    super({
-      prefix: config.prefix,
-      limiter: config.limiter,
-      timeout: config.timeout,
-      analytics: config.analytics,
-      ctx: {
-        redis: config.redis,
-        cache: config.ephemeralCache ? new Cache(config.ephemeralCache) : undefined,
-      },
-    });
-  }
-
-  /**
-   * 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<MultiRegionContext> {
-    const windowDuration = ms(window);
-    const script = `
-    local key     = KEYS[1]
-    local id      = ARGV[1]
-    local window  = ARGV[2]
-    
-    redis.call("SADD", key, id)
-    local members = redis.call("SMEMBERS", key)
-    if #members == 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 members
-`;
-
-    return async function (ctx: MultiRegionContext, identifier: string) {
-      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 requestId = randomId();
-      const bucket = Math.floor(Date.now() / windowDuration);
-      const key = [identifier, bucket].join(":");
-
-      const dbs: { redis: Redis; request: Promise<string[]> }[] = ctx.redis.map((redis) => ({
-        redis,
-        request: redis.eval(script, [key], [requestId, windowDuration]) as Promise<string[]>,
-      }));
-
-      const firstResponse = await Promise.any(dbs.map((s) => s.request));
-
-      const usedTokens = firstResponse.length;
-
-      const remaining = tokens - usedTokens - 1;
-
-      /**
-       * If the length between two databases does not match, we sync the two databases
-       */
-      async function sync() {
-        const individualIDs = await Promise.all(dbs.map((s) => s.request));
-        const allIDs = Array.from(new Set(individualIDs.flatMap((_) => _)).values());
-
-        for (const db of dbs) {
-          const ids = await db.request;
-          /**
-           * If the bucket in this db is already full, it doesn't matter which ids it contains.
-           * So we do not have to sync.
-           */
-          if (ids.length >= tokens) {
-            continue;
-          }
-          const diff = allIDs.filter((id) => !ids.includes(id));
-          /**
-           * Don't waste a request if there is nothing to send
-           */
-          if (diff.length === 0) {
-            continue;
-          }
-
-          await db.redis.sadd(key, ...allIDs);
-        }
-      }
-
-      /**
-       * Do not await sync. This should not run in the critical path.
-       */
-
-      const success = remaining > 0;
-      const reset = (bucket + 1) * windowDuration;
-
-      if (ctx.cache && !success) {
-        ctx.cache.blockUntil(identifier, reset);
-      }
-      return {
-        success,
-        limit: tokens,
-        remaining,
-        reset,
-        pending: sync(),
-      };
-    };
-  }
-
-  /**
-   * 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<MultiRegionContext> {
-    const windowSize = ms(window);
-    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 requestId   = ARGV[4]           -- uuid for this request
-
-
-      local currentMembers = redis.call("SMEMBERS", currentKey)
-      local requestsInCurrentWindow = #currentMembers
-      local previousMembers = redis.call("SMEMBERS", previousKey)
-      local requestsInPreviousWindow = #previousMembers
-
-      local percentageInCurrent = ( now % window) / window
-      if requestsInPreviousWindow * ( 1 - percentageInCurrent ) + requestsInCurrentWindow >= tokens then
-        return {currentMembers, previousMembers, false}
-      end
-
-      redis.call("SADD", currentKey, requestId)
-      table.insert(currentMembers, requestId)
-      if requestsInCurrentWindow == 0 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 {currentMembers, previousMembers, true}
-      `;
-    const windowDuration = ms(window);
-
-    return async function (ctx: MultiRegionContext, identifier: string) {
-      // 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 requestId = randomId();
-      const now = Date.now();
-
-      const currentWindow = Math.floor(now / windowSize);
-      const currentKey = [identifier, currentWindow].join(":");
-      const previousWindow = currentWindow - 1;
-      const previousKey = [identifier, previousWindow].join(":");
-
-      const dbs = ctx.redis.map((redis) => ({
-        redis,
-        request: redis.eval(
-          script,
-          [currentKey, previousKey],
-          [tokens, now, windowDuration, requestId],
-          // lua seems to return `1` for true and `null` for false
-        ) as Promise<[string[], string[], 1 | null]>,
-      }));
-
-      const percentageInCurrent = (now % windowDuration) / windowDuration;
-      const [current, previous, success] = await Promise.any(dbs.map((s) => s.request));
-
-      const previousPartialUsed = previous.length * (1 - percentageInCurrent);
-      const usedTokens = previousPartialUsed + current.length;
-
-      const remaining = tokens - usedTokens;
-
-      /**
-       * If a database differs from the consensus, we sync it
-       */
-      async function sync() {
-        const res = await Promise.all(dbs.map((s) => s.request));
-        const allCurrentIds = res.flatMap(([current]) => current);
-        for (const db of dbs) {
-          const [ids] = await db.request;
-          /**
-           * If the bucket in this db is already full, it doesn't matter which ids it contains.
-           * So we do not have to sync.
-           */
-          if (ids.length >= tokens) {
-            continue;
-          }
-          const diff = allCurrentIds.filter((id) => !ids.includes(id));
-          /**
-           * Don't waste a request if there is nothing to send
-           */
-          if (diff.length === 0) {
-            continue;
-          }
-
-          await db.redis.sadd(currentKey, ...diff);
-        }
-      }
-
-      // const success = remaining >= 0;
-      const reset = (currentWindow + 1) * windowDuration;
-      if (ctx.cache && !success) {
-        ctx.cache.blockUntil(identifier, reset);
-      }
-      return {
-        success: Boolean(success),
-        limit: tokens,
-        remaining,
-        reset,
-        pending: sync(),
-      };
-    };
-  }
-}

package/src/ratelimit.test.ts

@@ -1,155 +0,0 @@
-import { describe, expect, test } from "bun:test";
-import { log } from "console";
-import crypto from "node:crypto";
-import { Redis } from "@upstash/redis";
-import { Algorithm } from ".";
-import type { Duration } from "./duration";
-import { MultiRegionRatelimit } from "./multi";
-import { Ratelimit } from "./ratelimit";
-import { RegionRatelimit } from "./single";
-import { TestHarness } from "./test_utils";
-import type { Context, MultiRegionContext, RegionContext } from "./types";
-
-type TestCase = {
-  // requests per second
-  rps: number;
-  /**
-   * Multiplier for rate
-   *
-   * rate = 10, load = 0.5 -> attack rate will be 5
-   */
-  load: number;
-};
-const attackDuration = 60;
-const window = 5;
-const windowString: Duration = `${window} s`;
-
-const testcases: TestCase[] = [];
-
-for (const rps of [10, 100]) {
-  for (const load of [0.5, 1, 1.5]) {
-    testcases.push({ load, rps });
-  }
-}
-
-function run<TContext extends Context>(builder: (tc: TestCase) => Ratelimit<TContext>) {
-  for (const tc of testcases) {
-    const name = `${tc.rps.toString().padStart(4, " ")}/s - Load: ${(tc.load * 100)
-      .toString()
-      .padStart(3, " ")}% -> Sending ${(tc.rps * tc.load).toString().padStart(4, " ")}req/s`;
-    const ratelimit = builder(tc);
-
-    const isMultiRegion = ratelimit instanceof MultiRegionRatelimit;
-    const limits = {
-      lte: ((attackDuration * tc.rps) / window) * (isMultiRegion ? 1.5 : 1.2),
-      gte: ((attackDuration * tc.rps) / window) * (isMultiRegion ? 0.5 : 0.8),
-    };
-    describe(name, () => {
-      test(
-        `should be within ${limits.gte} - ${limits.lte}`,
-        async () => {
-          log(name);
-          const harness = new TestHarness(ratelimit);
-          await harness.attack(tc.rps * tc.load, attackDuration).catch((e) => {
-            console.error(e);
-          });
-          log(
-            "success:",
-            harness.metrics.success,
-            ", blocked:",
-            harness.metrics.rejected,
-            "out of:",
-            harness.metrics.requests,
-          );
-
-          expect(harness.metrics.success).toBeLessThanOrEqual(limits.lte);
-          expect(harness.metrics.success).toBeGreaterThanOrEqual(limits.gte);
-        },
-        attackDuration * 1000 * 2,
-      );
-    });
-  }
-}
-
-function newMultiRegion(limiter: Algorithm<MultiRegionContext>): Ratelimit<MultiRegionContext> {
-  function ensureEnv(key: string): string {
-    const value = process.env[key];
-    if (!value) {
-      throw new Error(`Environment variable ${key} not found`);
-    }
-    return value;
-  }
-
-  return new MultiRegionRatelimit({
-    prefix: crypto.randomUUID(),
-    redis: [
-      new Redis({
-        url: ensureEnv("EU2_UPSTASH_REDIS_REST_URL"),
-        token: ensureEnv("EU2_UPSTASH_REDIS_REST_TOKEN"),
-      }),
-      new Redis({
-        url: ensureEnv("APN_UPSTASH_REDIS_REST_URL"),
-        token: ensureEnv("APN_UPSTASH_REDIS_REST_TOKEN"),
-      }),
-      new Redis({
-        url: ensureEnv("US1_UPSTASH_REDIS_REST_URL"),
-        token: ensureEnv("US1_UPSTASH_REDIS_REST_TOKEN"),
-      }),
-    ],
-    limiter,
-  });
-}
-
-function newRegion(limiter: Algorithm<RegionContext>): Ratelimit<RegionContext> {
-  return new RegionRatelimit({
-    prefix: crypto.randomUUID(),
-    redis: Redis.fromEnv(),
-    limiter,
-  });
-}
-
-describe("timeout", () => {
-  test("pass after timeout", async () => {
-    const r = new RegionRatelimit({
-      prefix: crypto.randomUUID(),
-      // @ts-ignore - I just want to test the timeout
-      redis: {
-        ...Redis.fromEnv(),
-        eval: () => new Promise((r) => setTimeout(r, 2000)),
-      },
-      limiter: RegionRatelimit.fixedWindow(1, "1 s"),
-      timeout: 1000,
-    });
-    const start = Date.now();
-    const res = await r.limit("id");
-    const duration = Date.now() - start;
-    expect(res.success).toBe(true);
-    expect(res.limit).toBe(0);
-    expect(res.remaining).toBe(0);
-    expect(res.reset).toBe(0);
-    expect(duration).toBeGreaterThanOrEqual(900);
-    expect(duration).toBeLessThanOrEqual(1100);
-
-    // stop the test from leaking
-    await new Promise((r) => setTimeout(r, 5000));
-  }, 10000);
-});
-
-describe("fixedWindow", () => {
-  describe("region", () =>
-    run((tc) => newRegion(RegionRatelimit.fixedWindow(tc.rps, windowString))));
-
-  describe("multiRegion", () =>
-    run((tc) => newMultiRegion(MultiRegionRatelimit.fixedWindow(tc.rps, windowString))));
-});
-describe("slidingWindow", () => {
-  describe("region", () =>
-    run((tc) => newRegion(RegionRatelimit.slidingWindow(tc.rps, windowString))));
-  describe("multiRegion", () =>
-    run((tc) => newMultiRegion(MultiRegionRatelimit.slidingWindow(tc.rps, windowString))));
-});
-
-describe("tokenBucket", () => {
-  describe("region", () =>
-    run((tc) => newRegion(RegionRatelimit.tokenBucket(tc.rps, windowString, tc.rps))));
-});