bugfix(token-bucket): Check CurrentTokens Against Cost Variable (#119)

lib/hammer/atomic/token_bucket.ex

@@ -112,7 +112,7 @@ defmodule Hammer.Atomic.TokenBucket do
 
         current_tokens = min(capacity, current_fill + new_tokens)
 
-        if current_tokens >= 1 do
+        if current_tokens >= cost do
           final_level = current_tokens - cost
 
           :atomics.exchange(atomic, 1, final_level)

lib/hammer/ets/token_bucket.ex

@@ -110,7 +110,7 @@ defmodule Hammer.ETS.TokenBucket do
 
     current_tokens = min(capacity, current_level + new_tokens)
 
-    if current_tokens >= 1 do
+    if current_tokens >= cost do
       final_level = current_tokens - cost
       :ets.insert(table, {key, final_level, now})
       {:allow, final_level}

test/hammer/atomic/token_bucket_test.exs

@@ -62,6 +62,24 @@ defmodule Hammer.Atomic.TokenBucketTest do
       assert {:deny, _retry_after} =
                TokenBucket.hit(table, key, refill_rate, capacity, 1)
     end
+
+    test "handles costs greater than 1 correctly", %{table: table} do
+      key = "key"
+      refill_rate = 2
+      capacity = 10
+
+      # First hit with cost of 3 should succeed and leave 7 tokens
+      assert {:allow, 7} = TokenBucket.hit(table, key, refill_rate, capacity, 3)
+
+      # Second hit with cost of 4 should succeed and leave 3 tokens
+      assert {:allow, 3} = TokenBucket.hit(table, key, refill_rate, capacity, 4)
+
+      # Third hit with cost of 4 should be denied (only 3 tokens left)
+      assert {:deny, _retry_after} = TokenBucket.hit(table, key, refill_rate, capacity, 4)
+
+      # Small cost of 2 should still succeed since we have 3 tokens
+      assert {:allow, 1} = TokenBucket.hit(table, key, refill_rate, capacity, 2)
+    end
   end
 
   describe "get" do

test/hammer/ets/token_bucket_test.exs

@@ -62,6 +62,24 @@ defmodule Hammer.ETS.TokenBucketTest do
       assert {:deny, _retry_after} =
                TokenBucket.hit(table, key, refill_rate, capacity, 1)
     end
+
+    test "handles costs greater than 1 correctly", %{table: table} do
+      key = "key"
+      refill_rate = 2
+      capacity = 10
+
+      # First hit with cost of 3 should succeed and leave 7 tokens
+      assert {:allow, 7} = TokenBucket.hit(table, key, refill_rate, capacity, 3)
+
+      # Second hit with cost of 4 should succeed and leave 3 tokens
+      assert {:allow, 3} = TokenBucket.hit(table, key, refill_rate, capacity, 4)
+
+      # Third hit with cost of 4 should be denied (only 3 tokens left)
+      assert {:deny, _retry_after} = TokenBucket.hit(table, key, refill_rate, capacity, 4)
+
+      # Small cost of 2 should still succeed since we have 3 tokens
+      assert {:allow, 1} = TokenBucket.hit(table, key, refill_rate, capacity, 2)
+    end
   end
 
   describe "get" do