From 71d416ada09b0b423add548fccb240382b766615 Mon Sep 17 00:00:00 2001
From: Rhet <rhetprieto@gmail.com>
Date: Tue, 20 Apr 2021 16:20:13 -0600
Subject: [PATCH] Add new method to only increase counter

---
 CHANGELOG.md                     |  4 ++++
 README.md                        | 13 +++++++++++++
 lib/rollie/rate_limiter.rb       | 14 ++++++++++++--
 spec/rollie/rate_limiter_spec.rb | 23 +++++++++++++++++++++++
 4 files changed, 52 insertions(+), 2 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index ce61870..fa542e4 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,5 +1,9 @@
 # Changelog
 
+## 0.1.2
+
+* Add new method to only increase counter and decide what happens next
+
 ## 0.1.1
 
 * Remove maximum dependency versions
diff --git a/README.md b/README.md
index 0aaf6b4..6eacbb2 100644
--- a/README.md
+++ b/README.md
@@ -70,6 +70,19 @@ end
 The status will tell you the current state. You can also see the current count
 and how long until the bucket resets.
 
+If you don't want to block the request but control what happens next:
+
+```
+# limit 30 requests per second.
+twitter_rate = Rollie::RateLimiter.new("twitter_requests", limit: 30, interval: 1000)
+status = twitter_rate.increase_counter # This will always count blocked regardless on how you initialized the class
+if status.exceeded?
+  # Do Something here
+else 
+  # Do something different
+end
+```
+
 ```ruby
 status.exceeded?
 # => false
diff --git a/lib/rollie/rate_limiter.rb b/lib/rollie/rate_limiter.rb
index a0de9ae..ef371e7 100644
--- a/lib/rollie/rate_limiter.rb
+++ b/lib/rollie/rate_limiter.rb
@@ -34,6 +34,15 @@ def within_limit
       end
     end
 
+    # Increase counter and let you decide what to do next. Will always count blocked
+    #
+    # @return [Status] The current status for this RateLimiter.
+    def increase_counter
+      Rollie.redis do |conn|
+        inc(conn, true)
+      end
+    end
+
     # @return [Integer] The current count of this RateLimiter.
     def count
       Rollie.redis do |conn|
@@ -44,7 +53,8 @@ def count
 
     private
 
-    def inc(conn)
+    def inc(conn, count_blocked = nil)
+      count_blocked = count_blocked.nil? ? @count_blocked : count_blocked
       time = (Time.now.to_r * 1_000_000).round
       old = time - @interval
       range = conn.multi do
@@ -58,7 +68,7 @@ def inc(conn)
       current_count = range.length
       time_remaining = range.first.to_i - time + @interval
 
-      if exceeded && !@count_blocked
+      if exceeded && !count_blocked
         conn.zremrangebyscore(@key, time, time)
         current_count -= 1
       end
diff --git a/spec/rollie/rate_limiter_spec.rb b/spec/rollie/rate_limiter_spec.rb
index 01bc231..473d96b 100644
--- a/spec/rollie/rate_limiter_spec.rb
+++ b/spec/rollie/rate_limiter_spec.rb
@@ -85,6 +85,29 @@ module Rollie
       end
     end
 
+    describe "#increase_counter" do
+      before do
+        @limiter = RateLimiter.new(SecureRandom.hex(8), count_blocked: false)
+      end
+
+      it "should return status" do
+        status = @limiter.increase_counter
+        expect(status.count).to eq(1)
+        expect(status.exceeded?).to be(false)
+        expect(status.time_remaining).to eq(1000)
+      end
+
+      it "should increase count regardless of count_blocked setting" do
+        35.times do
+          @limiter.increase_counter
+        end
+
+        status = @limiter.increase_counter
+        expect(status.count).to eq(36)
+        expect(status.exceeded?).to be(true)
+      end
+    end
+
     describe '#count' do
       it 'returns the current count' do
         30.times do