pecorino 0.2.0 → 0.4.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: bfcca80bad8895a9b45ed8a3ed0afe06b52a0c6a851d8506b102a82b5375c2a3
-  data.tar.gz: 7f57dd803797acfdf29a8d7cae31854528012af249887ae49e398b992a48f9d4
+  metadata.gz: 445e9997824e9ef7857a31e626e2a6de981d466fd4f5187299cff86533596b13
+  data.tar.gz: 9b31fad0bf017b2a9ee1b4d4f0660376945db76970c65e61ecadeccf706425b6
 SHA512:
-  metadata.gz: dea8f3c693c1d9b5412cdc50cd333d1de6b92d90ec1358ae3c3f5d03fef685649bccd4f6fcc8757e04aa7aba5b13411b325b69b374be0bf30a86d10c16977613
-  data.tar.gz: ee00695a622947cbdc14a300d73d06a850abd2357b3c1357e7904e7c93f2ea5513a4edbb98349a896b5b39cc57b08e889a221e78319e5344f778000dd7059f06
+  metadata.gz: 9b82a54b4e6f721aa2c752d9f7d0503f5b621174826a76594ef9c91a1b943b881c1e377c8ef7c7d4c57d810a4f9be82ac17bad21c12f65570ebf3daba5c51c7d
+  data.tar.gz: 95436d3b317c43d08a6630ad5e4dba1494cdbcae2f3e98ffccfb9cecf896a8ad2f8a02861d987c0288c1cb04b29e7fb1d490bdf37821b7858b472cdb82e6092b

data/CHANGELOG.md

