halt-rate 0.1.1 → 0.2.0

package/dist/index.js

@@ -389,6 +389,11 @@ var LeakyBucket = class {
   }
 };
 
+// src/core/store.ts
+function isAtomicStore(store) {
+  return typeof store === "object" && store !== null && typeof store.evaluate === "function";
+}
+
 // src/core/limiter.ts
 var RateLimiter = class {
   constructor(options) {
@@ -433,6 +438,33 @@ var RateLimiter = class {
       return resp;
     }
     const storageKey = `halt:${policy.name}:${key}`;
+    if (isAtomicStore(this.store)) {
+      const span2 = this.otelTracer?.startSpan?.("halt.check", {
+        attributes: { policy: policy.name, key }
+      });
+      const decision2 = await this.store.evaluate({
+        key: storageKey,
+        algorithm: policy.algorithm,
+        limit: policy.limit,
+        window: policy.window,
+        burst: policy.burst,
+        cost: requestCost,
+        ttl: policy.window * 2
+      });
+      this.metricsRecorder?.(
+        "halt.request.checked",
+        { policy: policy.name, allowed: String(decision2.allowed) },
+        1
+      );
+      this.metricsRecorder?.(
+        decision2.allowed ? "halt.request.allowed" : "halt.request.blocked",
+        { policy: policy.name },
+        1
+      );
+      span2?.end?.();
+      return decision2;
+    }
+    const store = this.store;
     let algorithm = this.algorithmCache.get(policy.name);
     if (!algorithm) {
       if (policy.algorithm === "token_bucket" /* TOKEN_BUCKET */) {
@@ -450,7 +482,7 @@ var RateLimiter = class {
       this.algorithmCache.set(policy.name, algorithm);
     }
     const span = this.otelTracer?.startSpan?.("halt.check", { attributes: { policy: policy.name, key } });
-    const state = this.store.get(storageKey);
+    const state = store.get(storageKey);
     let decision;
     if (algorithm instanceof TokenBucket) {
       let tokens;
@@ -466,7 +498,7 @@ var RateLimiter = class {
       const result = algorithm.checkAndConsume(tokens, lastRefill, requestCost);
       decision = result.decision;
       const ttl = policy.window * 2;
-      this.store.set(storageKey, { tokens: result.newTokens, lastRefill: result.newLastRefill }, ttl);
+      store.set(storageKey, { tokens: result.newTokens, lastRefill: result.newLastRefill }, ttl);
     } else if (algorithm instanceof FixedWindow) {
       let count;
       let windowStart;
@@ -481,13 +513,13 @@ var RateLimiter = class {
       const result = algorithm.checkAndConsume(count, windowStart, requestCost);
       decision = result.decision;
       const ttl = policy.window * 2;
-      this.store.set(storageKey, { count: result.newCount, windowStart: result.newWindowStart }, ttl);
+      store.set(storageKey, { count: result.newCount, windowStart: result.newWindowStart }, ttl);
     } else if (algorithm instanceof SlidingWindow) {
       const buckets = state || algorithm.initialState();
       const result = algorithm.checkAndConsume(buckets, requestCost);
       decision = result.decision;
       const ttl = policy.window * 2;
-      this.store.set(storageKey, result.newBuckets, ttl);
+      store.set(storageKey, result.newBuckets, ttl);
     } else if (algorithm instanceof LeakyBucket) {
       let level;
       let lastLeak;
@@ -502,7 +534,7 @@ var RateLimiter = class {
       const result = algorithm.checkAndConsume(level, lastLeak, requestCost);
       decision = result.decision;
       const ttl = policy.window * 2;
-      this.store.set(storageKey, { level: result.newLevel, lastLeak: result.newLastLeak }, ttl);
+      store.set(storageKey, { level: result.newLevel, lastLeak: result.newLastLeak }, ttl);
     } else {
       throw new Error(`Algorithm ${typeof algorithm} not supported`);
     }
@@ -643,6 +675,272 @@ var InMemoryStore = class {
   }
 };
 
+// src/stores/redis-scripts.ts
+var NOW = `
+local t = redis.call('TIME')
+local now = tonumber(t[1]) + (tonumber(t[2]) / 1000000)
+`;
+var TOKEN_BUCKET_LUA = `
+${NOW}
+local limit = tonumber(ARGV[1])
+local window = tonumber(ARGV[2])
+local capacity = tonumber(ARGV[3])
+local cost = tonumber(ARGV[4])
+local ttl = tonumber(ARGV[5])
+local rate = limit / window
+
+local data = redis.call('HMGET', KEYS[1], 'tokens', 'ts')
+local tokens, last
+if data[1] then
+  tokens = tonumber(data[1])
+  last = tonumber(data[2])
+else
+  tokens = capacity
+  last = now
+end
+
+local elapsed = now - last
+if elapsed < 0 then elapsed = 0 end
+tokens = math.min(capacity, tokens + elapsed * rate)
+
+local needed = capacity - tokens
+local resetAt = math.floor(now + (needed / rate))
+
+if tokens >= cost then
+  tokens = tokens - cost
+  redis.call('HMSET', KEYS[1], 'tokens', tokens, 'ts', now)
+  redis.call('PEXPIRE', KEYS[1], ttl * 1000)
+  return {1, limit, math.floor(tokens), resetAt, -1}
+else
+  local deficit = cost - tokens
+  local retryAfter = math.floor(deficit / rate) + 1
+  redis.call('HMSET', KEYS[1], 'tokens', tokens, 'ts', now)
+  redis.call('PEXPIRE', KEYS[1], ttl * 1000)
+  return {0, limit, 0, resetAt, retryAfter}
+end
+`;
+var FIXED_WINDOW_LUA = `
+${NOW}
+local limit = tonumber(ARGV[1])
+local window = tonumber(ARGV[2])
+local cost = tonumber(ARGV[4])
+local ttl = tonumber(ARGV[5])
+
+local data = redis.call('HMGET', KEYS[1], 'count', 'start')
+local count, start
+if data[1] then
+  count = tonumber(data[1])
+  start = tonumber(data[2])
+else
+  count = 0
+  start = now
+end
+
+if (now - start) >= window then
+  count = 0
+  start = now
+end
+
+local resetAt = math.floor(start + window)
+
+if (count + cost) <= limit then
+  count = count + cost
+  redis.call('HMSET', KEYS[1], 'count', count, 'start', start)
+  redis.call('PEXPIRE', KEYS[1], ttl * 1000)
+  return {1, limit, limit - count, resetAt, -1}
+else
+  redis.call('HMSET', KEYS[1], 'count', count, 'start', start)
+  redis.call('PEXPIRE', KEYS[1], ttl * 1000)
+  local retryAfter = math.floor(resetAt - now) + 1
+  return {0, limit, 0, resetAt, retryAfter}
+end
+`;
+var LEAKY_BUCKET_LUA = `
+${NOW}
+local limit = tonumber(ARGV[1])
+local window = tonumber(ARGV[2])
+local capacity = tonumber(ARGV[3])
+local cost = tonumber(ARGV[4])
+local ttl = tonumber(ARGV[5])
+local leakRate = limit / window
+
+local data = redis.call('HMGET', KEYS[1], 'level', 'ts')
+local level, last
+if data[1] then
+  level = tonumber(data[1])
+  last = tonumber(data[2])
+else
+  level = 0
+  last = now
+end
+
+local elapsed = now - last
+if elapsed < 0 then elapsed = 0 end
+level = math.max(0, level - (elapsed * leakRate))
+
+local resetAt
+if level > 0 then
+  resetAt = math.floor(now + (level / leakRate))
+else
+  resetAt = math.floor(now)
+end
+
+if (level + cost) <= capacity then
+  level = level + cost
+  redis.call('HMSET', KEYS[1], 'level', level, 'ts', now)
+  redis.call('PEXPIRE', KEYS[1], ttl * 1000)
+  return {1, capacity, math.floor(capacity - level), resetAt, -1}
+else
+  local spaceNeeded = level + cost - capacity
+  local retryAfter = math.floor(spaceNeeded / leakRate) + 1
+  redis.call('HMSET', KEYS[1], 'level', level, 'ts', now)
+  redis.call('PEXPIRE', KEYS[1], ttl * 1000)
+  return {0, capacity, 0, resetAt, retryAfter}
+end
+`;
+var SLIDING_WINDOW_LUA = `
+${NOW}
+local limit = tonumber(ARGV[1])
+local window = tonumber(ARGV[2])
+local cost = tonumber(ARGV[4])
+local ttl = tonumber(ARGV[5])
+
+-- Drop entries older than the sliding window.
+local windowStart = now - window
+redis.call('ZREMRANGEBYSCORE', KEYS[1], '-inf', windowStart)
+
+local count = redis.call('ZCARD', KEYS[1])
+
+-- Oldest remaining entry determines when capacity frees up.
+local oldest = redis.call('ZRANGE', KEYS[1], 0, 0, 'WITHSCORES')
+local resetAt
+if oldest[2] then
+  resetAt = math.floor(tonumber(oldest[2]) + window)
+else
+  resetAt = math.floor(now + window)
+end
+
+if (count + cost) <= limit then
+  -- Unique members so concurrent requests in the same instant don't collide.
+  math.randomseed(math.floor(now * 1000000))
+  for i = 1, cost do
+    local member = string.format('%d-%d-%d', math.floor(now * 1000000), i, math.random(1, 1000000000))
+    redis.call('ZADD', KEYS[1], now, member)
+  end
+  redis.call('PEXPIRE', KEYS[1], ttl * 1000)
+  return {1, limit, limit - (count + cost), resetAt, -1}
+else
+  redis.call('PEXPIRE', KEYS[1], ttl * 1000)
+  local retryAfter
+  if oldest[2] then
+    retryAfter = math.floor(tonumber(oldest[2]) + window - now) + 1
+  else
+    retryAfter = math.floor(window) + 1
+  end
+  return {0, limit, 0, resetAt, retryAfter}
+end
+`;
+var SCRIPTS = {
+  token_bucket: TOKEN_BUCKET_LUA,
+  fixed_window: FIXED_WINDOW_LUA,
+  sliding_window: SLIDING_WINDOW_LUA,
+  leaky_bucket: LEAKY_BUCKET_LUA
+};
+
+// src/stores/redis.ts
+var RedisStore = class {
+  constructor(options) {
+    /** Cache of script source -> SHA loaded into Redis, to use EVALSHA. */
+    this.shaCache = /* @__PURE__ */ new Map();
+    this.client = options.client;
+    this.failMode = options.failMode ?? "open";
+    this.onError = options.onError;
+    this.metricsRecorder = options.metricsRecorder;
+  }
+  async evaluate(input) {
+    const script = SCRIPTS[input.algorithm];
+    if (!script) {
+      throw new Error(`Algorithm ${input.algorithm} not supported by RedisStore`);
+    }
+    const args = [
+      input.limit,
+      input.window,
+      input.burst,
+      input.cost,
+      input.ttl
+    ];
+    try {
+      const raw = await this.run(script, input.key, args);
+      return this.toDecision(raw, input);
+    } catch (err) {
+      this.onError?.(err);
+      this.metricsRecorder?.("halt.redis.error", { algorithm: input.algorithm }, 1);
+      return this.failDecision(input);
+    }
+  }
+  /** Run a script via EVALSHA, falling back to EVAL (and caching the SHA). */
+  async run(script, key, args) {
+    const sha = this.shaCache.get(script);
+    if (sha && this.client.evalsha) {
+      try {
+        return await this.client.evalsha(sha, 1, key, ...args);
+      } catch (err) {
+        if (!isNoScriptError(err)) throw err;
+      }
+    }
+    const result = await this.client.eval(script, 1, key, ...args);
+    if (!this.shaCache.has(script) && this.client.script) {
+      try {
+        const loaded = await this.client.script("LOAD", script);
+        if (typeof loaded === "string") this.shaCache.set(script, loaded);
+      } catch {
+      }
+    }
+    return result;
+  }
+  /** Map the Lua array [allowed, limit, remaining, resetAt, retryAfter] to a Decision. */
+  toDecision(raw, input) {
+    if (!Array.isArray(raw) || raw.length < 4) {
+      throw new Error(`Unexpected Redis script result: ${JSON.stringify(raw)}`);
+    }
+    const [allowed, limit, remaining, resetAt, retryAfter] = raw.map((v) => Number(v));
+    const decision = {
+      allowed: allowed === 1,
+      limit,
+      remaining,
+      resetAt
+    };
+    if (!decision.allowed && retryAfter !== void 0 && retryAfter >= 0) {
+      decision.retryAfter = retryAfter;
+    }
+    this.metricsRecorder?.(
+      "halt.request.checked",
+      { algorithm: input.algorithm, allowed: String(decision.allowed) },
+      1
+    );
+    return decision;
+  }
+  /** Decision returned when Redis is unreachable, per failMode. */
+  failDecision(input) {
+    const resetAt = Math.floor(Date.now() / 1e3 + input.window);
+    if (this.failMode === "open") {
+      this.metricsRecorder?.("halt.request.fail_open", { algorithm: input.algorithm }, 1);
+      return { allowed: true, limit: input.limit, remaining: input.limit, resetAt };
+    }
+    this.metricsRecorder?.("halt.request.fail_closed", { algorithm: input.algorithm }, 1);
+    return {
+      allowed: false,
+      limit: input.limit,
+      remaining: 0,
+      resetAt,
+      retryAfter: input.window
+    };
+  }
+};
+function isNoScriptError(err) {
+  return typeof err === "object" && err !== null && String(err.message ?? "").includes("NOSCRIPT");
+}
+
 // src/presets/index.ts
 var presets_exports = {};
 __export(presets_exports, {
@@ -774,7 +1072,9 @@ exports.Algorithm = Algorithm;
 exports.InMemoryStore = InMemoryStore;
 exports.KeyStrategy = KeyStrategy;
 exports.RateLimiter = RateLimiter;
+exports.RedisStore = RedisStore;
 exports.extractors = extractors_exports;
+exports.isAtomicStore = isAtomicStore;
 exports.normalizePolicy = normalizePolicy;
 exports.presets = presets_exports;
 exports.toHeaders = toHeaders;