Fix race condition between insert_new and lookup in ETS hit/5

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)
 

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)

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"

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"