diff --git a/lib/hammer.ex b/lib/hammer.ex index 6015872..50ca72d 100644 --- a/lib/hammer.ex +++ b/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} -> diff --git a/lib/hammer/backend.ex b/lib/hammer/backend.ex index 81cf0f5..9a717f2 100644 --- a/lib/hammer/backend.ex +++ b/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 ) :: diff --git a/lib/hammer/backend/ets.ex b/lib/hammer/backend/ets.ex index 3f3e044..c734a53 100644 --- a/lib/hammer/backend/ets.ex +++ b/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 diff --git a/lib/hammer/utils.ex b/lib/hammer/utils.ex index f226cb0..467c2bf 100644 --- a/lib/hammer/utils.ex +++ b/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 diff --git a/test/hammer_ets_test.exs b/test/hammer_ets_test.exs index 787763a..1dcd288 100644 --- a/test/hammer_ets_test.exs +++ b/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,9 +47,9 @@ 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 @@ -57,7 +57,7 @@ defmodule ETSTest 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) diff --git a/test/hammer_test.exs b/test/hammer_test.exs index d07177f..d0cee28 100644 --- a/test/hammer_test.exs +++ b/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