compile time everything

config/config.exs

@@ -1,12 +1,10 @@
-# This file is responsible for configuring your application
-# and its dependencies with the aid of the Mix.Config module.
 import Config
 
-config :hammer,
-  backend:
-    {Hammer.Backend.ETS,
-     [
-       ets_table_name: :hammer_backend_ets_buckets,
-       expiry_ms: 60_000 * 60 * 2,
-       cleanup_interval_ms: 60_000 * 2
-     ]}
+if config_env() == :test do
+  config :hammer,
+    backend:
+      {Hammer.Backend.ETS,
+       ets_table_name: :hammer_backend_ets_test_buckets,
+       expiry_ms: :timer.hours(2),
+       cleanup_interval_ms: :timer.minutes(2)}
+end

coveralls.json

@@ -1,5 +1,5 @@
 {
   "skip_files": [
-    "lib/hammer/application.ex"
+    "lib/hammer/backend.ex"
   ]
 }

lib/hammer.ex

@@ -8,111 +8,100 @@
 
   alias Hammer.Utils
 
-  @spec check_rate(id :: String.t(), scale_ms :: integer, limit :: integer) ::
-          {:allow, count :: integer}
-          | {:deny, limit :: integer}
-          | {:error, reason :: any}
+  backends =
+    case Application.compile_env!(:hammer, :backend) do
+      {_mod, _opts} = backend -> [single: backend]
+      [_ | _] = backends -> backends
+    end
+
+  {default_backend, _config} =
+    List.first(backends) ||
+      raise ArgumentError,
+            "expected at least one backend provided in :backends, got: #{inspect(backends)}"
+
+  @type id :: String.t()
+  @type scale_ms :: pos_integer
+  @type limit :: pos_integer
+  @type count :: pos_integer
+  @type increment :: non_neg_integer
+  @type backend :: atom
+
+  @type bucket_info :: {
+          count,
+          count_remaining :: non_neg_integer,
+          ms_to_next_bucket :: pos_integer,
+          created_at :: pos_integer | nil,
+          updated_at :: pos_integer | nil
+        }
+
   @doc """
   Check if the action you wish to perform is within the bounds of the rate-limit.
 
   Args:
   - `id`: String name of the bucket. Usually the bucket name is comprised of
   some fixed prefix, with some dynamic string appended, such as an IP address or
-  user id.
-  - `scale_ms`: Integer indicating size of bucket in milliseconds
+  user id
+  - `scale`: Integer indicating size of bucket in milliseconds
   - `limit`: Integer maximum count of actions within the bucket
 
   Returns either `{:allow,  count}`, `{:deny,   limit}` or `{:error,  reason}`
 
   Example:
 
       user_id = 42076
-      case check_rate("file_upload:\#{user_id}", 60_000, 5) do
+      case check_rate("file_upload:\#{user_id}", _scale = :timer.seconds(60), _limit = 5) do
         {:allow, _count} ->
           # do the file upload
         {:deny, _limit} ->
           # render an error page or something
       end
+
   """
-  def check_rate(id, scale_ms, limit) do
-    check_rate(:single, id, scale_ms, limit)
+  @spec check_rate(id, scale_ms, limit) :: {:allow, count} | {:deny, limit} | {:error, term}
+  def check_rate(id, scale, limit) do
+    check_rate_inc(id, scale, limit, 1)
   end
 
-  @spec check_rate(backend :: atom, id :: String.t(), scale_ms :: integer, limit :: integer) ::
-          {:allow, count :: integer}
-          | {:deny, limit :: integer}
-          | {:error, reason :: any}
-  @doc """
-  Same as `check_rate/3`, but allows specifying a backend.
-  """
+  @doc "Same as `check_rate/3`, but allows specifying a backend"
+  @spec check_rate(backend, id, scale_ms, limit) ::
+          {:allow, count} | {:deny, limit} | {:error, reason}
   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
-      {:ok, count} ->
-        if count > limit do
-          {:deny, limit}
-        else
-          {:allow, count}
-        end
-
-      {:error, reason} ->
-        {:error, reason}
-    end
+    check_rate_inc(backend, id, scale_ms, limit, 1)
   end
 
