diff --git a/lib/hammer/ets/leaky_bucket.ex b/lib/hammer/ets/leaky_bucket.ex index 441b190..224533c 100644 --- a/lib/hammer/ets/leaky_bucket.ex +++ b/lib/hammer/ets/leaky_bucket.ex @@ -117,7 +117,16 @@ defmodule Hammer.ETS.LeakyBucket do :ets.insert_new(table, {key, 0, now}) # Get current bucket state - [{^key, current_fill, last_update}] = :ets.lookup(table, key) + {current_fill, last_update} = + case :ets.lookup(table, key) do + [{^key, current_fill, last_update}] -> + {current_fill, last_update} + + [] -> + # Entry was deleted between insert_new and lookup (cleanup race or table restart) + :ets.insert(table, {key, 0, now}) + {0, now} + end leaked = trunc((now - last_update) * leak_rate) diff --git a/lib/hammer/ets/token_bucket.ex b/lib/hammer/ets/token_bucket.ex index fc0d2fb..b2d1fa8 100644 --- a/lib/hammer/ets/token_bucket.ex +++ b/lib/hammer/ets/token_bucket.ex @@ -119,7 +119,17 @@ defmodule Hammer.ETS.TokenBucket do # Try to insert new empty bucket if doesn't exist :ets.insert_new(table, {key, capacity, now}) - [{^key, current_level, last_update}] = :ets.lookup(table, key) + {current_level, last_update} = + case :ets.lookup(table, key) do + [{^key, current_level, last_update}] -> + {current_level, last_update} + + [] -> + # Entry was deleted between insert_new and lookup (cleanup race or table restart) + :ets.insert(table, {key, capacity, now}) + {capacity, now} + end + new_tokens = trunc((now - last_update) * refill_rate) current_tokens = min(capacity, current_level + new_tokens) diff --git a/test/hammer/ets/leaky_bucket_test.exs b/test/hammer/ets/leaky_bucket_test.exs index 77a1931..810881a 100644 --- a/test/hammer/ets/leaky_bucket_test.exs +++ b/test/hammer/ets/leaky_bucket_test.exs @@ -66,6 +66,32 @@ defmodule Hammer.ETS.LeakyBucketTest do end end + describe "race condition handling" do + test "hit recovers when entry is deleted between insert_new and lookup", %{table: table} do + key = "race_key" + leak_rate = 10 + capacity = 10 + + # Insert an entry, then delete it to simulate cleanup race + :ets.insert(table, {key, 5, System.system_time(:second)}) + :ets.delete(table, key) + + # hit should handle the missing entry gracefully + assert {:allow, 1} = LeakyBucket.hit(table, key, leak_rate, capacity, 1) + end + + test "hit works on a fresh empty table", %{table: table} do + key = "fresh_key" + leak_rate = 10 + capacity = 10 + + # Ensure key doesn't exist + assert :ets.lookup(table, key) == [] + + assert {:allow, 1} = LeakyBucket.hit(table, key, leak_rate, capacity, 1) + end + end + describe "get" do test "get returns current bucket level", %{table: table} do key = "key" diff --git a/test/hammer/ets/token_bucket_test.exs b/test/hammer/ets/token_bucket_test.exs index 4fa1243..b299674 100644 --- a/test/hammer/ets/token_bucket_test.exs +++ b/test/hammer/ets/token_bucket_test.exs @@ -82,6 +82,32 @@ defmodule Hammer.ETS.TokenBucketTest do end end + describe "race condition handling" do + test "hit recovers when entry is deleted between insert_new and lookup", %{table: table} do + key = "race_key" + refill_rate = 10 + capacity = 10 + + # Insert an entry, then delete it to simulate cleanup race + :ets.insert(table, {key, 5, System.system_time(:second)}) + :ets.delete(table, key) + + # hit should handle the missing entry gracefully + assert {:allow, 9} = TokenBucket.hit(table, key, refill_rate, capacity, 1) + end + + test "hit works on a fresh empty table", %{table: table} do + key = "fresh_key" + refill_rate = 10 + capacity = 10 + + # Ensure key doesn't exist + assert :ets.lookup(table, key) == [] + + assert {:allow, 9} = TokenBucket.hit(table, key, refill_rate, capacity, 1) + end + end + describe "get" do test "get returns current bucket level", %{table: table} do key = "key"