@@ -1,3 +1,19 @@
+## [0.4.0] - 2024-01-22
+
+- 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
+  "at most once in", where filling "exactly to capacity" is a requirement. It also provides for more accurate
+  and easier to understand throttling in general.
+- Make sure Bucket#able_to_accept? allows the bucket to be filled to capacity, not only to below capacity
+- Improve YARD documentation
+- 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
+
+- 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
 
 - [Add support for SQLite](https://github.com/cheddar-me/pecorino/pull/9)

data/README.md

@@ -22,10 +22,10 @@ And then execute:
 
 ## Usage
 
-Once the installation is done you can use Pecorino to start defining your throttles. Imagine you have a resource called `vault` and you want to limit the number of updates to it to 5 per second. To achieve that, instantiate a new `Throttle` in your controller or job code, and then trigger it using `Throttle#request!`. A call to `request!` registers 1 token getting added to the bucket. If the bucket is full, or the throttle is currently in "block" mode (has recently been triggered), a `Pecorino::Throttle::Throttled` exception will be raised.
+Once the installation is done you can use Pecorino to start defining your throttles. Imagine you have a resource called `vault` and you want to limit the number of updates to it to 5 per second. To achieve that, instantiate a new `Throttle` in your controller or job code, and then trigger it using `Throttle#request!`. A call to `request!` registers 1 token getting added to the bucket. If the bucket would overspill (your request would make it overflow), or the throttle is currently in "block" mode (has recently been triggered), a `Pecorino::Throttle::Throttled` exception will be raised.
 
 ```ruby
-throttle = Pecorino::Throttle.new(key: "vault", leak_rate: 5, capacity: 5)
+throttle = Pecorino::Throttle.new(key: "vault", over_time: 1.second, capacity: 5)
 throttle.request!
 ```
 In a Rails controller you can then rescue from this exception to render the appropriate response:
@@ -58,7 +58,7 @@ return render :capacity_exceeded unless throttle.able_to_accept?
 If you are dealing with a metered resource (like throughput, money, amount of storage...) you can supply the number of tokens to either `request!` or `able_to_accept?` to indicate the desired top-up of the leaky bucket. For example, if you are maintaining user wallets and want to ensure no more than 100 dollars may be taken from the wallet within a certain amount of time, you can do it like so:
 
 ```ruby
-throttle = Pecorino::Throttle.new(key: "wallet_t_#{current_user.id}", leak_rate: 100 / 60.0 / 60.0, capacity: 100, block_for: 60*60*3)
+throttle = Pecorino::Throttle.new(key: "wallet_t_#{current_user.id}", over_time_: 1.hour, capacity: 100, block_for: 3.hours)
 throttle.request!(20) # Attempt to withdraw 20 dollars
 throttle.request!(20) # Attempt to withdraw 20 dollars more
 throttle.request!(20) # Attempt to withdraw 20 dollars more
@@ -67,6 +67,8 @@ throttle.request!(20) # Attempt to withdraw 20 dollars more
 throttle.request!(2) # Attempt to withdraw 2 dollars more, will raise `Throttled` and block withdrawals for 3 hours
 ```
 
+## 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:
 
 
@@ -77,9 +79,10 @@ sleep 0.2
 b.state #=> Pecorino::LeakyBucket::State(full?: false, level: 1.8)
 ```
 
-Check out the inline YARD documentation for more options.
+Check out the inline YARD documentation for more options. Do take note of the differences between `fillup()` and `fillup_conditionally` as you
+might want to pick one or the other depending on your use case.
 
-## Cleaning out stale locks from the database
+## Cleaning out stale buckets and blocks from the database
 
 We recommend running the following bit of code every couple of hours (via cron or similar) to delete the stale blocks and leaky buckets from the system:
 

data/lib/pecorino/leaky_bucket.rb

@@ -25,69 +25,114 @@
 # The storage use is one DB row per leaky bucket you need to manage (likely - one throttled entity such
 # as a combination of an IP address + the URL you need to procect). The `key` is an arbitrary string you provide.
 class Pecorino::LeakyBucket
-  State = Struct.new(:level, :full) do
-    # Returns the level of the bucket after the operation on the LeakyBucket
-    # object has taken place. There is a guarantee that no tokens have leaked
-    # from the bucket between the operation and the freezing of the State
-    # struct.
-    #
-    # @!attribute [r] level
-    #   @return [Float]
+  # Returned from `.state` and `.fillup`
+  class State
+    def initialize(level, is_full)
+      @level = level.to_f
+      @full = !!is_full
+    end
+
+    # Returns the level of the bucket
+    # @return [Float]
+    attr_reader :level
 
     # Tells whether the bucket was detected to be full when the operation on
-    # the LeakyBucket was performed. There is a guarantee that no tokens have leaked
-    # from the bucket between the operation and the freezing of the State
-    # struct.
-    #
-    # @!attribute [r] full
-    #   @return [Boolean]
+    # the LeakyBucket was performed.
+    # @return [Boolean]
+    def full?
+      @full
+    end
 
-    alias_method :full?, :full
+    alias_method :full, :full?
+  end
 
-    # Returns the bucket level of the bucket state as a Float
-    #
-    # @return [Float]
-    def to_f
-      level.to_f
+  # Same as `State` but also communicates whether the write has been permitted or not. A conditional fillup
+  # may refuse a write if it would make the bucket overflow
+  class ConditionalFillupResult < State
+    def initialize(level, is_full, accepted)
+      super(level, is_full)
+      @accepted = !!accepted
     end
 
-    # Returns the bucket level of the bucket state rounded to an Integer
-    #
-    # @return [Integer]
-    def to_i
-      level.to_i
+    # Tells whether the bucket did accept the requested fillup
+    # @return [Boolean]
+    def accepted?
+      @accepted
     end
   end
 
+  # The key (name) of the leaky bucket
+  #   @return [String]
+  attr_reader :key
+
+  # The leak rate (tokens per second) of the bucket
+  #   @return [Float]
+  attr_reader :leak_rate
+
+  # The capacity of the bucket in tokens
+  #   @return [Float]
+  attr_reader :capacity
+
   # Creates a new LeakyBucket. The object controls 1 row in the database is
   # specific to the bucket key.
   #
   # @param key[String] the key for the bucket. The key also gets used
   #   to derive locking keys, so that operations on a particular bucket
   #   are always serialized.
-  # @param leak_rate[Float] the leak rate of the bucket, in tokens per second
+  # @param leak_rate[Float] the leak rate of the bucket, in tokens per second.
+  #   Either `leak_rate` or `over_time` can be used, but not both.
+  # @param over_time[#to_f] over how many seconds the bucket will leak out to 0 tokens.
+  #   The value is assumed to be the number of seconds
+  #   - or a duration which returns the number of seconds from `to_f`.
+  #   Either `leak_rate` or `over_time` can be used, but not both.
   # @param capacity[Numeric] how many tokens is the bucket capped at.
   #   Filling up the bucket using `fillup()` will add to that number, but
   #   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:, leak_rate:, capacity:)
+  def initialize(key:, capacity:, 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
-    @leak_rate = leak_rate.to_f
     @capacity = capacity.to_f
   end
 
-  # Places `n` tokens in the bucket. Once tokens are placed, the bucket is set to expire
-  # within 2 times the time it would take it to leak to 0, regardless of how many tokens
-  # get put in - since the amount of tokens put in the bucket will always be capped
-  # to the `capacity:` value you pass to the constructor. Calling `fillup` also deletes
-  # leaky buckets which have expired.
+  # Places `n` tokens in the bucket. If the bucket has less capacity than `n` tokens, the bucket will be filled to capacity.
+  # If the bucket has less capacity than `n` tokens, it will be filled to capacity. If the bucket is already full
+  # when the fillup is requested, the bucket stays at capacity.
   #
-  # @param n_tokens[Float]
+  # Once tokens are placed, the bucket is set to expire within 2 times the time it would take it to leak to 0,
+  # regardless of how many tokens get put in - since the amount of tokens put in the bucket will always be capped
+  # to the `capacity:` value you pass to the constructor.
+  #
+  # @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, did_overflow = Pecorino.adapter.add_tokens(capacity: @capacity, key: @key, leak_rate: @leak_rate, n_tokens: n_tokens)
-    State.new(capped_level_after_fillup, did_overflow)
+    capped_level_after_fillup, is_full = Pecorino.adapter.add_tokens(capacity: @capacity, key: @key, leak_rate: @leak_rate, n_tokens: n_tokens)
+    State.new(capped_level_after_fillup, is_full)
+  end
+
+  # Places `n` tokens in the bucket. If the bucket has less capacity than `n` tokens, the fillup will be rejected.
+  # This can be used for "exactly once" semantics or just more precise rate limiting. Note that if the bucket has
+  # _exactly_ `n` tokens of capacity the fillup will be accepted.
+  #
+  # Once tokens are placed, the bucket is set to expire within 2 times the time it would take it to leak to 0,
+  # regardless of how many tokens get put in - since the amount of tokens put in the bucket will always be capped
+  # to the `capacity:` value you pass to the constructor.
+  #
+  # @example
+  #    withdrawals = LeakyBuket.new(key: "wallet-#{user.id}", capacity: 200, over_time: 1.day)
+  #    if withdrawals.fillup_conditionally(amount_to_withdraw).accepted?
+  #      user.wallet.withdraw(amount_to_withdraw)
+  #    else
+  #      raise "You need to wait a bit before withdrawing more"
+  #    end
+  # @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)
+    ConditionalFillupResult.new(capped_level_after_fillup, is_full, did_accept)
   end
 
   # Returns the current state of the bucket, containing the level and whether the bucket is full.
@@ -108,6 +153,6 @@ class Pecorino::LeakyBucket
   # @param n_tokens[Float]
   # @return [boolean]
   def able_to_accept?(n_tokens)
-    (state.level + n_tokens) < @capacity
+    (state.level + n_tokens) <= @capacity
   end
 end

data/lib/pecorino/postgres.rb

@@ -72,7 +72,7 @@ Pecorino::Postgres = Struct.new(:model_class) do
       RETURNING
         level,
         -- Compare level to the capacity inside the DB so that we won't have rounding issues
-        level >= :capacity AS did_overflow
+        level >= :capacity AS at_capacity
     SQL
 
     # Note the use of .uncached here. The AR query cache will actually see our
@@ -80,8 +80,67 @@ Pecorino::Postgres = Struct.new(:model_class) do
     # 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) }
-    capped_level_after_fillup, did_overflow = upserted.fetch("level"), upserted.fetch("did_overflow")
-    [capped_level_after_fillup, did_overflow]
+    capped_level_after_fillup, at_capacity = upserted.fetch("level"), upserted.fetch("at_capacity")
+    [capped_level_after_fillup, at_capacity]
+  end
+
+  def add_tokens_conditionally(key:, capacity:, leak_rate:, n_tokens:)
+    # Take double the time it takes the bucket to empty under normal circumstances
+    # until the bucket may be deleted.
+    may_be_deleted_after_seconds = (capacity.to_f / leak_rate.to_f) * 2.0
+
+    # Create the leaky bucket if it does not exist, and update
+    # to the new level, taking the leak rate into account - if the bucket exists.
+    query_params = {
+      key: key.to_s,
+      capacity: capacity.to_f,
+      delete_after_s: may_be_deleted_after_seconds,
+      leak_rate: leak_rate.to_f,
+      fillup: n_tokens.to_f
+    }
+
+    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,
+          -- then we also clamp the above + fillup to not go below 0
+          GREATEST(0.0, 
+              GREATEST(0.0, level - (EXTRACT(EPOCH FROM (clock_timestamp() - last_touched_at)) * :leak_rate)) + :fillup
+          ) AS level_post_with_uncapped_fillup,
+          GREATEST(0.0,
+              level - (EXTRACT(EPOCH FROM (clock_timestamp() - last_touched_at)) * :leak_rate)
+          ) AS level_post
+        FROM pecorino_leaky_buckets
+        WHERE key = :key
+      )
+      INSERT INTO pecorino_leaky_buckets AS t
+        (key, last_touched_at, may_be_deleted_after, level)
+      VALUES
+        (
+          :key,
+          clock_timestamp(),
+          clock_timestamp() + ':delete_after_s second'::interval,
+          GREATEST(0.0,
+            (CASE WHEN :fillup > :capacity THEN 0.0 ELSE :fillup END)
+          )
+        )
+      ON CONFLICT (key) DO UPDATE SET
+        last_touched_at = EXCLUDED.last_touched_at,
+        may_be_deleted_after = EXCLUDED.may_be_deleted_after,
+        level = CASE WHEN (SELECT level_post_with_uncapped_fillup FROM pre) <= :capacity THEN
+          (SELECT level_post_with_uncapped_fillup FROM pre)
+        ELSE
+          (SELECT level_post FROM pre)
+        END
+      RETURNING
+        COALESCE((SELECT level_post FROM pre), 0.0) AS level_before,
+        level AS level_after
+    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:)
@@ -90,17 +149,17 @@ Pecorino::Postgres = Struct.new(:model_class) do
       INSERT INTO pecorino_blocks AS t
         (key, blocked_until)
       VALUES
-        (:key, NOW() + ':block_for seconds'::interval)
+        (:key, clock_timestamp() + ':block_for seconds'::interval)
       ON CONFLICT (key) DO UPDATE SET
         blocked_until = GREATEST(EXCLUDED.blocked_until, t.blocked_until)
-      RETURNING blocked_until;
+      RETURNING blocked_until
     SQL
     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])
-      SELECT blocked_until FROM pecorino_blocks WHERE key = ? AND blocked_until >= NOW() LIMIT 1
+      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) }
   end

data/lib/pecorino/sqlite.rb

@@ -94,6 +94,74 @@ Pecorino::Sqlite = Struct.new(:model_class) do
     [capped_level_after_fillup, one_if_did_overflow == 1]
   end
 
+  def add_tokens_conditionally(key:, capacity:, leak_rate:, n_tokens:)
+    # Take double the time it takes the bucket to empty under normal circumstances
+    # until the bucket may be deleted.
+    may_be_deleted_after_seconds = (capacity.to_f / leak_rate.to_f) * 2.0
+
+    # Create the leaky bucket if it does not exist, and update
+    # to the new level, taking the leak rate into account - if the bucket exists.
+    query_params = {
+      key: key.to_s,
+      capacity: capacity.to_f,
+      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
+    }
+
+    # 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 INTO pecorino_leaky_buckets AS t
+        (id, 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
+          0.0
+        )
+      ON CONFLICT (key) DO UPDATE SET
+      -- Make sure we extend the lifetime of the row
+      -- 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)
+
+    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
+          -- Note the double clamping here. First we clamp the "current level - leak" to not go below zero,
+          -- then we also clamp the above + fillup to not go below 0
+          MAX(0.0, MAX(0.0, level - ((:now_s - last_touched_at) * :leak_rate)) + :fillup) AS level_post_with_uncapped_fillup,
+          MAX(0.0, level - ((:now_s - last_touched_at) * :leak_rate)) AS level_post
+        FROM
+          pecorino_leaky_buckets
+        WHERE key = :key
+      ) UPDATE pecorino_leaky_buckets SET
+        last_touched_at = :now_s,
+        may_be_deleted_after = DATETIME('now', '+:delete_after_s seconds'),
+        level = CASE WHEN (SELECT level_post_with_uncapped_fillup FROM pre) <= :capacity THEN
+          (SELECT level_post_with_uncapped_fillup FROM pre)
+        ELSE
+          (SELECT level_post FROM pre)
+        END
+      RETURNING
+        (SELECT level_post FROM pre) AS level_before,
+        level AS level_after
+    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])

data/lib/pecorino/throttle.rb

@@ -14,6 +14,10 @@ class Pecorino::Throttle
       blocked_until ? true : false
     end
 
+    # Returns the number of seconds until the block will be lifted, rouded up to the closest
+    # whole second. This value can be used in a "Retry-After" HTTP response header.
+    #
+    # @return [Integer]
     def retry_after
       (blocked_until - Time.now.utc).ceil
     end
@@ -23,11 +27,14 @@ class Pecorino::Throttle
     # Returns the throttle which raised the exception. Can be used to disambiguiate between
     # multiple Throttled exceptions when multiple throttles are applied in a layered fashion:
     #
+    # @example
+    #    begin
     #      ip_addr_throttle.request!
     #      user_email_throttle.request!
     #      db_insert_throttle.request!(n_items_to_insert)
     #    rescue Pecorino::Throttled => e
     #      deliver_notification(user) if e.throttle == user_email_throttle
+    #    end
     #
     # @return [Throttle]
     attr_reader :throttle
@@ -43,19 +50,21 @@ class Pecorino::Throttle
   end
 
   # @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
+  # @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 leaky_bucket_options Options for `Pecorino::LeakyBucket.new`
   # @see PecorinoLeakyBucket.new
-  def initialize(key:, block_for: 30, **)
-    @key = key.to_s
-    @block_for = block_for.to_f
+  def initialize(key:, block_for: nil, **)
     @bucket = Pecorino::LeakyBucket.new(key:, **)
+    @key = key.to_s
+    @block_for = block_for ? block_for.to_f : (@bucket.capacity / @bucket.leak_rate)
   end
 
   # Tells whether the throttle will let this number of requests pass without raising
   # a Throttled. Note that this is not race-safe. Another request could overflow the bucket
   # after you call `able_to_accept?` but before you call `throttle!`. So before performing
-  # the action you still need to call `throttle!`
+  # the action you still need to call `throttle!`. You may still use `able_to_accept?` to
+  # provide better UX to your users before they cause an action that would otherwise throttle.
   #
   # @param n_tokens[Float]
   # @return [boolean]
@@ -70,10 +79,13 @@ class Pecorino::Throttle
   # The exception can be rescued later to provide a 429 response. This method is better
   # to use before performing the unit of work that the throttle is guarding:
   #
-  # @example      t.request!
-  #               Note.create!(note_params)
-  #            rescue Pecorino::Throttle::Throttled => e
-  #               [429, {"Retry-After" => e.retry_after.to_s}, []]
+  # @example
+  #   begin
+  #      t.request!
+  #      Note.create!(note_params)
+  #   rescue Pecorino::Throttle::Throttled => e
+  #      [429, {"Retry-After" => e.retry_after.to_s}, []]
+  #   end
   #
   # If the method call succeeds it means that the request is not getting throttled.
   #
@@ -81,30 +93,33 @@ class Pecorino::Throttle
   def request!(n = 1)
     state = request(n)
     raise Throttled.new(self, state) if state.blocked?
+    nil
   end
 
   # Register that a request is being performed. Will not raise any exceptions but return
   # the time at which the block will be lifted if a block resulted from this request or
   # was already in effect. Can be used for registering actions which already took place,
-  # but should result in subsequent actions being blocked in subsequent requests later.
+  # but should result in subsequent actions being blocked.
   #
-  # @example    unless t.able_to_accept?
-  #       Note.create!(note_params)
-  #       t.request
-  #     else
-  #       raise "Throttled or block in effect"
-  #     end
+  # @example
+  #   if t.able_to_accept?
+  #     Entry.create!(entry_params)
+  #     t.request
+  #   end
   #
   # @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)
     return State.new(existing_blocked_until.utc) if existing_blocked_until
 
-    # Topup the leaky bucket
-    return State.new(nil) unless @bucket.fillup(n.to_f).full?
-
-    # and set the block if we reached it
-    fresh_blocked_until = Pecorino.adapter.set_block(key: @key, block_for: @block_for)
-    State.new(fresh_blocked_until.utc)
+    # Topup the leaky bucket, and if the topup gets rejected - block the caller
+    fillup = @bucket.fillup_conditionally(n)
+    if fillup.accepted?
+      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)
+      State.new(fresh_blocked_until.utc)
+    end
   end
 end

data/lib/pecorino/version.rb

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

metadata

@@ -1,14 +1,14 @@
 --- !ruby/object:Gem::Specification
 name: pecorino
 version: !ruby/object:Gem::Version
-  version: 0.2.0
+  version: 0.4.0
 platform: ruby
 authors:
 - Julik Tarkhanov
 autorequire:
 bindir: exe
 cert_chain: []
-date: 2024-01-09 00:00:00.000000000 Z
+date: 2024-01-22 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: activerecord