Add expires_at/2 API for fix_window algorithm

lib/hammer.ex

@@ -68,7 +68,15 @@ defmodule Hammer do
   """
   @callback get(key, scale) :: count
 
-  @optional_callbacks hit: 4, inc: 2, inc: 3, set: 3, get: 2
+  @doc """
+  Optional callback for getting the expiration time of the current window for a key.
+
+  Returns the expiration timestamp in milliseconds, or `0` if the key has no active window.
+  Only available for the `:fix_window` algorithm.
+  """
+  @callback expires_at(key, scale) :: non_neg_integer()
+
+  @optional_callbacks hit: 4, inc: 2, inc: 3, set: 3, get: 2, expires_at: 2
 
   @doc """
   Use the Hammer library in a module to create a rate limiter.

lib/hammer/atomic.ex

@@ -150,6 +150,12 @@ defmodule Hammer.Atomic do
           @algorithm.get(@table, key)
         end
       end
+
+      if function_exported?(@algorithm, :expires_at, 3) do
+        def expires_at(key, scale) do
+          @algorithm.expires_at(@table, key, scale)
+        end
+      end
     end
   end
 

lib/hammer/atomic/fix_window.ex

@@ -186,6 +186,22 @@ defmodule Hammer.Atomic.FixWindow do
     end
   end
 
+  @doc """
+  Returns the expiration time (in milliseconds) of the current window for a given key.
+
+  Returns `0` if the key has no active window.
+  """
+  @spec expires_at(table :: atom(), key :: term(), scale :: pos_integer()) :: non_neg_integer()
+  def expires_at(table, key, scale) do
+    window = div(Atomic.now(), scale)
+    full_key = {key, window}
+
+    case :ets.lookup(table, full_key) do
+      [{_, atomic}] -> :atomics.get(atomic, 2)
+      [] -> 0
+    end
+  end
+
   @doc false
   @spec normalize_entry(key :: term(), atomic :: reference()) :: map()
   def normalize_entry(key, atomic) do

lib/hammer/ets.ex

@@ -149,6 +149,12 @@ defmodule Hammer.ETS do
           @algorithm.get(@table, key)
         end
       end
+
+      if function_exported?(@algorithm, :expires_at, 3) do
+        def expires_at(key, scale) do
+          @algorithm.expires_at(@table, key, scale)
+        end
+      end
     end
   end
 

lib/hammer/ets/fix_window.ex

@@ -164,6 +164,22 @@ defmodule Hammer.ETS.FixWindow do
     :ets.select_delete(table, match_spec)
   end
 
+  @doc """
+  Returns the expiration time (in milliseconds) of the current window for a given key.
+
+  Returns `0` if the key has no active window.
+  """
+  @spec expires_at(table :: atom(), key :: term(), scale :: pos_integer()) :: non_neg_integer()
+  def expires_at(table, key, scale) do
+    window = div(ETS.now(), scale)
+    full_key = {key, window}
+
+    case :ets.lookup(table, full_key) do
+      [{_full_key, _count, expires_at}] -> expires_at
+      [] -> 0
+    end
+  end
+
   @doc false
   @spec select_expired(config :: ETS.config()) :: list()
   def select_expired(config) do

test/hammer/atomic/fix_window_test.exs

@@ -128,6 +128,25 @@ defmodule Hammer.Atomic.FixWindowTest do
     end
   end
 
+  describe "expires_at" do
+    test "returns 0 for unknown key", %{table: table} do
+      assert FixWindow.expires_at(table, "unknown", :timer.seconds(10)) == 0
+    end
+
+    test "returns the expiration timestamp after a hit", %{table: table} do
+      key = "key"
+      scale = :timer.seconds(10)
+
+      assert {:allow, 1} = FixWindow.hit(table, key, scale, 10, 1)
+
+      expires_at = FixWindow.expires_at(table, key, scale)
+      now = System.system_time(:millisecond)
+
+      assert expires_at > now
+      assert expires_at <= now + scale
+    end
+  end
+
   describe "get/set" do
     test "get returns the count set for the given key and scale", %{table: table} do
       key = "key"

test/hammer/ets/fix_window_test.exs

@@ -100,6 +100,25 @@ defmodule Hammer.ETS.FixWindowTest do
     end
   end
 
+  describe "expires_at" do
+    test "returns 0 for unknown key", %{table: table} do
+      assert FixWindow.expires_at(table, "unknown", :timer.seconds(10)) == 0
+    end
+
+    test "returns the expiration timestamp after a hit", %{table: table} do
+      key = "key"
+      scale = :timer.seconds(10)
+
+      assert {:allow, 1} = FixWindow.hit(table, key, scale, 10, 1)
+
+      expires_at = FixWindow.expires_at(table, key, scale)
+      now = System.system_time(:millisecond)
+
+      assert expires_at > now
+      assert expires_at <= now + scale
+    end
+  end
+
   describe "get/set" do
     test "get returns the count set for the given key and scale", %{table: table} do
       key = "key"