diff --git a/lib/hammer.ex b/lib/hammer.ex index 6015872..228e849 100644 --- a/lib/hammer.ex +++ b/lib/hammer.ex @@ -4,14 +4,77 @@ defmodule Hammer do This is the main API for the Hammer rate-limiter. This module assumes a backend pool has been started, most likely by the Hammer application. + + defmodule MyApp.RateLimit do + use Hammer, backend: Hammer.Backend.ETS + end + + MyApp.RateLimit.start_link(expiry_ms: :timer.minutes(10), cleanup_interval_ms: :timer.minutes(1)) + {:allow, 1] = MyApp.RateLimit.hit("some-key", _scale = :timer.seconds(60), _limit = 10) + {:allow, 1} = MyApp.RateLimit.get("some-key", _scale = :timer.seconds(60), _limit = 10) + MyApp.RateLimit.wait("some-key", _scale = :timer.seconds(60), _limit = 10) + + backend = Hammer.backend(Hammer.Backend.ETS, opts) + Hammer.check_rate(backend, ) + """ - alias Hammer.Utils + @doc false + defmacro __using__(opts) do + quote bind_quoted: [opts: opts] do + backend = Keyword.fetch!(opts, :backend) + + unless is_atom(backend) do + raise ArgumentError, + "expected :backend to be a module name like Hammer.Backend.ETS, got: #{inspect(backend)}" + end + + @behaviour Hammer + @backend backend + + def __backend__ do + @backend + end + + def start_link(opts) do + @backend.start_link(opts) + end + + def child_spec(opts) do + @backend.child_spec(opts) + end + + def check_rate(id, scale_ms, limit) do + @backend.check_rate_inc(id, scale_ms, limit, 1) + end + + def check_rate_inc(id, scale_ms, limit, increment) do + @backend.check_rate_inc(id, scale_ms, limit, increment) + end + + def inspect_bucket(id, scale_ms, limit) do + {stamp, key} = Hammer.Utils.stamp_key(id, scale_ms) + ms_to_next_bucket = elem(key, 0) * scale_ms + scale_ms - stamp + + case @backend.get_bucket(key) do + {:ok, nil} -> + {:ok, {0, limit, ms_to_next_bucket, nil, nil}} + + {: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}} + + {:error, reason} -> + {:error, reason} + end + end + + def delete_buckets(id) do + @backend.delete_buckets(id) + end + end + end - @spec check_rate(id :: String.t(), scale_ms :: integer, limit :: integer) :: - {:allow, count :: integer} - | {:deny, limit :: integer} - | {:error, reason :: any} @doc """ Check if the action you wish to perform is within the bounds of the rate-limit. @@ -27,92 +90,34 @@ defmodule Hammer do Example: user_id = 42076 - case check_rate("file_upload:\#{user_id}", 60_000, 5) do + case MyHammer.check_rate("file_upload:\#{user_id}", 60_000, 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) - 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. """ - 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 + @callback check_rate(id :: String.t(), scale_ms :: integer, limit :: integer) :: + {:allow, count :: integer} + | {:deny, limit :: integer} + | {:error, reason :: any} - {:error, reason} -> - {:error, reason} - end - 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. """ - def check_rate_inc(id, scale_ms, limit, increment) do - check_rate_inc(:single, id, scale_ms, limit, increment) - end + @callback check_rate_inc( + id :: String.t(), + scale_ms :: integer, + limit :: integer, + increment :: integer + ) :: + {:allow, count :: integer} + | {:deny, limit :: integer} + | {:error, reason :: any} - @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 - - {:error, reason} -> - {:error, reason} - 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} @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 +142,12 @@ defmodule Hammer do {:ok, {1, 2499, 29381612, 1450281014468, 1450281014468}} """ - def inspect_bucket(id, scale_ms, limit) do - inspect_bucket(:single, 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}} - - {: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}} - - {:error, reason} -> - {:error, reason} - end - end + @callback 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} - @spec delete_buckets(id :: String.t()) :: - {:ok, count :: integer} - | {:error, reason :: any} @doc """ Delete all buckets belonging to the provided id, including the current one. Effectively resets the rate-limit for the id. @@ -183,82 +162,8 @@ defmodule Hammer do Example: user_id = 2406 - {:ok, _count} = delete_buckets("file_uploads:\#{user_id}") - - """ - def delete_buckets(id) do - delete_buckets(:single, 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]) - 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. + {:ok, _count} = MyHammer.delete_buckets("file_uploads:\#{user_id}") - 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) - 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 - end - - defp call_backend(which, function, args) do - pool = Utils.pool_name(which) - backend = Utils.get_backend_module(which) - - :poolboy.transaction( - pool, - fn pid -> apply(backend, function, [pid | args]) end, - 60_000 - ) - end + @callback delete_buckets(id :: String.t()) :: {:ok, count :: integer} | {:error, reason :: any} end diff --git a/lib/hammer/application.ex b/lib/hammer/application.ex index d406d96..7b6be5d 100644 --- a/lib/hammer/application.ex +++ b/lib/hammer/application.ex @@ -53,16 +53,15 @@ defmodule Hammer.Application do """ use Application - require Logger def start(_type, _args) do - config = + children = Application.get_env( :hammer, :backend, {Hammer.Backend.ETS, []} ) - Hammer.Supervisor.start_link(config, name: Hammer.Supervisor) + Hammer.Supervisor.start_link(children, name: Hammer.Supervisor) end end diff --git a/lib/hammer/backend.ex b/lib/hammer/backend.ex index 81cf0f5..76f3cae 100644 --- a/lib/hammer/backend.ex +++ b/lib/hammer/backend.ex @@ -7,35 +7,12 @@ defmodule Hammer.Backend do @type bucket_info :: {key :: bucket_key, count :: integer, created :: integer, updated :: integer} - @callback count_hit( - pid :: pid(), - key :: bucket_key, - now :: integer - ) :: - {:ok, count :: integer} - | {:error, reason :: any} + @callback count_hit(key :: bucket_key, now :: integer, increment :: integer, opts :: keyword) :: + {:ok, count :: integer} | {:error, reason :: any} - @callback count_hit( - pid :: pid(), - key :: bucket_key, - now :: integer, - increment :: integer - ) :: - {:ok, count :: integer} - | {:error, reason :: any} + @callback get_bucket(key :: bucket_key, opts :: keyword) :: + {:ok, info :: bucket_info | nil} | {:error, reason :: any} - @callback get_bucket( - pid :: pid(), - key :: bucket_key - ) :: - {:ok, info :: bucket_info} - | {:ok, nil} - | {:error, reason :: any} - - @callback delete_buckets( - pid :: pid(), - id :: String.t() - ) :: - {:ok, count_deleted :: integer} - | {:error, reason :: any} + @callback delete_buckets(id :: String.t(), opts :: keyword) :: + {:ok, count_deleted :: integer} | {:error, reason :: any} end diff --git a/lib/hammer/backend/ets.ex b/lib/hammer/backend/ets.ex index 3f3e044..db1f016 100644 --- a/lib/hammer/backend/ets.ex +++ b/lib/hammer/backend/ets.ex @@ -62,32 +62,8 @@ defmodule Hammer.Backend.ETS do GenServer.call(__MODULE__, :stop) end - @doc """ - Record a hit in the bucket identified by `key` - """ - @spec count_hit( - pid :: pid(), - key :: bucket_key, - now :: integer - ) :: - {:ok, count :: integer} - | {:error, reason :: any} - def count_hit(pid, key, now) do - count_hit(pid, key, now, 1) - end - - @doc """ - Record a hit in the bucket identified by `key`, with a custom increment - """ - @spec count_hit( - pid :: pid(), - key :: bucket_key, - now :: integer, - increment :: integer - ) :: - {:ok, count :: integer} - | {:error, reason :: any} - def count_hit(_pid, key, now, increment) do + @impl Hammer.Backend + def count_hit(key, now, increment, _opts) do if :ets.member(@ets_table_name, key) do [count, _] = :ets.update_counter(@ets_table_name, key, [ @@ -108,17 +84,8 @@ defmodule Hammer.Backend.ETS do {:error, e} end - @doc """ - Retrieve information about the bucket identified by `key` - """ - @spec get_bucket( - pid :: pid(), - key :: bucket_key - ) :: - {:ok, info :: bucket_info} - | {:ok, nil} - | {:error, reason :: any} - def get_bucket(_pid, key) do + @impl Hammer.Backend + def get_bucket(key, _opts) do result = case :ets.lookup(@ets_table_name, key) do [] -> @@ -134,16 +101,8 @@ defmodule Hammer.Backend.ETS do {:error, e} end - @doc """ - Delete all buckets associated with `id`. - """ - @spec delete_buckets( - pid :: pid(), - id :: String.t() - ) :: - {:ok, count_deleted :: integer} - | {:error, reason :: any} - def delete_buckets(_pid, id) do + @impl Hammer.Backend + def delete_buckets(id, _opts) do # Compiled from: # fun do {{bucket_number, bid},_,_,_} when bid == ^id -> true end count_deleted = @@ -159,6 +118,7 @@ defmodule Hammer.Backend.ETS do ## GenServer Callbacks + @impl GenServer def init(args) do cleanup_interval_ms = Keyword.get(args, :cleanup_interval_ms) expiry_ms = Keyword.get(args, :expiry_ms) @@ -193,10 +153,12 @@ defmodule Hammer.Backend.ETS do {:ok, state} end + @impl GenServer def handle_call(:stop, _from, state) do {:stop, :normal, :ok, state} end + @impl GenServer def handle_info(:prune, state) do %{expiry_ms: expiry_ms} = state now = Utils.timestamp() diff --git a/lib/hammer/supervisor.ex b/lib/hammer/supervisor.ex index 3b37fc9..df9f6a0 100644 --- a/lib/hammer/supervisor.ex +++ b/lib/hammer/supervisor.ex @@ -15,37 +15,13 @@ defmodule Hammer.Supervisor do # Single backend def init(config) when is_tuple(config) do - children = [ - to_pool_spec(:hammer_backend_single_pool, config) - ] - + children = [config] Supervisor.init(children, strategy: :one_for_one) end # Multiple backends def init(config) when is_list(config) do - children = - Enum.map(config, fn {k, c} -> - "hammer_backend_#{k}_pool" - |> String.to_atom() - |> to_pool_spec(c) - end) - + children = Enum.map(config, fn {_k, c} -> c end) Supervisor.init(children, strategy: :one_for_one) end - - # Private helpers - defp to_pool_spec(name, {mod, args}) do - pool_size = args[:pool_size] || 4 - pool_max_overflow = args[:pool_max_overflow] || 0 - - opts = [ - name: {:local, name}, - worker_module: mod, - size: pool_size, - max_overflow: pool_max_overflow - ] - - :poolboy.child_spec(name, opts, args) - end end diff --git a/lib/hammer/utils.ex b/lib/hammer/utils.ex index f226cb0..112cc76 100644 --- a/lib/hammer/utils.ex +++ b/lib/hammer/utils.ex @@ -1,14 +1,6 @@ defmodule Hammer.Utils do @moduledoc false - def pool_name do - pool_name(:single) - end - - def pool_name(name) do - String.to_atom("hammer_backend_#{name}_pool") - end - # Returns Erlang Time as milliseconds since 00:00 GMT, January 1, 1970 def timestamp do DateTime.to_unix(DateTime.utc_now(), :millisecond) diff --git a/mix.exs b/mix.exs index b37b064..209317c 100644 --- a/mix.exs +++ b/mix.exs @@ -28,8 +28,7 @@ defmodule Hammer.Mixfile do [ {:credo, "~> 1.7", only: [:dev, :test]}, {:ex_doc, "~> 0.30", only: :dev}, - {:dialyxir, "~> 1.3", only: [:dev, :test], runtime: false}, - {:poolboy, "~> 1.5"} + {:dialyxir, "~> 1.3", only: [:dev, :test], runtime: false} ] end diff --git a/mix.lock b/mix.lock index 18cce91..79aa696 100644 --- a/mix.lock +++ b/mix.lock @@ -11,5 +11,4 @@ "makeup_elixir": {:hex, :makeup_elixir, "0.16.2", "627e84b8e8bf22e60a2579dad15067c755531fea049ae26ef1020cad58fe9578", [:mix], [{:makeup, "~> 1.0", [hex: :makeup, repo: "hexpm", optional: false]}, {:nimble_parsec, "~> 1.2.3 or ~> 1.3", [hex: :nimble_parsec, repo: "hexpm", optional: false]}], "hexpm", "41193978704763f6bbe6cc2758b84909e62984c7752b3784bd3c218bb341706b"}, "makeup_erlang": {:hex, :makeup_erlang, "0.1.5", "e0ff5a7c708dda34311f7522a8758e23bfcd7d8d8068dc312b5eb41c6fd76eba", [:mix], [{:makeup, "~> 1.0", [hex: :makeup, repo: "hexpm", optional: false]}], "hexpm", "94d2e986428585a21516d7d7149781480013c56e30c6a233534bedf38867a59a"}, "nimble_parsec": {:hex, :nimble_parsec, "1.4.0", "51f9b613ea62cfa97b25ccc2c1b4216e81df970acd8e16e8d1bdc58fef21370d", [:mix], [], "hexpm", "9c565862810fb383e9838c1dd2d7d2c437b3d13b267414ba6af33e50d2d1cf28"}, - "poolboy": {:hex, :poolboy, "1.5.2", "392b007a1693a64540cead79830443abf5762f5d30cf50bc95cb2c1aaafa006b", [:rebar3], [], "hexpm", "dad79704ce5440f3d5a3681c8590b9dc25d1a561e8f5a9c995281012860901e3"}, }