-  @spec check_rate_inc(
-          id :: String.t(),
-          scale_ms :: integer,
-          limit :: integer,
-          increment :: integer
-        ) ::
-          {:allow, count :: integer}
-          | {:deny, limit :: integer}
-          | {:error, reason :: any}
   @doc """
   Same as check_rate/3, but allows the increment number to be specified.
   This is useful for limiting apis which have some idea of 'cost', where the cost
   of each hit can be specified.
   """
+  @spec check_rate_inc(id, scale_ms, limit, increment) ::
+          {:allow, count} | {:deny, limit} | {:error, term}
   def check_rate_inc(id, scale_ms, limit, increment) do
-    check_rate_inc(:single, id, scale_ms, limit, increment)
+    check_rate_inc(unquote(default_backend), id, scale_ms, limit, increment)
   end
 
-  @spec check_rate_inc(
-          backend :: atom,
-          id :: String.t(),
-          scale_ms :: integer,
-          limit :: integer,
-          increment :: integer
-        ) ::
-          {:allow, count :: integer}
-          | {:deny, limit :: integer}
-          | {:error, reason :: any}
-  @doc """
-  Same as check_rate_inc/4, but allows specifying a backend.
-  """
-  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
-      {:ok, count} ->
-        if count > limit do
-          {:deny, limit}
-        else
-          {:allow, count}
-        end
+  @doc "Same as check_rate_inc/4, but allows specifying a backend"
+  @spec check_rate_inc(backend, id, scale_ms, limit, increment) ::
+          {:allow, count} | {:deny, limit} | {:error, term}
+  def check_rate_inc(backend, id, scale_ms, limit, increment)
+
+  for {backend, {mod, opts}} <- backends do
+    def check_rate_inc(unquote(backend), id, scale_ms, limit, increment) do
+      {stamp, key} = stamp_key(id, scale_ms)
 
-      {:error, reason} ->
-        {:error, reason}
+      try do
+        count = unquote(mod).count_hit(key, stamp, increment, unquote(opts))
+        if count > limit, do: {:deny, limit}, else: {:allow, count}
+      rescue
+        e -> {:error, e}
+      end
     end
   end
 
-  @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}}
-          | {:error, reason :: any}
+  def check_rate_inc(backend, _id, _scale_ms, _limit, _increment) do
+    unknown_backend(backend)
+  end
+
   @doc """
   Inspect bucket to get count, count_remaining, ms_to_next_bucket, created_at,
   updated_at. This function is free of side-effects and should be called with
@@ -137,38 +126,46 @@
       {:ok, {1, 2499, 29381612, 1450281014468, 1450281014468}}
 
   """
+  @spec inspect_bucket(id, scale_ms, limit) :: {:ok, bucket_info} | {:error, term}
   def inspect_bucket(id, scale_ms, limit) do
-    inspect_bucket(:single, id, scale_ms, limit)
+    inspect_bucket(unquote(default_backend), id, scale_ms, limit)
   end
 
