Fix bug that allows check_rate to succeed 2x in sequence even if the bucket limit is 1 attempt per hour

lib/hammer.ex

@@ -48,7 +48,7 @@ defmodule Hammer do
   def check_rate(backend, id, scale_ms, limit) do
     {stamp, key} = Utils.stamp_key(id, scale_ms)
 
-    case call_backend(backend, :count_hit, [key, stamp]) do
+    case call_backend(backend, :count_hit, [key, scale_ms, stamp]) do
       {:ok, count} ->
         if count > limit do
           {:deny, limit}
@@ -95,7 +95,7 @@ defmodule Hammer do
   def check_rate_inc(backend, id, scale_ms, limit, increment) do
     {stamp, key} = Utils.stamp_key(id, scale_ms)
 
-    case call_backend(backend, :count_hit, [key, stamp, increment]) do
+    case call_backend(backend, :count_hit, [key, scale_ms, stamp, increment]) do
       {:ok, count} ->
         if count > limit do
           {:deny, limit}
@@ -110,8 +110,8 @@ defmodule Hammer do
 
   @spec inspect_bucket(id :: String.t(), scale_ms :: integer, limit :: integer) ::
           {:ok,
-           {count :: integer, count_remaining :: integer, ms_to_next_bucket :: integer,
-            created_at :: integer | nil, updated_at :: integer | nil}}
+           {count :: integer, count_remaining :: integer, ms_to_next_bucket :: integer, created_at :: integer | nil,
+            updated_at :: integer | nil}}
           | {:error, reason :: any}
   @doc """
   Inspect bucket to get count, count_remaining, ms_to_next_bucket, created_at,
@@ -143,22 +143,25 @@ defmodule Hammer do
 
   @spec inspect_bucket(backend :: atom, id :: String.t(), scale_ms :: integer, limit :: integer) ::
           {:ok,
-           {count :: integer, count_remaining :: integer, ms_to_next_bucket :: integer,
-            created_at :: integer | nil, updated_at :: integer | nil}}
+           {count :: integer, count_remaining :: integer, ms_to_next_bucket :: integer, created_at :: integer | nil,
+            updated_at :: integer | nil}}
           | {:error, reason :: any}
   @doc """
   Same as inspect_bucket/3, but allows specifying a backend
   """
   def inspect_bucket(backend, id, scale_ms, limit) do
-    {stamp, key} = Utils.stamp_key(id, scale_ms)
-    ms_to_next_bucket = elem(key, 0) * scale_ms + scale_ms - stamp
+    {_stamp, key} = Utils.stamp_key(id, scale_ms)
 
     case call_backend(backend, :get_bucket, [key]) do
       {:ok, nil} ->
-        {:ok, {0, limit, ms_to_next_bucket, nil, nil}}
+        {:ok, {0, limit, scale_ms, nil, nil}}
 
       {:ok, {_, count, created_at, updated_at}} ->
+        now = Utils.timestamp()
+        time_passed = now - created_at
+        ms_to_next_bucket = max(scale_ms - time_passed, 0)
         count_remaining = if limit > count, do: limit - count, else: 0
+
         {:ok, {count, count_remaining, ms_to_next_bucket, created_at, updated_at}}
 
       {:error, reason} ->

lib/hammer/backend.ex

@@ -10,6 +10,7 @@ defmodule Hammer.Backend do
   @callback count_hit(
               pid :: pid(),
               key :: bucket_key,
+              scale_ms :: integer,
               now :: integer
             ) ::
               {:ok, count :: integer}
@@ -18,6 +19,7 @@ defmodule Hammer.Backend do
   @callback count_hit(
               pid :: pid(),
               key :: bucket_key,
+              scale_ms :: integer,
               now :: integer,
               increment :: integer
             ) ::

lib/hammer/backend/ets.ex

@@ -30,6 +30,7 @@ defmodule Hammer.Backend.ETS do
   @behaviour Hammer.Backend
 
   use GenServer
+
   alias Hammer.Utils
 
   @type bucket_key :: {bucket :: integer, id :: String.t()}
@@ -68,12 +69,13 @@ defmodule Hammer.Backend.ETS do
   @spec count_hit(
           pid :: pid(),
           key :: bucket_key,
+          scale_ms :: integer,
           now :: integer
         ) ::
           {:ok, count :: integer}
           | {:error, reason :: any}
-  def count_hit(pid, key, now) do
-    count_hit(pid, key, now, 1)
+  def count_hit(pid, key, scale_ms, now) do
+    count_hit(pid, key, scale_ms, now, 1)
   end
 
   @doc """
