co-limit 0.1.2 → 0.1.4

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 5f54e6c9ae0b841090b1fe2eae972242c4a416eb504de4940cc807a500f7d27a
-  data.tar.gz: ebc7e45c7c29b425a4391230173f8db3a44c288482a917f2ef8e57be75552a31
+  metadata.gz: 6485d06837c5fb0b4b0fad6a0d8d4cdaede3b73dbd056b89af116e39f0be121a
+  data.tar.gz: 3d237294612352598260a83054ee370d7dacd297543f257189c1e0d18ea84771
 SHA512:
-  metadata.gz: 8535547535b239b227de6f78dfa0c37380e33c77f74db7c85ad91792068847b44f9074b02b3e034f9d94c8c89d517ac884ab18f111aaaca8bdeee02fc774b2ba
-  data.tar.gz: b37ee6f384eed0728a8b4a29b3271475bbe42607c0d879471081947ec4d2c34ef7d14be28f95b305094cf2cb8c548c6aa3a5a1b77fdc5fbf0eac52711fe0ac53
+  metadata.gz: c30142a96da7574db4946c78a37d3716d5f69e4fd5c73a86e6bfffddb30f8c1bacf67206a7b22159c94c5322c5f3be821044a4c1c43ce94051fe9c3f25174611
+  data.tar.gz: 02c49b15b93c06208221ff6d4013c31dfe0053f9326c0e22f3d04c187906d17627b1c72e20ee7bea78f7ebedc169ac4074759aa9342c5eff8fef824dcef62d5e

data/README.md

@@ -1,9 +1,9 @@
 
 # Limit Gem
 
-Gem that provides flexible, Redis-backed rate limiting utilities. It supports both Fixed Window and Rolling Window (Sliding Log) strategies, to easily control the number of allowed requests for a given identifier within a time window.
+Gem that provides flexible, Redis-backed rate limiting utilities. It supports Fixed Window, Rolling Window (Sliding Log), and Token Bucket strategies to control the number of allowed requests for a given identifier within a time window or with a steady refill rate.
 
-You can define rate-limiting rules dynamically using a lambda, and configure Redis via environment variables (REDIS_HOST, REDIS_PORT, REDIS_PASSWORD) or by passing connection details directly.
+You can define rate-limiting rules dynamically using a lambda and configure Redis via environment variables (REDIS_HOST, REDIS_PORT, REDIS_PASSWORD) or bypassing connection details directly.
 
 This gem is ideal for APIs, background jobs, or any system that needs simple, efficient throttling logic.
 
@@ -21,11 +21,23 @@ If you are not using Bundler, you can install the gem directly by running:
 $ gem install co-limit
 ```
 
+### Supported Rate Limiters:
+
+- **Fixed Window Rate Limiter**:
+  Allows a specified number of requests within a fixed time window. This method can cause burst traffic as it doesn't account for requests made outside of the window until it resets.
+
+- **Rolling Window Rate Limiter**:
+  Uses a sliding window mechanism, where only requests made within the last `n` seconds are counted. It provides more consistent traffic flow but can be more resource-intensive.
+
+- **Token Bucket Rate Limiter**:
+  Implements the classic token bucket algorithm. A bucket with fixed capacity refills tokens at a steady rate, allowing short bursts up to the bucket capacity while enforcing long‑term average throughput.
+
 ## Usage
 
-### Example Usage
+### Examples
 
-Here's an example of how to use the rate limiter in your application:
+#### Rolling Window Example
+Here's an example of how to use the rolling window rate limiter in your application:
 
 ```ruby
 sync_limit_calculator = lambda do |key| 
