pecorino 0.5.0 → 0.7.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 2cd65abe4917de817a0b0b08672376d5a7dd1e16c01b3a4e0c47ff1467600a1e
-  data.tar.gz: 58ca1578813b7a5bcc058d9a37a869e778a1614b736becced4bc82f78efba2c5
+  metadata.gz: 4f002108b01cca0f3d6a3315fe63f56ba5801916fa935a2aa60f608da1f54eee
+  data.tar.gz: e463591f7289505c6ba321d2db5fe7dd6c746036acac617515ab99dba011f88f
 SHA512:
-  metadata.gz: bbbcdc936bef119b1b02695dc626f8421576a619414f4493a541876562ca3cf5136dcdd1a796c9c08cd6cc2a8345be8ac7fecd8f6371fddc9b480bd26d017296
-  data.tar.gz: cd4e2ba40164eb0c9f56e17aeef656727703cb98f7e70cedaafa3d9af9202195f9c0e1a1dbc3dfdb411f277f31e84a6a2641d5f81bf85d6d07b3b92937877cbd
+  metadata.gz: 6847b163ed5857db1c04659e04d8698c734050b1cf9547dea52d7122e924e02663e68466682c7864ce115c54b1bceda6563c716b857437fab0a2463a223f7c02
+  data.tar.gz: daca5b4118c52b9df77cab06a4fa1d6a7ce4d3678da0cc962c90c3955d3b28c1b586536fc407811f9ede89807ab144d159eeaa50830bd34023a7bceaf3404331

data/.github/workflows/ci.yml

@@ -2,37 +2,11 @@ name: CI
 
 on:
   - push
-  - pull_request
 
 env:
   BUNDLE_PATH: vendor/bundle
 
 jobs:
-  # lint:
-  #   name: Code Style
-  #   runs-on: ubuntu-22.04
-  #   if: github.event_name == 'push' || github.event.pull_request.head.repo.full_name != github.repository
-  #   strategy:
-  #     matrix:
-  #       ruby:
-  #         - '2.7'
-  #   steps:
-  #     - name: Checkout
-  #       uses: actions/checkout@v4
-  #     - name: Setup Ruby
-  #       uses: ruby/setup-ruby@v1
-  #       with:
-  #         ruby-version: ${{ matrix.ruby }}
-  #         bundler-cache: true
-  #     - name: Rubocop Cache
-  #       uses: actions/cache@v3
-  #       with:
-  #         path: ~/.cache/rubocop_cache
-  #         key: ${{ runner.os }}-rubocop-${{ hashFiles('.rubocop.yml') }}
-  #         restore-keys: |
-  #           ${{ runner.os }}-rubocop-
-  #     - name: Rubocop
-  #       run: bundle exec rubocop
   test:
     name: Tests
     runs-on: ubuntu-22.04
@@ -50,6 +24,11 @@ jobs:
         options: --health-cmd pg_isready --health-interval 10s --health-timeout 5s --health-retries 5
         ports:
           - 5432:5432
+      redis:
+        image: redis
+        options: --health-cmd "redis-cli ping" --health-interval 10s --health-timeout 5s --health-retries 5
+        ports:
+          - 6379:6379
     steps:
       - name: Checkout
         uses: actions/checkout@v4

data/CHANGELOG.md

@@ -1,4 +1,16 @@
-## [0.5.0] - 2024-02-11
+## 0.7.0
+
+- Allow `Pecorino.adapter` to be assigned, and add `adapter:` to all classes. This allows the adapter for Pecorino to be configured manually and overridden in an initializer.
+- Add Redis-based adapter derived from Prorate
+- Formalize and test the adapter API
+- Add a memory-based adapter for single-process applications (and as a reference)
+- For SQLite tables, do not use UUID primary keys - there is no need for that, and SQLite does not have a antive UUID gen functin that is enabled on all builds
+
+## 0.6.0
+
+- Add `Pecorino::Block` for setting blocks directly. These are available both to `Throttle` with the same key and on their own. This can be used to set arbitrary blocks without having to configure a `Throttle` first.
+
+## 0.5.0
 
 - Add `CachedThrottle` for caching the throttle blocks. This allows protection to the database when the throttle is in a blocked state.
 - Add `Throttle#throttled` for silencing alerts
@@ -6,11 +18,11 @@
 - Allow accessing `Throttle::State` from the `Throttled` exception so that the blocked throttle state can be cached downstream (in Rails cache, for example)
 - Make `Throttle#request!` return the new state if there was no exception raised
 
-## [0.4.1] - 2024-02-11
+## 0.4.1
 
 - Make sure Pecorino works on Ruby 2.7 as well by removing 3.x-exclusive syntax
 
-## [0.4.0] - 2024-01-22
+## 0.4.0
 
 - Use Bucket#connditional_fillup inside Throttle and throttle only when the capacity _would_ be exceeded, as opposed
   to throttling when capacity has already been exceeded. This allows for finer-grained throttles such as
@@ -21,17 +33,17 @@
 - Allow "conditional fillup" - only add tokens to the leaky bucket if the bucket has enough space.
 - Fix `over_time` leading to incorrect `leak_rate`. The divider/divisor were swapped, leading to the inverse leak rate getting computed.
 
-## [0.3.0] - 2024-01-18
+## 0.3.0
 
 - Allow `over_time` in addition to `leak_rate`, which is a more intuitive parameter to tweak
 - Set default `block_for` to the time it takes the bucket to leak out completely instead of 30 seconds
 
-## [0.2.0] - 2024-01-09
+## 0.2.0
 
 - [Add support for SQLite](https://github.com/cheddar-me/pecorino/pull/9)
 - [Use comparisons in SQL to determine whether the leaky bucket did overflow](https://github.com/cheddar-me/pecorino/pull/8)
 - [Change the way Structs are defined to appease Tapioca/Sorbet](https://github.com/cheddar-me/pecorino/pull/6)
 
-## [0.1.0] - 2023-10-30
+## 0.1.0
 
 - Initial release

data/README.md

@@ -1,10 +1,14 @@
 # Pecorino
 
-Pecorino is a rate limiter based on the concept of leaky buckets. It uses your DB as the storage backend for the throttles. It is compact, easy to install, and does not require additional infrastructure. The approach used by Pecorino has been previously used by [prorate](https://github.com/WeTransfer/prorate) with Redis, and that approach has proven itself.
+Pecorino is a rate limiter based on the concept of leaky buckets, or more specifically - based on the [generic cell rate](https://brandur.org/rate-limiting) algorithm. It uses your DB as the storage backend for the throttles. It is compact, easy to install, and does not require additional infrastructure. The approach used by Pecorino has been previously used by [prorate](https://github.com/WeTransfer/prorate) with Redis, and that approach has proven itself.
 
-Pecorino is designed to integrate seamlessly into any Rails application using a PostgreSQL or SQLite database (at the moment there is no MySQL support, we would be delighted if you could add it).
+Pecorino is designed to integrate seamlessly into any Rails application, and will use either:
 
-If you would like to know more about the leaky bucket algorithm: [this article](http://live.julik.nl/2022/08/the-unreasonable-effectiveness-of-leaky-buckets) or the [Wikipedia article](https://en.wikipedia.org/wiki/Leaky_bucket) are both good starting points.
+* A memory store (good enough if you have just 1 process)
+* A PostgreSQL or SQLite database (at the moment there is no MySQL support, we would be delighted if you could add it)
+* A Redis instance
+
+If you would like to know more about the leaky bucket algorithm: [this article](http://live.julik.nl/2022/08/the-unreasonable-effectiveness-of-leaky-buckets) or the [Wikipedia article](https://en.wikipedia.org/wiki/Leaky_bucket) are both good starting points. [This Wikipedia article](https://en.wikipedia.org/wiki/Generic_cell_rate_algorithm) describes the generic cell rate algorithm in more detail as well.
 
 ## Installation
 
@@ -103,6 +107,58 @@ end
 
 This way, every time there is an error on the "fancy AI service" the throttle will be triggered, and if it overflows - a subsequent request will be blocked.
 
+## A note on database transactions
+
+Pecorino uses your main database. When calling the `Throttle` or `LeakyBucket` objects, SQL queries will be performed by Pecorino and those queries may result in changes to data. If you are currently inside a database transaction, your bucket topups or set blocks may get reverted. For example, imagine you have a controller like this:
+
+```ruby
+class WalletController < ApplicationController
+  rescue_from Pecorino::Throttle::Throttled do |e|
+    response.set_header('Retry-After', e.retry_after.to_s)
+    render nothing: true, status: 429
+  end
+
+  def withdraw
+     Wallet.transaction do
+       t = Pecorino::Throttle.new("wallet_#{current_user.id}_max_withdrawal", capacity: 200_00, over_time: 5.minutes)
+       t.request!(10_00)
+       current_user.wallet.withdraw(Money.new(10, "EUR"))
+     end
+  end
+end
+```
+
+what will happen is that even though the `withdraw()` call is not going to be performed, the increment of the throttle will not either, because the exception will result in a `ROLLBACK`.
+
+If you need to use Pecorino in combination with transactions, you will need to design with that in mind. Either call `Throttle` before entering the `transaction do`:
+
+```ruby
+def withdraw
+  t = Pecorino::Throttle.new("wallet_#{current_user.id}_max_withdrawal", capacity: 200_00, over_time: 5.minutes)
+  t.request!(10_00)
+  Wallet.transaction do
+    current_user.wallet.withdraw(Money.new(10, "EUR"))
+  end
+end
+```
+
+or use the `request()` method instead to still commit:
+
+```ruby
+def withdraw
+  Wallet.transaction do
+    t = Pecorino::Throttle.new("wallet_#{current_user.id}_max_withdrawal", capacity: 200_00, over_time: 5.minutes)
+    throttle_state = t.request(10_00)
+    return render(nothing: true, status: 429) if throttle_state.blocked?
+
+    current_user.wallet.withdraw(Money.new(10, "EUR"))
+  end
+end
+```
+
+Note also that this behaviour might be desirable for your use case (that the throttle and the data update together in
+a transactional manner) – it just helps to be aware of it.
+
 ## Using just the leaky bucket
 
 Sometimes you don't want to use a throttle, but you want to track the amount added to the leaky bucket over time. A lower-level abstraction is available for that purpose in the form of the `LeakyBucket` class. It will not raise any exceptions and will not install blocks, but will permit you to track a bucket's state over time:

data/lib/pecorino/adapters/base_adapter.rb

@@ -0,0 +1,66 @@
+# frozen_string_literal: true
+
+# An adapter allows Pecorino throttles, leaky buckets and other
+# resources to interfact to a data storage backend - a database, usually.
+class Pecorino::Adapters::BaseAdapter
+  # Returns the state of a leaky bucket. The state should be a tuple of two
+  # values: the current level (Float) and whether the bucket is now at capacity (Boolean)
+  #
+  # @param key[String] the key of the leaky bucket
+  # @param capacity[Float] the capacity of the leaky bucket to limit to
+  # @param leak_rate[Float] how many tokens leak out of the bucket per second
+  # @return [Array]
+  def state(key:, capacity:, leak_rate:)
+    [0, false]
+  end
+
+  # Adds tokens to the leaky bucket. The return value is a tuple of two
+  # values: the current level (Float) and whether the bucket is now at capacity (Boolean)
+  #
+  # @param key[String] the key of the leaky bucket
+  # @param capacity[Float] the capacity of the leaky bucket to limit to
+  # @param leak_rate[Float] how many tokens leak out of the bucket per second
+  # @param n_tokens[Float] how many tokens to add
+  # @return [Array]
+  def add_tokens(key:, capacity:, leak_rate:, n_tokens:)
+    [0, false]
+  end
+
+  # Adds tokens to the leaky bucket conditionally. If there is capacity, the tokens will
+  # be added. If there isn't - the fillup will be rejected. The return value is a triplet of
+  # the current level (Float), whether the bucket is now at capacity (Boolean)
+  # and whether the fillup was accepted (Boolean)
+  #
+  # @param key[String] the key of the leaky bucket
+  # @param capacity[Float] the capacity of the leaky bucket to limit to
+  # @param leak_rate[Float] how many tokens leak out of the bucket per second
+  # @param n_tokens[Float] how many tokens to add
+  # @return [Array]
+  def add_tokens_conditionally(key:, capacity:, leak_rate:, n_tokens:)
+    [0, false, false]
+  end
+
+  # Sets a timed block for the given key - this is used when a throttle fires. The return value
+  # is not defined - the call should always succeed.
+  # @param key[String] the key of the block
+  # @param block_for[#to_f, Active Support Duration] the duration of the block, in seconds
+  def set_block(key:, block_for:)
+  end
+
+  # Returns the time until which a block for a given key is in effect. If there is no block in
+  # effect, the method should return `nil`. The return value is either a `Time` or `nil`
+  # @param key[String] the key of the block
+  def blocked_until(key:)
+  end
+
+  # Deletes leaky buckets which have an expiry value prior to now and throttle blocks which have
+  # now lapsed
+  # @return [void]
+  def prune
+  end
+
+  # Creates the database tables for Pecorino to operate, or initializes other
+  # schema-like resources the adapter needs to operate
+  def create_tables(active_record_schema)
+  end
+end

data/lib/pecorino/adapters/memory_adapter.rb

@@ -0,0 +1,147 @@
+# frozen_string_literal: true
+
+# A memory store for leaky buckets and blocks
+class Pecorino::Adapters::MemoryAdapter
+  class KeyedLock
+    def initialize
+      @locked_keys = Set.new
+      @lock_mutex = Mutex.new
+    end
+
+    def lock(key)
+      loop do
+        @lock_mutex.synchronize do
+          next if @locked_keys.include?(key)
+          @locked_keys << key
+          return
+        end
+      end
+    end
+
+    def unlock(key)
+      @lock_mutex.synchronize do
+        @locked_keys.delete(key)
+      end
+    end
+
+    def with(key)
+      lock(key)
+      yield
+    ensure
+      unlock(key)
+    end
+  end
+
+  def initialize
+    @buckets = {}
+    @blocks = {}
+    @lock = KeyedLock.new
+  end
+
+  # Returns the state of a leaky bucket. The state should be a tuple of two
+  # values: the current level (Float) and whether the bucket is now at capacity (Boolean)
+  def state(key:, capacity:, leak_rate:)
+    @lock.lock(key)
+    level, ts = @buckets[key]
+    @lock.unlock(key)
+
+    return [0, false] unless level
+
+    dt = get_mono_time - ts
+    level_after_leak = [0, level - (leak_rate * dt)].max
+    [level_after_leak.to_f, (level_after_leak - capacity) >= 0]
+  end
+
+  # Adds tokens to the leaky bucket. The return value is a tuple of two
+  # values: the current level (Float) and whether the bucket is now at capacity (Boolean)
+  def add_tokens(key:, capacity:, leak_rate:, n_tokens:)
+    add_tokens_with_lock(key, capacity, leak_rate, n_tokens, _conditionally = false)
+  end
+
+  # Adds tokens to the leaky bucket conditionally. If there is capacity, the tokens will
+  # be added. If there isn't - the fillup will be rejected. The return value is a triplet of
+  # the current level (Float), whether the bucket is now at capacity (Boolean)
+  # and whether the fillup was accepted (Boolean)
+  def add_tokens_conditionally(key:, capacity:, leak_rate:, n_tokens:)
+    add_tokens_with_lock(key, capacity, leak_rate, n_tokens, _conditionally = true)
+  end
+
+  # Sets a timed block for the given key - this is used when a throttle fires. The return value
+  # is not defined - the call should always succeed.
+  def set_block(key:, block_for:)
+    raise ArgumentError, "block_for must be positive" unless block_for > 0
+    @lock.lock(key)
+    @blocks[key] = get_mono_time + block_for.to_f
+    Time.now + block_for.to_f
+  ensure
+    @lock.unlock(key)
+  end
+
+  # Returns the time until which a block for a given key is in effect. If there is no block in
+  # effect, the method should return `nil`. The return value is either a `Time` or `nil`
+  def blocked_until(key:)
+    blocked_until_monotonic = @blocks[key]
+    return unless blocked_until_monotonic
+
+    now_monotonic = get_mono_time
+    return unless blocked_until_monotonic > now_monotonic
+
+    Time.now + (blocked_until_monotonic - now_monotonic)
+  end
+
+  # Deletes leaky buckets which have an expiry value prior to now and throttle blocks which have
+  # now lapsed
+  def prune
+    now_monotonic = get_mono_time
+
+    @blocks.keys.each do |key|
+      @lock.with(key) do
+        @blocks.delete(key) if @blocks[key] && @blocks[key] < now_monotonic
+      end
+    end
+
+    @buckets.keys.each do |key|
+      @lock.with(key) do
+        _level, expire_at_monotonic = @buckets[key]
+        @buckets.delete(key) if expire_at_monotonic && expire_at_monotonic < now_monotonic
+      end
+    end
+  end
+
+  # No-op
+  def create_tables(active_record_schema)
+  end
+
+  private
+
+  def add_tokens_with_lock(key, capacity, leak_rate, n_tokens, conditionally)
+    @lock.lock(key)
+    now = get_mono_time
+    level, ts, _ = @buckets[key] || [0.0, now]
+
+    dt = now - ts
+    level_after_leak = clamp(0, level - (leak_rate * dt), capacity)
+    level_after_fillup = level_after_leak + n_tokens
+    if level_after_fillup > capacity && conditionally
+      return [level_after_leak, level_after_leak >= capacity, _did_accept = false]
+    end
+
+    clamped_level_after_fillup = clamp(0, level_after_fillup, capacity)
+    expire_after = now + (level_after_fillup / leak_rate)
+    @buckets[key] = [clamped_level_after_fillup, now, expire_after]
+
+    [clamped_level_after_fillup, clamped_level_after_fillup == capacity, _did_accept = true]
+  ensure
+    @lock.unlock(key)
+  end
+
+  def get_mono_time
+    Process.clock_gettime(Process::CLOCK_MONOTONIC)
+  end
+
+  def clamp(min, value, max)
+    return min if value < min
+    return max if value > max
+    value
+  end
+end

data/lib/pecorino/{postgres.rb → adapters/postgres_adapter.rb}

@@ -1,6 +1,10 @@
 # frozen_string_literal: true
 
-Pecorino::Postgres = Struct.new(:model_class) do
+class Pecorino::Adapters::PostgresAdapter
+  def initialize(model_class)
+    @model_class = model_class
+  end
+
   def state(key:, capacity:, leak_rate:)
     query_params = {
       key: key.to_s,
@@ -10,7 +14,7 @@ Pecorino::Postgres = Struct.new(:model_class) do
     # The `level` of the bucket is what got stored at `last_touched_at` time, and we can
     # extrapolate from it to see how many tokens have leaked out since `last_touched_at` -
     # we don't need to UPDATE the value in the bucket here
-    sql = model_class.sanitize_sql_array([<<~SQL, query_params])
+    sql = @model_class.sanitize_sql_array([<<~SQL, query_params])
       SELECT
         GREATEST(
           0.0, LEAST(
@@ -26,7 +30,7 @@ Pecorino::Postgres = Struct.new(:model_class) do
 
     # If the return value of the query is a NULL it means no such bucket exists,
     # so we assume the bucket is empty
-    current_level = model_class.connection.uncached { model_class.connection.select_value(sql) } || 0.0
+    current_level = @model_class.connection.uncached { @model_class.connection.select_value(sql) } || 0.0
     [current_level, capacity - current_level.abs < 0.01]
   end
 
@@ -45,7 +49,7 @@ Pecorino::Postgres = Struct.new(:model_class) do
       fillup: n_tokens.to_f
     }
 
-    sql = model_class.sanitize_sql_array([<<~SQL, query_params])
+    sql = @model_class.sanitize_sql_array([<<~SQL, query_params])
       INSERT INTO pecorino_leaky_buckets AS t
         (key, last_touched_at, may_be_deleted_after, level)
       VALUES
@@ -79,7 +83,7 @@ Pecorino::Postgres = Struct.new(:model_class) do
     # query as a repeat (since we use "select_one" for the RETURNING bit) and will not call into Postgres
     # correctly, thus the clock_timestamp() value would be frozen between calls. We don't want that here.
     # See https://stackoverflow.com/questions/73184531/why-would-postgres-clock-timestamp-freeze-inside-a-rails-unit-test
-    upserted = model_class.connection.uncached { model_class.connection.select_one(sql) }
+    upserted = @model_class.connection.uncached { @model_class.connection.select_one(sql) }
     capped_level_after_fillup, at_capacity = upserted.fetch("level"), upserted.fetch("at_capacity")
     [capped_level_after_fillup, at_capacity]
   end
@@ -99,7 +103,7 @@ Pecorino::Postgres = Struct.new(:model_class) do
       fillup: n_tokens.to_f
     }
 
-    sql = model_class.sanitize_sql_array([<<~SQL, query_params])
+    sql = @model_class.sanitize_sql_array([<<~SQL, query_params])
       WITH pre AS MATERIALIZED (
         SELECT
           -- Note the double clamping here. First we clamp the "current level - leak" to not go below zero,
@@ -137,15 +141,16 @@ Pecorino::Postgres = Struct.new(:model_class) do
         level AS level_after
     SQL
 
-    upserted = model_class.connection.uncached { model_class.connection.select_one(sql) }
+    upserted = @model_class.connection.uncached { @model_class.connection.select_one(sql) }
     level_after = upserted.fetch("level_after")
     level_before = upserted.fetch("level_before")
     [level_after, level_after >= capacity, level_after != level_before]
   end
 
   def set_block(key:, block_for:)
+    raise ArgumentError, "block_for must be positive" unless block_for > 0
     query_params = {key: key.to_s, block_for: block_for.to_f}
-    block_set_query = model_class.sanitize_sql_array([<<~SQL, query_params])
+    block_set_query = @model_class.sanitize_sql_array([<<~SQL, query_params])
       INSERT INTO pecorino_blocks AS t
         (key, blocked_until)
       VALUES
@@ -154,13 +159,36 @@ Pecorino::Postgres = Struct.new(:model_class) do
         blocked_until = GREATEST(EXCLUDED.blocked_until, t.blocked_until)
       RETURNING blocked_until
     SQL
-    model_class.connection.uncached { model_class.connection.select_value(block_set_query) }
+    @model_class.connection.uncached { @model_class.connection.select_value(block_set_query) }
   end
 
   def blocked_until(key:)
-    block_check_query = model_class.sanitize_sql_array([<<~SQL, key])
+    block_check_query = @model_class.sanitize_sql_array([<<~SQL, key])
       SELECT blocked_until FROM pecorino_blocks WHERE key = ? AND blocked_until >= clock_timestamp() LIMIT 1
     SQL
-    model_class.connection.uncached { model_class.connection.select_value(block_check_query) }
+    @model_class.connection.uncached { @model_class.connection.select_value(block_check_query) }
+  end
+
+  def prune
+    @model_class.connection.execute("DELETE FROM pecorino_blocks WHERE blocked_until < NOW()")
+    @model_class.connection.execute("DELETE FROM pecorino_leaky_buckets WHERE may_be_deleted_after < NOW()")
+  end
+
+  def create_tables(active_record_schema)
+    active_record_schema.create_table :pecorino_leaky_buckets, id: :uuid do |t|
+      t.string :key, null: false
+      t.float :level, null: false
+      t.datetime :last_touched_at, null: false
+      t.datetime :may_be_deleted_after, null: false
+    end
+    active_record_schema.add_index :pecorino_leaky_buckets, [:key], unique: true
+    active_record_schema.add_index :pecorino_leaky_buckets, [:may_be_deleted_after]
+
+    active_record_schema.create_table :pecorino_blocks, id: :uuid do |t|
+      t.string :key, null: false
+      t.datetime :blocked_until, null: false
+    end
+    active_record_schema.add_index :pecorino_blocks, [:key], unique: true
+    active_record_schema.add_index :pecorino_blocks, [:blocked_until]
   end
 end

data/lib/pecorino/adapters/redis_adapter/add_tokens_conditionally.lua

@@ -0,0 +1,90 @@
+-- Single threaded Leaky Bucket implementation (without blocking).
+-- args: key_base, leak_rate, bucket_ttl, fillup. To just verify the state of the bucket leak_rate of 0 may be passed.
+-- returns: the leve of the bucket in number of tokens.
+-- This script is largely adapted from Prorate https://github.com/WeTransfer/prorate
+
+-- this is required to be able to use TIME and writes; basically it lifts the script into IO
+redis.replicate_commands()
+
+-- Redis documentation recommends passing the keys separately so that Redis
+-- can - in the future - verify that they live on the same shard of a cluster, and
+-- raise an error if they are not. As far as can be understood this functionality is not
+-- yet present, but if we can make a little effort to make ourselves more future proof
+-- we should.
+local bucket_level_key = KEYS[1]
+local last_updated_key = KEYS[2]
+
+local leak_rate = tonumber(ARGV[1])
+local fillup = tonumber(ARGV[2]) -- How many tokens this call adds to the bucket.
+local bucket_capacity = tonumber(ARGV[3]) -- How many tokens is the bucket allowed to contain
+local conditional_fillup = tonumber(ARGV[4]) -- Whether to fillup conditionally
+
+-- Compute the key TTL for the bucket. We are interested in how long it takes the bucket
+-- to leak all the way to 0, as this is the time when the values stay relevant. We pad with 1 second
+-- to have a little cushion.
+local key_lifetime = math.ceil((bucket_capacity / leak_rate) + 1)
+
+-- Take a timestamp
+local redis_time = redis.call("TIME") -- Array of [seconds, microseconds]
+local now = tonumber(redis_time[1]) + (tonumber(redis_time[2]) / 1000000)
+
+-- get current bucket level. The throttle key might not exist yet in which
+-- case we default to 0
+local bucket_level = tonumber(redis.call("GET", bucket_level_key)) or 0
+
+-- ...and then perform the leaky bucket fillup/leak. We need to do this also when the bucket has
+-- just been created because the initial fillup to add might be so high that it will
+-- immediately overflow the bucket and trigger the throttle, on the first call.
+local last_updated = tonumber(redis.call("GET", last_updated_key)) or now -- use sensible default of 'now' if the key does not exist
+
+-- Subtract the number of tokens leaked since last call
+local dt = now - last_updated
+local bucket_level_after_leaking = math.max(0, math.min(bucket_level - (leak_rate * dt), bucket_capacity))
+local bucket_level_after_fillup = bucket_level_after_leaking + fillup
+local did_accept = 0
+
+-- Figure out whether the fillup would overflow the bucket
+if conditional_fillup == 1 and bucket_level_after_fillup > bucket_capacity then
+  local at_capacity = bucket_level_after_leaking >= bucket_capacity
+  -- See below about string return
+  return {string.format("%.9f", bucket_level_after_leaking), at_capacity, did_accept}
+end
+
+-- and _then_ and add the tokens we fillup with. Cap the value to be 0 < capacity
+local new_bucket_level = math.max(0, math.min(bucket_capacity, bucket_level_after_fillup))
+
+-- Since we return a floating point number string-formatted even if the bucket is full we
+-- have some loss of precision in the formatting, even if the bucket was actually full.
+-- This bit of information is useful to preserve.
+local at_capacity = 0
+if new_bucket_level == bucket_capacity then
+  at_capacity = 1
+end
+
+did_accept = 1
+
+-- If both the initial level was 0, and the level after putting tokens in is 0 we
+-- can avoid setting keys in Redis at all as this was only a level check.
+if new_bucket_level == 0 and bucket_level == 0 then
+  return {"0.0", at_capacity, did_accept}
+end
+
+-- Save the new bucket level
+redis.call("SETEX", bucket_level_key, key_lifetime, new_bucket_level)
+
+-- Record when we updated the bucket so that the amount of tokens leaked
+-- can be correctly determined on the next invocation
+redis.call("SETEX", last_updated_key, key_lifetime, now)
+
+-- Most Redis adapters when used with the Lua interface truncate floats
+-- to integers (at least in Python that is documented to be the case in
+-- the Redis ebook here
+-- https://redislabs.com/ebook/part-3-next-steps/chapter-11-scripting-redis-with-lua/11-1-adding-functionality-without-writing-c
+-- We need access to the bucket level as a float value since our leak rate might as well be floating point, and to achieve that
+-- we can go two ways. We can turn the float into a Lua string, and then parse it on the other side, or we can convert it to
+-- a tuple of two integer values - one for the integer component and one for fraction.
+-- Now, the unpleasant aspect is that when we do this we will lose precision - the number is not going to be
+-- exactly equal to capacity, thus we lose the bit of information which tells us whether we filled up the bucket or not.
+-- Also since the only moment we can register whether the bucket is above capacity is now - in this script, since
+-- by the next call some tokens will have leaked.
+return {string.format("%.9f", new_bucket_level), at_capacity, did_accept}

data/lib/pecorino/adapters/redis_adapter.rb

@@ -0,0 +1,95 @@
+# frozen_string_literal: true
+
+require_relative "base_adapter"
+require "digest"
+require "redis"
+
+# An adapter for storing Pecorino leaky buckets and blocks in Redis. It uses Lua
+# to enforce atomicity for leaky bucket operations
+class Pecorino::Adapters::RedisAdapter < Pecorino::Adapters::BaseAdapter
+  class RedisScript
+    def initialize(script_filename)
+      @script_body = File.read(File.dirname(__FILE__) + "/redis_adapter/" + script_filename)
+      @sha = Digest::SHA1.hexdigest(@script_body)
+    end
+
+    def load_and_eval(redis, keys, argv)
+      redis.evalsha(@sha, keys: keys, argv: argv)
+    rescue Redis::CommandError => e
+      if e.message.include? "NOSCRIPT"
+        redis.script(:load, @script_body)
+        retry
+      else
+        raise e
+      end
+    end
+  end
+
+  ADD_TOKENS_SCRIPT = RedisScript.new("add_tokens_conditionally.lua")
+
+  def initialize(redis_connection_or_connection_pool, key_prefix: "pecorino")
+    @redis_pool = redis_connection_or_connection_pool
+    @key_prefix = key_prefix
+  end
+
+  # Returns the state of a leaky bucket. The state should be a tuple of two
+  # values: the current level (Float) and whether the bucket is now at capacity (Boolean)
+  def state(key:, capacity:, leak_rate:)
+    add_tokens(key: key, capacity: capacity, leak_rate: leak_rate, n_tokens: 0)
+  end
+
+  # Adds tokens to the leaky bucket. The return value is a tuple of two
+  # values: the current level (Float) and whether the bucket is now at capacity (Boolean)
+  def add_tokens(key:, capacity:, leak_rate:, n_tokens:)
+    keys = ["#{@key_prefix}:leaky_bucket:#{key}:level", "#{@key_prefix}:leaky_bucket:#{key}:last_touched"]
+    argv = [leak_rate, n_tokens, capacity, _conditional = 0]
+    decimal_float_level, at_capacity_int, _ = with_redis do |redis|
+      ADD_TOKENS_SCRIPT.load_and_eval(redis, keys, argv)
+    end
+    [decimal_float_level.to_f, at_capacity_int == 1]
+  end
+
+  # Adds tokens to the leaky bucket conditionally. If there is capacity, the tokens will
+  # be added. If there isn't - the fillup will be rejected. The return value is a triplet of
+  # the current level (Float), whether the bucket is now at capacity (Boolean)
+  # and whether the fillup was accepted (Boolean)
+  def add_tokens_conditionally(key:, capacity:, leak_rate:, n_tokens:)
+    keys = ["#{@key_prefix}:leaky_bucket:#{key}:level", "#{@key_prefix}:leaky_bucket:#{key}:last_touched"]
+    argv = [leak_rate, n_tokens, capacity, _conditional = 1]
+    decimal_float_level, at_capacity_int, did_accept_int = with_redis do |redis|
+      ADD_TOKENS_SCRIPT.load_and_eval(redis, keys, argv)
+    end
+    [decimal_float_level.to_f, at_capacity_int == 1, did_accept_int == 1]
+  end
+
+  # Sets a timed block for the given key - this is used when a throttle fires. The return value
+  # is not defined - the call should always succeed.
+  def set_block(key:, block_for:)
+    raise ArgumentError, "block_for must be positive" unless block_for > 0
+    blocked_until = Time.now + block_for
+    with_redis do |r|
+      r.setex("#{@key_prefix}:leaky_bucket:#{key}:block", block_for.to_f.ceil, blocked_until.to_f)
+    end
+    blocked_until
+  end
+
+  # Returns the time until which a block for a given key is in effect. If there is no block in
+  # effect, the method should return `nil`. The return value is either a `Time` or `nil`
+  def blocked_until(key:)
+    seconds_from_epoch = with_redis do |r|
+      r.get("#{@key_prefix}:leaky_bucket:#{key}:block")
+    end
+    return unless seconds_from_epoch
+    Time.at(seconds_from_epoch.to_f).utc
+  end
+
+  private
+
+  def with_redis
+    if @redis_pool.respond_to?(:with)
+      @redis_pool.with { |conn| yield(conn) }
+    else
+      yield @redis_pool
+    end
+  end
+end

data/lib/pecorino/{sqlite.rb → adapters/sqlite_adapter.rb}

@@ -1,6 +1,10 @@
 # frozen_string_literal: true
 
-Pecorino::Sqlite = Struct.new(:model_class) do
+class Pecorino::Adapters::SqliteAdapter
+  def initialize(model_class)
+    @model_class = model_class
+  end
+
   def state(key:, capacity:, leak_rate:)
     # With a server database, it is really important to use the clock of the database itself so
     # that concurrent requests will see consistent bucket level calculations. Since SQLite is
@@ -17,7 +21,7 @@ Pecorino::Sqlite = Struct.new(:model_class) do
     # The `level` of the bucket is what got stored at `last_touched_at` time, and we can
     # extrapolate from it to see how many tokens have leaked out since `last_touched_at` -
     # we don't need to UPDATE the value in the bucket here
-    sql = model_class.sanitize_sql_array([<<~SQL, query_params])
+    sql = @model_class.sanitize_sql_array([<<~SQL, query_params])
       SELECT
         MAX(
           0.0, MIN(
@@ -33,7 +37,7 @@ Pecorino::Sqlite = Struct.new(:model_class) do
 
     # If the return value of the query is a NULL it means no such bucket exists,
     # so we assume the bucket is empty
-    current_level = model_class.connection.uncached { model_class.connection.select_value(sql) } || 0.0
+    current_level = @model_class.connection.uncached { @model_class.connection.select_value(sql) } || 0.0
     [current_level, capacity - current_level.abs < 0.01]
   end
 
@@ -50,16 +54,14 @@ Pecorino::Sqlite = Struct.new(:model_class) do
       delete_after_s: may_be_deleted_after_seconds,
       leak_rate: leak_rate.to_f,
       now_s: Time.now.to_f, # See above as to why we are using a time value passed in
-      fillup: n_tokens.to_f,
-      id: SecureRandom.uuid # SQLite3 does not autogenerate UUIDs
+      fillup: n_tokens.to_f
     }
 
-    sql = model_class.sanitize_sql_array([<<~SQL, query_params])
+    sql = @model_class.sanitize_sql_array([<<~SQL, query_params])
       INSERT INTO pecorino_leaky_buckets AS t
-        (id, key, last_touched_at, may_be_deleted_after, level)
+        (key, last_touched_at, may_be_deleted_after, level)
       VALUES
         (
-          :id,
           :key,
           :now_s, -- Precision loss must be avoided here as it is used for calculations
           DATETIME('now', '+:delete_after_s seconds'), -- Precision loss is acceptable here
@@ -89,7 +91,7 @@ Pecorino::Sqlite = Struct.new(:model_class) do
     # query as a repeat (since we use "select_one" for the RETURNING bit) and will not call into Postgres
     # correctly, thus the clock_timestamp() value would be frozen between calls. We don't want that here.
     # See https://stackoverflow.com/questions/73184531/why-would-postgres-clock-timestamp-freeze-inside-a-rails-unit-test
-    upserted = model_class.connection.uncached { model_class.connection.select_one(sql) }
+    upserted = @model_class.connection.uncached { @model_class.connection.select_one(sql) }
     capped_level_after_fillup, one_if_did_overflow = upserted.fetch("level"), upserted.fetch("did_overflow")
     [capped_level_after_fillup, one_if_did_overflow == 1]
   end
@@ -107,19 +109,17 @@ Pecorino::Sqlite = Struct.new(:model_class) do
       delete_after_s: may_be_deleted_after_seconds,
       leak_rate: leak_rate.to_f,
       now_s: Time.now.to_f, # See above as to why we are using a time value passed in
-      fillup: n_tokens.to_f,
-      id: SecureRandom.uuid # SQLite3 does not autogenerate UUIDs
+      fillup: n_tokens.to_f
     }
 
     # Sadly with SQLite we need to do an INSERT first, because otherwise the inserted row is visible
     # to the WITH clause, so we cannot combine the initial fillup and the update into one statement.
     # This shuld be fine however since we will suppress the INSERT on a key conflict
-    insert_sql = model_class.sanitize_sql_array([<<~SQL, query_params])
+    insert_sql = @model_class.sanitize_sql_array([<<~SQL, query_params])
       INSERT INTO pecorino_leaky_buckets AS t
-        (id, key, last_touched_at, may_be_deleted_after, level)
+        (key, last_touched_at, may_be_deleted_after, level)
       VALUES
         (
-          :id,
           :key,
           :now_s, -- Precision loss must be avoided here as it is used for calculations
           DATETIME('now', '+:delete_after_s seconds'), -- Precision loss is acceptable here
@@ -130,9 +130,9 @@ Pecorino::Sqlite = Struct.new(:model_class) do
       -- so that it can't be deleted between our INSERT and our UPDATE
         may_be_deleted_after = EXCLUDED.may_be_deleted_after
     SQL
-    model_class.connection.execute(insert_sql)
+    @model_class.connection.execute(insert_sql)
 
-    sql = model_class.sanitize_sql_array([<<~SQL, query_params])
+    sql = @model_class.sanitize_sql_array([<<~SQL, query_params])
       -- With SQLite MATERIALIZED has to be used so that level_post is calculated before the UPDATE takes effect
       WITH pre(level_post_with_uncapped_fillup, level_post) AS MATERIALIZED (
         SELECT
@@ -156,30 +156,31 @@ Pecorino::Sqlite = Struct.new(:model_class) do
         level AS level_after
     SQL
 
-    upserted = model_class.connection.uncached { model_class.connection.select_one(sql) }
+    upserted = @model_class.connection.uncached { @model_class.connection.select_one(sql) }
     level_after = upserted.fetch("level_after")
     level_before = upserted.fetch("level_before")
     [level_after, level_after >= capacity, level_after != level_before]
   end
 
   def set_block(key:, block_for:)
-    query_params = {id: SecureRandom.uuid, key: key.to_s, block_for: block_for.to_f, now_s: Time.now.to_f}
-    block_set_query = model_class.sanitize_sql_array([<<~SQL, query_params])
+    raise ArgumentError, "block_for must be positive" unless block_for > 0
+    query_params = {key: key.to_s, block_for: block_for.to_f, now_s: Time.now.to_f}
+    block_set_query = @model_class.sanitize_sql_array([<<~SQL, query_params])
       INSERT INTO pecorino_blocks AS t
-        (id, key, blocked_until)
+        (key, blocked_until)
       VALUES
-        (:id, :key, :now_s + :block_for)
+        (:key, :now_s + :block_for)
       ON CONFLICT (key) DO UPDATE SET
         blocked_until = MAX(EXCLUDED.blocked_until, t.blocked_until)
       RETURNING blocked_until;
     SQL
-    blocked_until_s = model_class.connection.uncached { model_class.connection.select_value(block_set_query) }
+    blocked_until_s = @model_class.connection.uncached { @model_class.connection.select_value(block_set_query) }
     Time.at(blocked_until_s)
   end
 
   def blocked_until(key:)
     now_s = Time.now.to_f
-    block_check_query = model_class.sanitize_sql_array([<<~SQL, {now_s: now_s, key: key}])
+    block_check_query = @model_class.sanitize_sql_array([<<~SQL, {now_s: now_s, key: key}])
       SELECT
         blocked_until
       FROM
@@ -187,7 +188,31 @@ Pecorino::Sqlite = Struct.new(:model_class) do
       WHERE
         key = :key AND blocked_until >= :now_s LIMIT 1
     SQL
-    blocked_until_s = model_class.connection.uncached { model_class.connection.select_value(block_check_query) }
+    blocked_until_s = @model_class.connection.uncached { @model_class.connection.select_value(block_check_query) }
     blocked_until_s && Time.at(blocked_until_s)
   end
+
+  def prune
+    now_s = Time.now.to_f
+    @model_class.connection.execute("DELETE FROM pecorino_blocks WHERE blocked_until < ?", now_s)
+    @model_class.connection.execute("DELETE FROM pecorino_leaky_buckets WHERE may_be_deleted_after < ?", now_s)
+  end
+
+  def create_tables(active_record_schema)
+    active_record_schema.create_table :pecorino_leaky_buckets do |t|
+      t.string :key, null: false
+      t.float :level, null: false
+      t.datetime :last_touched_at, null: false
+      t.datetime :may_be_deleted_after, null: false
+    end
+    active_record_schema.add_index :pecorino_leaky_buckets, [:key], unique: true
+    active_record_schema.add_index :pecorino_leaky_buckets, [:may_be_deleted_after]
+
+    active_record_schema.create_table :pecorino_blocks do |t|
+      t.string :key, null: false
+      t.datetime :blocked_until, null: false
+    end
+    active_record_schema.add_index :pecorino_blocks, [:key], unique: true
+    active_record_schema.add_index :pecorino_blocks, [:blocked_until]
+  end
 end

data/lib/pecorino/block.rb

@@ -0,0 +1,29 @@
+# frozen_string_literal: true
+
+# Provides access to Pecorino blocks - same blocks which get set when a throttle triggers. The blocks
+# are just keys in the data store which have an expiry value. This can be useful if you want to restrict
+# access to a resource for an arbitrary timespan.
+class Pecorino::Block
+  # Sets a block for the given key. The block will also be seen by the Pecorino::Throttle with the same key
+  #
+  # @param key[String] the key to set the block for
+  # @param block_for[Float] the number of seconds or a time interval to block for
+  # @param adapter[Pecorino::Adapters::BaseAdapter] the adapter to set the value in.
+  # @return [Time] the time when the block will be released
+  def self.set!(key:, block_for:, adapter: Pecorino.adapter)
+    adapter.set_block(key: key, block_for: block_for)
+    Time.now + block_for
+  rescue ArgumentError # negative block
+    nil
+  end
+
+  # Returns the time until a certain block is in effect
+  #
+  # @param key[String] the key to get the expiry time for
+  # @param adapter[Pecorino::Adapters::BaseAdapter] the adapter to get the value from
+  # @return [Time,nil] the time when the block will be released
+  def self.blocked_until(key:, adapter: Pecorino.adapter)
+    t = adapter.blocked_until(key: key)
+    (t && t > Time.now) ? t : nil
+  end
+end

data/lib/pecorino/install_generator.rb

@@ -5,8 +5,8 @@ require "rails/generators/active_record"
 
 module Pecorino
   #
-  # Rails generator used for setting up GoodJob in a Rails application.
-  # Run it with +bin/rails g good_job:install+ in your console.
+  # Rails generator used for setting up Pecorino in a Rails application.
+  # Run it with +bin/rails g pecorino:install+ in your console.
   #
   class InstallGenerator < Rails::Generators::Base
     include ActiveRecord::Generators::Migration

data/lib/pecorino/leaky_bucket.rb

@@ -90,12 +90,14 @@ class Pecorino::LeakyBucket
   #   the bucket contents will then be capped at this value. So with
   #   bucket_capacity set to 12 and a `fillup(14)` the bucket will reach the level
   #   of 12, and will then immediately start leaking again.
-  def initialize(key:, capacity:, leak_rate: nil, over_time: nil)
+  # @param adapter[Pecorino::Adapters::BaseAdapter] a compatible adapter
+  def initialize(key:, capacity:, adapter: Pecorino.adapter, leak_rate: nil, over_time: nil)
     raise ArgumentError, "Either leak_rate: or over_time: must be specified" if leak_rate.nil? && over_time.nil?
     raise ArgumentError, "Either leak_rate: or over_time: may be specified, but not both" if leak_rate && over_time
     @leak_rate = leak_rate || (capacity / over_time.to_f)
     @key = key
     @capacity = capacity.to_f
+    @adapter = adapter
   end
 
   # Places `n` tokens in the bucket. If the bucket has less capacity than `n` tokens, the bucket will be filled to capacity.
@@ -109,7 +111,7 @@ class Pecorino::LeakyBucket
   # @param n_tokens[Float] How many tokens to fillup by
   # @return [State] the state of the bucket after the operation
   def fillup(n_tokens)
-    capped_level_after_fillup, is_full = Pecorino.adapter.add_tokens(capacity: @capacity, key: @key, leak_rate: @leak_rate, n_tokens: n_tokens)
+    capped_level_after_fillup, is_full = @adapter.add_tokens(capacity: @capacity, key: @key, leak_rate: @leak_rate, n_tokens: n_tokens)
     State.new(capped_level_after_fillup, is_full)
   end
 
@@ -131,7 +133,7 @@ class Pecorino::LeakyBucket
   # @param n_tokens[Float] How many tokens to fillup by
   # @return [ConditionalFillupResult] the state of the bucket after the operation and whether the operation succeeded
   def fillup_conditionally(n_tokens)
-    capped_level_after_fillup, is_full, did_accept = Pecorino.adapter.add_tokens_conditionally(capacity: @capacity, key: @key, leak_rate: @leak_rate, n_tokens: n_tokens)
+    capped_level_after_fillup, is_full, did_accept = @adapter.add_tokens_conditionally(capacity: @capacity, key: @key, leak_rate: @leak_rate, n_tokens: n_tokens)
     ConditionalFillupResult.new(capped_level_after_fillup, is_full, did_accept)
   end
 
@@ -140,7 +142,7 @@ class Pecorino::LeakyBucket
   #
   # @return [State] the snapshotted state of the bucket at time of query
   def state
-    current_level, is_full = Pecorino.adapter.state(key: @key, capacity: @capacity, leak_rate: @leak_rate)
+    current_level, is_full = @adapter.state(key: @key, capacity: @capacity, leak_rate: @leak_rate)
     State.new(current_level, is_full)
   end
 

data/lib/pecorino/throttle.rb

@@ -99,10 +99,13 @@ class Pecorino::Throttle
   # @param key[String] the key for both the block record and the leaky bucket
   # @param block_for[Numeric] the number of seconds to block any further requests for. Defaults to time it takes
   #   the bucket to leak out to the level of 0
+  # @param adapter[Pecorino::Adapters::BaseAdapter] a compatible adapter
   # @param leaky_bucket_options Options for `Pecorino::LeakyBucket.new`
   # @see PecorinoLeakyBucket.new
-  def initialize(key:, block_for: nil, **leaky_bucket_options)
-    @bucket = Pecorino::LeakyBucket.new(key: key, **leaky_bucket_options)
+  def initialize(key:, block_for: nil, adapter: Pecorino.adapter, **leaky_bucket_options)
+    @adapter = adapter
+    leaky_bucket_options.delete(:adapter)
+    @bucket = Pecorino::LeakyBucket.new(key: key, adapter: @adapter, **leaky_bucket_options)
     @key = key.to_s
     @block_for = block_for ? block_for.to_f : (@bucket.capacity / @bucket.leak_rate)
   end
@@ -116,7 +119,7 @@ class Pecorino::Throttle
   # @param n_tokens[Float]
   # @return [boolean]
   def able_to_accept?(n_tokens = 1)
-    Pecorino.adapter.blocked_until(key: @key).nil? && @bucket.able_to_accept?(n_tokens)
+    @adapter.blocked_until(key: @key).nil? && @bucket.able_to_accept?(n_tokens)
   end
 
   # Register that a request is being performed. Will raise Throttled
@@ -156,7 +159,7 @@ class Pecorino::Throttle
   #
   # @return [State] the state of the throttle after filling up the leaky bucket / trying to pass the block
   def request(n = 1)
-    existing_blocked_until = Pecorino.adapter.blocked_until(key: @key)
+    existing_blocked_until = Pecorino::Block.blocked_until(key: @key, adapter: @adapter)
     return State.new(existing_blocked_until.utc) if existing_blocked_until
 
     # Topup the leaky bucket, and if the topup gets rejected - block the caller
@@ -165,7 +168,7 @@ class Pecorino::Throttle
       State.new(nil)
     else
       # and set the block if the fillup was rejected
-      fresh_blocked_until = Pecorino.adapter.set_block(key: @key, block_for: @block_for)
+      fresh_blocked_until = Pecorino::Block.set!(key: @key, block_for: @block_for, adapter: @adapter)
       State.new(fresh_blocked_until.utc)
     end
   end

data/lib/pecorino/version.rb

@@ -1,5 +1,5 @@
 # frozen_string_literal: true
 
 module Pecorino
-  VERSION = "0.5.0"
+  VERSION = "0.7.0"
 end

data/lib/pecorino.rb

@@ -4,26 +4,27 @@ require "active_support/concern"
 require "active_record/sanitization"
 
 require_relative "pecorino/version"
-require_relative "pecorino/leaky_bucket"
-require_relative "pecorino/throttle"
 require_relative "pecorino/railtie" if defined?(Rails::Railtie)
-require_relative "pecorino/cached_throttle"
 
 module Pecorino
-  autoload :Postgres, "pecorino/postgres"
-  autoload :Sqlite, "pecorino/sqlite"
+  autoload :LeakyBucket, "pecorino/leaky_bucket"
+  autoload :Block, "pecorino/block"
+  autoload :Throttle, "pecorino/throttle"
+  autoload :CachedThrottle, "pecorino/cached_throttle"
+
+  module Adapters
+    autoload :MemoryAdapter, "pecorino/adapters/memory_adapter"
+    autoload :PostgresAdapter, "pecorino/adapters/postgres_adapter"
+    autoload :SqliteAdapter, "pecorino/adapters/sqlite_adapter"
+    autoload :RedisAdapter, "pecorino/adapters/redis_adapter"
+  end
 
   # Deletes stale leaky buckets and blocks which have expired. Run this method regularly to
   # avoid accumulating too many unused rows in your tables.
   #
   # @return void
   def self.prune!
-    # Delete all the old blocks here (if we are under a heavy swarm of requests which are all
-    # blocked it is probably better to avoid the big delete)
-    ActiveRecord::Base.connection.execute("DELETE FROM pecorino_blocks WHERE blocked_until < NOW()")
-
-    # Prune buckets which are no longer used. No "uncached" needed here since we are using "execute"
-    ActiveRecord::Base.connection.execute("DELETE FROM pecorino_leaky_buckets WHERE may_be_deleted_after < NOW()")
+    adapter.prune
   end
 
   # Creates the tables and indexes needed for Pecorino. Call this from your migrations like so:
@@ -37,36 +38,41 @@ module Pecorino
   # @param active_record_schema[ActiveRecord::SchemaMigration] the migration through which we will create the tables
   # @return void
   def self.create_tables(active_record_schema)
-    active_record_schema.create_table :pecorino_leaky_buckets, id: :uuid do |t|
-      t.string :key, null: false
-      t.float :level, null: false
-      t.datetime :last_touched_at, null: false
-      t.datetime :may_be_deleted_after, null: false
-    end
-    active_record_schema.add_index :pecorino_leaky_buckets, [:key], unique: true
-    active_record_schema.add_index :pecorino_leaky_buckets, [:may_be_deleted_after]
+    adapter.create_tables(active_record_schema)
+  end
 
-    active_record_schema.create_table :pecorino_blocks, id: :uuid do |t|
-      t.string :key, null: false
-      t.datetime :blocked_until, null: false
-    end
-    active_record_schema.add_index :pecorino_blocks, [:key], unique: true
-    active_record_schema.add_index :pecorino_blocks, [:blocked_until]
+  # Allows assignment of an adapter for storing throttles. Normally this would be a subclass of `Pecorino::Adapters::BaseAdapter`, but
+  # you can assign anything you like. Set this in an initializer. By default Pecorino will use the adapter configured from your main
+  # database, but you can also create a separate database for it - or use Redis or memory storage.
+  #
+  # @param adapter[Pecorino::Adapters::BaseAdapter]
+  # @return [Pecorino::Adapters::BaseAdapter]
+  def self.adapter=(adapter)
+    @adapter = adapter
+  end
+
+  # Returns the currently configured adapter, or the default adapter from the main database
+  #
+  # @return [Pecorino::Adapters::BaseAdapter]
+  def self.adapter
+    @adapter || default_adapter_from_main_database
   end
 
   # Returns the database implementation for setting the values atomically. Since the implementation
   # differs per database, this method will return a different adapter depending on which database is
   # being used
-  def self.adapter
+  #
+  # @param adapter[Pecorino::Adapters::BaseAdapter]
+  def self.default_adapter_from_main_database
     model_class = ActiveRecord::Base
     adapter_name = model_class.connection.adapter_name
     case adapter_name
     when /postgres/i
-      Pecorino::Postgres.new(model_class)
+      Pecorino::Adapters::PostgresAdapter.new(model_class)
     when /sqlite/i
-      Pecorino::Sqlite.new(model_class)
+      Pecorino::Adapters::SqliteAdapter.new(model_class)
     else
-      raise "Pecorino does not support #{adapter_name} just yet"
+      raise "Pecorino does not support the #{adapter_name} database just yet"
     end
   end
 end

data/pecorino.gemspec

@@ -39,6 +39,7 @@ Gem::Specification.new do |spec|
   spec.add_development_dependency "standard"
   spec.add_development_dependency "magic_frozen_string_literal"
   spec.add_development_dependency "minitest-fail-fast"
+  spec.add_development_dependency "redis", "~> 5", "< 6"
 
   # For more information and examples about making a new gem, checkout our
   # guide at: https://bundler.io/guides/creating_gem.html

metadata

@@ -1,14 +1,14 @@
 --- !ruby/object:Gem::Specification
 name: pecorino
 version: !ruby/object:Gem::Version
-  version: 0.5.0
+  version: 0.7.0
 platform: ruby
 authors:
 - Julik Tarkhanov
 autorequire:
 bindir: exe
 cert_chain: []
-date: 2024-02-11 00:00:00.000000000 Z
+date: 2024-05-20 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: activerecord
@@ -136,6 +136,26 @@ dependencies:
     - - ">="
       - !ruby/object:Gem::Version
         version: '0'
+- !ruby/object:Gem::Dependency
+  name: redis
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '5'
+    - - "<"
+      - !ruby/object:Gem::Version
+        version: '6'
+  type: :development
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '5'
+    - - "<"
+      - !ruby/object:Gem::Version
+        version: '6'
 description: Pecorino allows you to define throttles and rate meters for your metered
   resources, all through your standard DB
 email:
@@ -154,13 +174,18 @@ files:
 - README.md
 - Rakefile
 - lib/pecorino.rb
+- lib/pecorino/adapters/base_adapter.rb
+- lib/pecorino/adapters/memory_adapter.rb
+- lib/pecorino/adapters/postgres_adapter.rb
+- lib/pecorino/adapters/redis_adapter.rb
+- lib/pecorino/adapters/redis_adapter/add_tokens_conditionally.lua
+- lib/pecorino/adapters/sqlite_adapter.rb
+- lib/pecorino/block.rb
 - lib/pecorino/cached_throttle.rb
 - lib/pecorino/install_generator.rb
 - lib/pecorino/leaky_bucket.rb
 - lib/pecorino/migrations/create_pecorino_tables.rb.erb
-- lib/pecorino/postgres.rb
 - lib/pecorino/railtie.rb
-- lib/pecorino/sqlite.rb
 - lib/pecorino/throttle.rb
 - lib/pecorino/version.rb
 - pecorino.gemspec
@@ -186,7 +211,7 @@ required_rubygems_version: !ruby/object:Gem::Requirement
     - !ruby/object:Gem::Version
       version: '0'
 requirements: []
-rubygems_version: 3.4.10
+rubygems_version: 3.3.7
 signing_key:
 specification_version: 4
 summary: Database-based rate limiter using leaky buckets