@@ -82,12 +84,13 @@ defmodule Hammer.Backend.ETS do
   @spec count_hit(
           pid :: pid(),
           key :: bucket_key,
+          scale_ms :: integer,
           now :: integer,
           increment :: integer
         ) ::
           {:ok, count :: integer}
           | {:error, reason :: any}
-  def count_hit(_pid, key, now, increment) do
+  def count_hit(pid, key, scale_ms, now, increment) do
     if :ets.member(@ets_table_name, key) do
       [count, _] =
         :ets.update_counter(@ets_table_name, key, [
@@ -101,6 +104,7 @@ defmodule Hammer.Backend.ETS do
     else
       # Insert {key, count, created_at, updated_at}
       true = :ets.insert(@ets_table_name, {key, increment, now, now})
+      Process.send_after(pid, {:delete_bucket, key}, scale_ms)
       {:ok, increment}
     end
   rescue
@@ -208,4 +212,15 @@ defmodule Hammer.Backend.ETS do
 
     {:noreply, state}
   end
+
+  def handle_info({:delete_bucket, {_bucket, id}}, state) do
+    :ets.select_delete(@ets_table_name, [
+      {{{:"$1", :"$2"}, :_, :_, :_}, [{:==, :"$2", id}], [true]}
+    ])
+
+    {:noreply, state}
+  rescue
+    _e ->
+      {:noreply, state}
+  end
 end

lib/hammer/utils.ex

@@ -17,9 +17,7 @@ defmodule Hammer.Utils do
   # Returns tuple of {timestamp, key}, where key is {bucket_number, id}
   def stamp_key(id, scale_ms) do
     stamp = timestamp()
-    # with scale_ms = 1 bucket changes every millisecond
-    bucket_number = trunc(stamp / scale_ms)
-    key = {bucket_number, id}
+    key = {scale_ms, id}
     {stamp, key}
   end
 

test/hammer_ets_test.exs

@@ -23,9 +23,9 @@ defmodule ETSTest do
   test "count_hit", context do
     pid = context[:pid]
     {stamp, key} = Utils.stamp_key("one", 200_000)
-    assert {:ok, 1} = ETS.count_hit(pid, key, stamp)
-    assert {:ok, 2} = ETS.count_hit(pid, key, stamp)
-    assert {:ok, 3} = ETS.count_hit(pid, key, stamp)
+    assert {:ok, 1} = ETS.count_hit(pid, key, 200_000, stamp)
+    assert {:ok, 2} = ETS.count_hit(pid, key, 200_000, stamp)
+    assert {:ok, 3} = ETS.count_hit(pid, key, 200_000, stamp)
   end
 
   test "get_bucket", context do
@@ -34,10 +34,10 @@ defmodule ETSTest do
     # With no hits
     assert {:ok, nil} = ETS.get_bucket(pid, key)
     # With one hit
-    assert {:ok, 1} = ETS.count_hit(pid, key, stamp)
+    assert {:ok, 1} = ETS.count_hit(pid, key, 200_000, stamp)
     assert {:ok, {{_, "two"}, 1, _, _}} = ETS.get_bucket(pid, key)
     # With two hits
-    assert {:ok, 2} = ETS.count_hit(pid, key, stamp)
+    assert {:ok, 2} = ETS.count_hit(pid, key, 200_000, stamp)
     assert {:ok, {{_, "two"}, 2, _, _}} = ETS.get_bucket(pid, key)
   end
 
@@ -47,17 +47,17 @@ defmodule ETSTest do
     # With no hits
     assert {:ok, 0} = ETS.delete_buckets(pid, "three")
     # With three hits in same bucket
-    assert {:ok, 1} = ETS.count_hit(pid, key, stamp)
-    assert {:ok, 2} = ETS.count_hit(pid, key, stamp)
-    assert {:ok, 3} = ETS.count_hit(pid, key, stamp)
+    assert {:ok, 1} = ETS.count_hit(pid, key, 200_000, stamp)
+    assert {:ok, 2} = ETS.count_hit(pid, key, 200_000, stamp)
+    assert {:ok, 3} = ETS.count_hit(pid, key, 200_000, stamp)
     assert {:ok, 1} = ETS.delete_buckets(pid, "three")
   end
 
   test "timeout pruning", context do
     pid = context[:pid]
     expiry_ms = context[:expiry_ms]
     {stamp, key} = Utils.stamp_key("something-pruned", 200_000)
-    assert {:ok, 1} = ETS.count_hit(pid, key, stamp)
+    assert {:ok, 1} = ETS.count_hit(pid, key, 200_000, stamp)
     assert {:ok, {{_, "something-pruned"}, 1, _, _}} = ETS.get_bucket(pid, key)
     :timer.sleep(expiry_ms * 5)
     assert {:ok, nil} = ETS.get_bucket(pid, key)

test/hammer_test.exs

@@ -84,18 +84,20 @@ defmodule HammerTest do
   end
 
   test "returns expected tuples on inspect_bucket" do
-    assert {:ok, {0, 2, _, nil, nil}} = Hammer.inspect_bucket("inspect_bucket11", 1000, 2)
-    assert {:allow, 1} = Hammer.check_rate("inspect_bucket11", 1000, 2)
-    assert {:ok, {1, 1, _, _, _}} = Hammer.inspect_bucket("inspect_bucket11", 1000, 2)
-    assert {:allow, 2} = Hammer.check_rate("inspect_bucket11", 1000, 2)
-    assert {:allow, 1} = Hammer.check_rate("inspect_bucket22", 1000, 2)
-    assert {:ok, {2, 0, _, _, _}} = Hammer.inspect_bucket("inspect_bucket11", 1000, 2)
-    assert {:deny, 2} = Hammer.check_rate("inspect_bucket11", 1000, 2)
+    assert {:ok, {0, 2, _, nil, nil}} = Hammer.inspect_bucket("inspect_bucket11", 10_000, 2)
+    assert {:allow, 1} = Hammer.check_rate("inspect_bucket11", 10_000, 2)
+    assert {:ok, {1, 1, _, _, _}} = Hammer.inspect_bucket("inspect_bucket11", 10_000, 2)
+    assert {:allow, 2} = Hammer.check_rate("inspect_bucket11", 10_000, 2)
+    assert {:allow, 1} = Hammer.check_rate("inspect_bucket22", 10_000, 2)
+    assert {:ok, {2, 0, _, _, _}} = Hammer.inspect_bucket("inspect_bucket11", 10_000, 2)
+    assert {:deny, 2} = Hammer.check_rate("inspect_bucket11", 10_000, 2)
+
+    :timer.sleep(1)
 
     assert {:ok, {3, 0, ms_to_next_bucket, _, _}} =
-             Hammer.inspect_bucket("inspect_bucket11", 1000, 2)
+             Hammer.inspect_bucket("inspect_bucket11", 10_000, 2)
 
-    assert ms_to_next_bucket < 1000
+    assert ms_to_next_bucket < 10_000
   end
 
   test "returns expected tuples on delete_buckets" do
@@ -126,4 +128,24 @@ defmodule HammerTest do
     assert {:allow, 10} = Hammer.check_rate("cost-bucket2", 1000, 10)
     assert {:deny, 10} = Hammer.check_rate_inc("cost-bucket2", 1000, 10, 2)
   end
+
+  test "ms_to_next_bucket should be equal to scale_ms on first check_rate", %{bucket: bucket} do
+    scale_ms = :timer.seconds(20)
+    assert {:ok, {0, 2, _, nil, nil}} = Hammer.inspect_bucket(bucket, scale_ms, 2)
+    assert {:allow, 1} = Hammer.check_rate(bucket, scale_ms, 2)
+    assert {:ok, {1, 1, ms_to_next_bucket, _created, _updated}} = Hammer.inspect_bucket(bucket, scale_ms, 2)
+
+    assert ms_to_next_bucket > :timer.seconds(19)
+    assert ms_to_next_bucket <= :timer.seconds(20)
+  end
+
+  test "ms_to_next_bucket should be equal to scale_ms on first check_rate_inc", %{bucket: bucket} do
+    scale_ms = :timer.seconds(20)
+    assert {:ok, {0, 2, _, nil, nil}} = Hammer.inspect_bucket(bucket, scale_ms, 2)
+    assert {:allow, 2} = Hammer.check_rate_inc(bucket, scale_ms, 2, 2)
+    assert {:ok, {2, 0, ms_to_next_bucket, _created, _updated}} = Hammer.inspect_bucket(bucket, scale_ms, 2)
+
+    assert ms_to_next_bucket > :timer.seconds(19)
+    assert ms_to_next_bucket <= :timer.seconds(20)
+  end
 end