@@ -56,6 +68,91 @@ success_count += 1 if allowed
 puts "Success count: #{success_count}"  # Expected to be 11
 ```
 
+#### Fixed Window Example
+
+```ruby
+# Map keys to a fixed-window limit of N requests per M seconds
+SITE_LIMITS = {
+  x: { max_requests: 10, window_seconds: 5 },   # 10 requests per 5 seconds
+  default: { max_requests: 10, window_seconds: 60 }
+}.freeze
+
+calc = lambda do |key|
+  plan = key.split(":").last.to_sym rescue :default
+  SITE_LIMITS.fetch(plan, SITE_LIMITS[:default])
+end
+
+limiter = Limit::FixedWindowRateLimiter.new(
+  identifier_prefix: "access",
+  limit_calculator: calc
+)
+
+key = "user123:x"  # => will map to { max_requests: 10, window_seconds: 5 }
+
+# Consume up to the window limit
+allowed = 0
+12.times { allowed += 1 if limiter.allowed?(key) }
+puts allowed  # => 10
+
+# Wait until the next aligned fixed window begins to reset the counter
+window = 5
+now = Time.now
+sleep(window - (now.to_i % window) + 0.5)
+
+puts limiter.allowed?(key)  # => true (window reset)
+```
+
+#### Token Bucket Example
+
+```ruby
+# Choose a plan via the key suffix, then map it to a token bucket config
+PLANS = {
+  basic:   { bucket_capacity: 3, refill_rate: 1, refill_interval: 1 },  # 3 burst, ~1 req/s
+  premium: { bucket_capacity: 10, refill_rate: 5, refill_interval: 1 }  # 10 burst, ~5 req/s
+}.freeze
+
+calc = lambda do |key|
+  plan = key.split(":").last.to_sym rescue :basic
+  PLANS.fetch(plan, PLANS[:basic])
+end
+
+limiter = Limit::TokenBucketRateLimiter.new(
+  identifier_prefix: "access",   # identifier_prefix is used to namespace keys in Redis (consistent across all limiters)
+  limit_calculator: calc
+)
+
+key = "user123:basic"            # the final Redis key will be "access:user123:basic" via identifier_prefix
+
+# Use up to capacity immediately
+3.times { puts limiter.allowed?(key) }  # => true, true, true
+puts limiter.allowed?(key)              # => false (bucket empty)
+
+sleep 1.2                               # wait ~1s to refill ~1 token
+puts limiter.allowed?(key)              # => true
+```
+
+### Key Points:
+
+- **identifier_prefix**: A namespace prefix for Redis keys (e.g., `"access"`).
+- **limit_calculator**: A `Lambda` that takes a key (e.g., `"user_id:site_name"`) and returns a hash with `max_requests` and `window_seconds`.
+
+### Limit Hash Structure by Limiter
+
+Your `limit_calculator` lambda must return a hash whose expected keys depend on the limiter you use:
+
+- FixedWindowRateLimiter and RollingWindowRateLimiter expect:
+  `{ max_requests: Integer, window_seconds: Integer }`
+
+- TokenBucketRateLimiter expects:
+  `{ bucket_capacity: Integer, refill_rate: Integer, refill_interval: Integer }`
+
+  Meaning: `refill_rate` tokens are added every `refill_interval` seconds (e.g., `refill_rate: 1, refill_interval: 1` = 1 token/second; `refill_rate: 3, refill_interval: 2` = 1.5 tokens/second effective rate). All values must be positive integers.
+
+Notes:
+- Redis keys are stored as "#{identifier_prefix}:#{key}". Choose concise, stable keys; e.g., identifier_prefix: "access" and key: "user_id:plan" → "access:user_id:plan".
+- Ensure the keys you pass are unique per entity you want to limit (e.g., "access:user_id:plan").
+- For Token Bucket, the Redis TTL is kept around the refill interval so buckets expire when inactive.
+
 ### Redis Configuration
 
 You can configure the Redis connection either by passing the connection details as arguments or by setting environment variables.
@@ -82,19 +179,6 @@ You can configure the Redis connection either by passing the connection details
 
   In this case, the gem will use these environment variables to establish the Redis connection.
 
-### Key Points:
-
-- **identifier_prefix**: A namespace prefix for Redis keys (e.g., `"access"`).
-- **limit_calculator**: A `Lambda` that takes a key (e.g., `"user_id:site_name"`) and returns a hash with `max_requests` and `window_seconds`.
-
-### Supported Rate Limiters:
-
-- **Fixed Window Rate Limiter**:
-  Allows a specified number of requests within a fixed time window. This method can cause burst traffic as it doesn't account for requests made outside of the window until it resets.
-
-- **Rolling Window Rate Limiter**:
-  Uses a sliding window mechanism, where only requests made within the last `n` seconds are counted. It provides more consistent traffic flow but can be more resource-intensive.
-
 ## Development
 
 After checking out the repo, install the dependencies by running:

data/lib/limit/version.rb

@@ -1,5 +1,5 @@
 # frozen_string_literal: true
 
 module Limit
-  VERSION = '0.1.2'
+  VERSION = '0.1.4'
 end

data/lib/limit.rb

@@ -9,12 +9,14 @@ module Limit
   # Base Class, implementing:
   # - Method to connect with redis
   # - Signature for limit_calculator
-  class BaseLimiter
+  class BaseRateLimiter
     attr_reader :limit_calculator, :identifier_prefix
 
     @redis = nil
     @logger = Logger.new($stdout)
 
+    REQUIRED_KEYS = %i[max_requests window_seconds].freeze
+
     class << self
 
       def connection
@@ -28,6 +30,14 @@ module Limit
         @logger ||= Logger.new($stdout)
       end
 
+      def load_script
+        script_name = name.split('::')[1].chomp('RateLimiter')
+        script_path = File.expand_path("scripts/#{script_name}.lua", __dir__)
+        @script_sha = @redis.script(:load, File.read(script_path))
+      end
+
+      attr_reader :script_sha
+
       private
 
       def create_connection(host:, port:, password:)
@@ -44,7 +54,8 @@ module Limit
     def initialize(identifier_prefix:, limit_calculator:, host: nil, port: nil, password: nil)
 
       # @param identifier_prefix: [String] A namespace prefix for redis keys for this limiter instance
-      # @param limit_calculator: [Lambda] A method that takes a key(String) and returns hash: {max_requests: Integer, window_seconds: Integer}
+      # @param limit_calculator: Lambda that takes a key(String) and returns hash: {max_requests: Integer, window_seconds: Integer}
+      # But diff implementation can update the hash struct
 
       unless identifier_prefix.is_a?(String) && !identifier_prefix.empty?
         raise ArgumentError, 'identifier_prefix must be a non-empty String'
@@ -52,7 +63,7 @@ module Limit
 
       raise ArgumentError, 'limit_calculator must be a Lambda' unless limit_calculator.lambda?
 
-      # Will be using the same connection across all instance unless wanted to connect to diff instance of redis
+      # Will be using the same connection across all instances unless wanted to connect to diff instance of redis
 
       @redis = if host && port && password
                  self.class.send(:create_connection, host: host, port: port, password: password)
@@ -72,6 +83,7 @@ module Limit
         raise e
       end
 
+      self.class.load_script
     end
 
     def allowed?(key)
@@ -79,37 +91,54 @@ module Limit
     end
 
     def get_key(key)
-      raise NotImplementedError "#{self.class.name} must implement the get_key() method"
+      "#{@identifier_prefix}:#{key}"
     end
 
 
     protected
 
+    def required_keys
+      self.class::REQUIRED_KEYS
+    end
+
     def get_current_limit(key)
       limit_data = @limit_calculator.call(key)
 
-      unless limit_data.is_a?(Hash) && limit_data[:max_requests].is_a?(Integer) && limit_data[:max_requests].positive? &&
-             limit_data[:window_seconds].is_a?(Integer) && limit_data[:window_seconds].positive?
-
-        raise ArgumentError, "Limit calculator for key '#{key}' returned invalid data: #{limit_data.inspect}. Expected { max_requests: Integer > 0, window_seconds: Integer > 0 }"
+      required_keys.all? do |key|
+        unless limit_data[key].is_a?(Integer) && limit_data[key].positive?
+          raise ArgumentError, "Limit calculator for key '#{key}' returned invalid data: #{limit_data.inspect}.\n
+          Expected #{required_keys}"
+        end
       end
 
       limit_data
     end
 
-    def redis_pipeline(&block)
+    def eval_sha(key, argv)
+      attempts ||= 0
       begin
-        @redis.pipelined { |pipe| block.call(pipe) }
+        @redis.evalsha(self.class.script_sha, key, argv)
+      rescue Redis::CannotConnectError, Redis::TimeoutError, Redis::ConnectionError => e
+        if attempts < 3
+          @logger.warn(e.message)
+          @redis = self.class.connection
+          attempts += 1
+          retry
+        else
+          @logger.error("Error connecting to Redis: #{e.message}")
+        end
       rescue Redis::CommandError => e
+        if e.message.start_with?("NOSCRIPT")
+          self.class.load_script
+          retry
+        else
+          @logger.error(e.message)
+        end
+      rescue Redis::BaseError => e
         @logger.error(e.message)
       end
     end
 
-    def current_micros
-      (Time.now.to_f * 1_000_000).to_i
-    end
-
-
   end
 
 
@@ -117,60 +146,56 @@ module Limit
 
   # Fixed Window Rate Limiter, allows n of request in a fixed window
   # ALERT: There is a chance of bursts/spike in this method, so use it with caution
-  class FixedWindowRateLimiter < BaseLimiter
+  class FixedWindowRateLimiter < BaseRateLimiter
     def allowed?(key)
       limit_data = get_current_limit(key)
-      max_requests = limit_data[:max_requests]
-      window_seconds = limit_data[:window_seconds]
-
-      window_key = get_key(key, window_seconds)
-
-      results = redis_pipeline do |pipe|
-        # This is for simplicity as incr handles both creation and incrementing, rather than waiting on some read
-        # using the pipeline, the whole operation would also be atomic, as redis is single threaded and both queries are send in one trip
-        # https://redis.io/docs/latest/develop/use/pipelining/
-        pipe.incr(window_key)
-        pipe.expire(window_key, window_seconds)
-      end
-
-      results[0] <= max_requests
-    end
-
-    def get_key(key, window_seconds)
-      time_window = (Time.now.to_i / window_seconds) * window_seconds
-      "#{@identifier_prefix}:#{key}:#{time_window.to_s}"
+      result = eval_sha(
+        [get_key(key)],
+        [
+          limit_data[:window_seconds],
+          limit_data[:max_requests]
+        ]
+      )
+      result == 1
     end
   end
 
-  # RollingWindow Rate limiter implemented using Sliding Log, allows n no of requests in rolling window
-  class RollingWindowRateLimiter < BaseLimiter
+  # RollingWindow Rate limiter, implemented using Sliding Log, allows n no of requests in a rolling window
+  class RollingWindowRateLimiter < BaseRateLimiter
     def allowed?(key)
       limit_data = get_current_limit(key)
-      max_requests = limit_data[:max_requests]
-      window_seconds = limit_data[:window_seconds]
-
-      set_key = get_key(key)
+      result = eval_sha(
+        [get_key(key)],
+        [
+          limit_data[:window_seconds],
+          limit_data[:max_requests]
+        ]
+      )
+      result == 1
+    end
+  end
 
-      curr_micros = current_micros
+  # TokenBucketRateLimiter implements the token bucket algorithm that'll allow the steady refill rate
+  class TokenBucketRateLimiter < BaseRateLimiter
 
-      window_start_micros = curr_micros - (window_seconds*1_000_000)
+    # limit_calculator: Lambda that takes a key(String) and returns hash:
+    # {bucket_capacity: Integer, refill_rate: Integer, refill_interval: Integer}
 
-      results = redis_pipeline do |pipe|
-        # uses sorted set
-        # https://redis.io/glossary/redis-sorted-sets/
-        pipe.zremrangebyscore(set_key, 0, window_start_micros)
-        pipe.zadd(set_key, curr_micros, curr_micros.to_s)
-        pipe.zcard(set_key)
-        pipe.expire(set_key, window_seconds)
-      end
+    REQUIRED_KEYS = %i[bucket_capacity refill_rate refill_interval].freeze
 
-      results[2] <= max_requests
+    def allowed?(key)
+      limit_data = get_current_limit(key)
+      result = eval_sha(
+        [get_key(key)],
+        [
+          limit_data[:bucket_capacity],
+          limit_data[:refill_rate],
+          limit_data[:refill_interval]
+        ]
+      )
+      result == 1
     end
 
-    def get_key(key)
-      "#{@identifier_prefix}:#{key}"
-    end
   end
 
-
 end

data/lib/scripts/FixedWindow.lua

@@ -0,0 +1,21 @@
+local key = KEYS[1]
+local window_seconds = tonumber(ARGV[1])
+local max_requests = tonumber(ARGV[2])
+
+local curr_time = redis.call('TIME')
+local curr_sec = tonumber(curr_time[1])
+
+local time_window = math.floor(curr_sec / window_seconds) * window_seconds
+
+local window_key = key .. ":" .. tostring(time_window)
+
+-- This is for simplicity as incr handles both creation and incrementing, rather than waiting on some read
+local count = redis.call('incr', window_key)
+
+redis.call('expire', window_key, window_seconds)
+
+if count <= max_requests then
+    return 1
+else
+    return 0
+end

data/lib/scripts/RollingWindow.lua

@@ -0,0 +1,28 @@
+local key = KEYS[1]
+local window_seconds = tonumber(ARGV[1])
+local max_requests = tonumber(ARGV[2])
+
+local time_data = redis.call('TIME')
+local curr_micros = tonumber(time_data[1]) * 1000000 + tonumber(time_data[2])
+local window_start_micros = curr_micros - (window_seconds * 1000000)
+
+-- uses sorted set
+-- https://redis.io/glossary/redis-sorted-sets/
+-- Remove old entries
+redis.call('ZREMRANGEBYSCORE', key, 0, window_start_micros)
+
+-- Add current request
+redis.call('ZADD', key, curr_micros, tostring(curr_micros))
+
+-- Count requests
+local count = redis.call('ZCARD', key)
+
+-- Set expiry
+redis.call('EXPIRE', key, window_seconds)
+
+if count <= max_requests then
+    return 1
+else
+    return 0
+end
+

data/lib/scripts/TokenBucket.lua

@@ -0,0 +1,23 @@
+local key = KEYS[1]
+local bucket_capacity = tonumber(ARGV[1])
+local refill_rate = tonumber(ARGV[2])
+local refill_interval = tonumber(ARGV[3])
+
+local time_data = redis.call('TIME')
+local current_time = tonumber(time_data[1]) * 1000000 + tonumber(time_data[2])
+
+local bucket_data = redis.call('HMGET', key, 'tokens', 'last_updated')
+local tokens = tonumber(bucket_data[1]) or bucket_capacity
+local last_updated = tonumber(bucket_data[2]) or current_time
+
+local elapsed_micros = current_time - last_updated
+local elapsed_seconds = elapsed_micros / 1000000
+local tokens_to_add = math.floor(elapsed_seconds * refill_rate / refill_interval)
+local new_tokens = math.min(bucket_capacity, tokens + tokens_to_add)
+
+if new_tokens > 0 then
+redis.call('HMSET', key, 'tokens', new_tokens - 1, 'last_updated', current_time)
+redis.call('EXPIRE', key, refill_interval * 2)
+return 1
+end
+return 0

data/sig/limit/limit.rbs

@@ -2,9 +2,11 @@ type redis = untyped
 
 module Limit
 
-  class BaseLimiter
+  class BaseRateLimiter
+    REQUIRED_KEYS: Array[Symbol]
     self.@logger: Logger
     self.@redis: redis
+    self.@script_sha: untyped
     @logger: Logger
 
     @redis: redis | nil
@@ -17,6 +19,10 @@ module Limit
 
     def self.logger: -> Logger
 
+    def self.load_script: -> String
+
+    def self.script_sha: -> String
+
     attr_reader identifier_prefix: String
     attr_reader limit_calculator: Proc
 
@@ -32,6 +38,8 @@ module Limit
 
     def log: (String, String) -> nil
 
-    def redis_pipeline: () { (redis) -> untyped } -> untyped
-end
+    def eval_sha: (Array[String], Array[Integer]) -> bool
+
+    def required_keys: -> Array[Symbol]
+  end
 end

metadata

@@ -1,14 +1,14 @@
 --- !ruby/object:Gem::Specification
 name: co-limit
 version: !ruby/object:Gem::Version
-  version: 0.1.2
+  version: 0.1.4
 platform: ruby
 authors:
 - CosmicOppai
 autorequire:
 bindir: exe
 cert_chain: []
-date: 2025-05-30 00:00:00.000000000 Z
+date: 2025-09-21 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: redis
@@ -39,6 +39,9 @@ files:
 - Steepfile
 - lib/limit.rb
 - lib/limit/version.rb
+- lib/scripts/FixedWindow.lua
+- lib/scripts/RollingWindow.lua
+- lib/scripts/TokenBucket.lua
 - sig/limit/limit.rbs
 homepage: https://github.com/cosmicoppai/limit
 licenses: