From 35fd837ce14f789eb82fd9858250ff5ea070aecb Mon Sep 17 00:00:00 2001
From: Julik Tarkhanov <me@julik.nl>
Date: Tue, 12 Mar 2024 15:28:37 +0100
Subject: [PATCH] Add naked Block

for blocking for an arbitrary timespan
---
 CHANGELOG.md             | 16 ++++++++++------
 lib/pecorino.rb          |  7 ++++---
 lib/pecorino/block.rb    | 24 ++++++++++++++++++++++++
 lib/pecorino/throttle.rb |  4 ++--
 lib/pecorino/version.rb  |  2 +-
 test/block_test.rb       | 31 +++++++++++++++++++++++++++++++
 6 files changed, 72 insertions(+), 12 deletions(-)
 create mode 100644 lib/pecorino/block.rb
 create mode 100644 test/block_test.rb

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 00786a6..2df074c 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,4 +1,8 @@
-## [0.5.0] - 2024-02-11
+## 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 +10,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 +25,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
diff --git a/lib/pecorino.rb b/lib/pecorino.rb
index 3250ae8..00eea4f 100644
--- a/lib/pecorino.rb
+++ b/lib/pecorino.rb
@@ -4,14 +4,15 @@
 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"
 
   # Deletes stale leaky buckets and blocks which have expired. Run this method regularly to
   # avoid accumulating too many unused rows in your tables.
diff --git a/lib/pecorino/block.rb b/lib/pecorino/block.rb
new file mode 100644
index 0000000..66a1c92
--- /dev/null
+++ b/lib/pecorino/block.rb
@@ -0,0 +1,24 @@
+# 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
+  # @return [Time] the time when the block will be released
+  def self.set!(key:, block_for:)
+    Pecorino.adapter.set_block(key: key, block_for: block_for)
+    Time.now + block_for
+  end
+
+  # Returns the time until a certain block is in effect
+  #
+  # @return [Time,nil] the time when the block will be released
+  def self.blocked_until(key:)
+    t = Pecorino.adapter.blocked_until(key: key)
+    (t && t > Time.now) ? t : nil
+  end
+end
diff --git a/lib/pecorino/throttle.rb b/lib/pecorino/throttle.rb
index 99edfed..a9ce060 100644
--- a/lib/pecorino/throttle.rb
+++ b/lib/pecorino/throttle.rb
@@ -156,7 +156,7 @@ def request!(n = 1)
   #
   # @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)
     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 +165,7 @@ def request(n = 1)
       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)
       State.new(fresh_blocked_until.utc)
     end
   end
diff --git a/lib/pecorino/version.rb b/lib/pecorino/version.rb
index f856fb4..2ef89ea 100644
--- a/lib/pecorino/version.rb
+++ b/lib/pecorino/version.rb
@@ -1,5 +1,5 @@
 # frozen_string_literal: true
 
 module Pecorino
-  VERSION = "0.5.0"
+  VERSION = "0.6.0"
 end
diff --git a/test/block_test.rb b/test/block_test.rb
new file mode 100644
index 0000000..f8b6c9e
--- /dev/null
+++ b/test/block_test.rb
@@ -0,0 +1,31 @@
+# frozen_string_literal: true
+
+require "test_helper"
+
+class BlockTest < ActiveSupport::TestCase
+  def setup
+    create_postgres_database
+  end
+
+  def teardown
+    drop_postgres_database
+  end
+
+  test "sets a block" do
+    k = Base64.strict_encode64(Random.bytes(4))
+    assert_nil Pecorino::Block.blocked_until(key: k)
+    assert Pecorino::Block.set!(key: k, block_for: 30.minutes)
+
+    blocked_until = Pecorino::Block.blocked_until(key: k)
+    assert_in_delta Time.now + 30.minutes, blocked_until, 10
+  end
+
+  test "does not return a block which has lapsed" do
+    k = Base64.strict_encode64(Random.bytes(4))
+    assert_nil Pecorino::Block.blocked_until(key: k)
+    assert Pecorino::Block.set!(key: k, block_for: -30.minutes)
+
+    blocked_until = Pecorino::Block.blocked_until(key: k)
+    assert_nil blocked_until
+  end
+end