-  @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}}
-          | {: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
-
-    case call_backend(backend, :get_bucket, [key]) do
-      {:ok, nil} ->
-        {:ok, {0, limit, ms_to_next_bucket, nil, nil}}
+  @spec inspect_bucket(backend, id, scale_ms, limit) :: {:ok, bucket_info} | {:error, term}
+  def inspect_bucket(backend, id, scale_ms, limit)
+
+  for {backend, {mod, opts}} <- backends do
+    def inspect_bucket(unquote(backend), id, scale_ms, limit) do
+      {stamp, key} = stamp_key(id, scale_ms)
+      ms_to_next_bucket = elem(key, 0) * scale_ms + scale_ms - stamp
+
+      result =
+        try do
+          {:ok, unquote(mod).get_bucket(key, unquote(opts))}
+        rescue
+          e -> {:error, e}
+        end
 
-      {:ok, {_, count, created_at, updated_at}} ->
-        count_remaining = if limit > count, do: limit - count, else: 0
-        {:ok, {count, count_remaining, ms_to_next_bucket, created_at, updated_at}}
+      with {:ok, bucket} <- result do
+        case bucket do
+          {_, count, created_at, updated_at} ->
+            count_remaining = if limit > count, do: limit - count, else: 0
+            {:ok, {count, count_remaining, ms_to_next_bucket, created_at, updated_at}}
 
-      {:error, reason} ->
-        {:error, reason}
+          nil ->
+            {:ok, {0, limit, ms_to_next_bucket, nil, nil}}
+        end
+      end
     end
   end
 
-  @spec delete_buckets(id :: String.t()) ::
-          {:ok, count :: integer}
-          | {:error, reason :: any}
+  def inspect_bucket(backend, _id, _scale_ms, _limit) do
+    unknown_backend(backend)
+  end
+
   @doc """
   Delete all buckets belonging to the provided id, including the current one.
   Effectively resets the rate-limit for the id.
@@ -186,79 +183,42 @@
       {:ok, _count} = delete_buckets("file_uploads:\#{user_id}")
 
   """
+  @spec delete_buckets(id) :: {:ok, count_deleted :: non_neg_integer} | {:error, term}
   def delete_buckets(id) do
-    delete_buckets(:single, id)
+    delete_buckets(unquote(default_backend), id)
   end
 
-  @spec delete_buckets(backend :: atom, id :: String.t()) ::
-          {:ok, count :: integer}
-          | {:error, reason :: any}
   @doc """
   Same as delete_buckets/1, but allows specifying a backend
   """
-  def delete_buckets(backend, id) do
-    call_backend(backend, :delete_buckets, [id])
+  @spec delete_buckets(backend, id) :: {:ok, count_deleted :: non_neg_integer} | {:error, term}
+  def delete_buckets(backend, id)
+
+  for {backend, {mod, opts}} <- backends do
+    def delete_buckets(unquote(backend), id) do
+      try do
+        {:ok, unquote(mod).delete_buckets(id, unquote(opts))}
+      rescue
+        e -> {:error, e}
+      end
+    end
   end
 
-  @spec make_rate_checker(id_prefix :: String.t(), scale_ms :: integer, limit :: integer) ::
-          (id :: String.t() ->
-             {:allow, count :: integer}
-             | {:deny, limit :: integer}
-             | {:error, reason :: any})
-  @doc """
-  Make a rate-checker function, with the given `id` prefix, scale_ms and limit.
-
-  Arguments:
-
-  - `id_prefix`: String prefix to the `id`
-  - `scale_ms`: Integer indicating size of bucket in milliseconds
-  - `limit`: Integer maximum count of actions within the bucket
-
-  Returns a function which accepts an `id` suffix, which will be combined with
-  the `id_prefix`. Calling this returned function is equivalent to:
-  `Hammer.check_rate("\#{id_prefix}\#{id}", scale_ms, limit)`
-
-  Example:
-
-      chat_rate_limiter = make_rate_checker("send_chat_message:", 60_000, 20)
-      user_id = 203517
-      case chat_rate_limiter.(user_id) do
-        {:allow, _count} ->
-          # allow chat message
-        {:deny, _limit} ->
-          # deny
-      end
-  """
-  def make_rate_checker(id_prefix, scale_ms, limit) do
-    make_rate_checker(:single, id_prefix, scale_ms, limit)
+  def delete_buckets(backend, _id) do
+    unknown_backend(backend)
   end
 
-  @spec make_rate_checker(
-          backend :: atom,
-          id_prefix :: String.t(),
-          scale_ms :: integer,
-          limit :: integer
-        ) ::
-          (id :: String.t() ->
-             {:allow, count :: integer}
-             | {:deny, limit :: integer}
-             | {:error, reason :: any})
-  @doc """
-  """
-  def make_rate_checker(backend, id_prefix, scale_ms, limit) do
-    fn id ->
-      check_rate(backend, "#{id_prefix}#{id}", scale_ms, limit)
-    end
+  @compile inline: [unknown_backend: 1]
+  defp unknown_backend(backend) do
+    raise ArgumentError, "unknown backend #{inspect(backend)}"
   end
 
-  defp call_backend(which, function, args) do
-    pool = Utils.pool_name(which)
-    backend = Utils.get_backend_module(which)
+  @compile inline: [timestamp: 0]
+  defp timestamp, do: System.system_time(:millisecond)
 
-    :poolboy.transaction(
-      pool,
-      fn pid -> apply(backend, function, [pid | args]) end,
-      60_000
-    )
+  @compile inline: [full_key: 2]
+  defp full_key(key, scale) do
+    bucket = div(now(), scale)
+    {key, bucket}
   end
 end