diff --git a/.formatter.exs b/.formatter.exs index d2cda26..7f59b1f 100644 --- a/.formatter.exs +++ b/.formatter.exs @@ -1,4 +1,4 @@ # Used by "mix format" [ - inputs: ["{mix,.formatter}.exs", "{config,lib,test}/**/*.{ex,exs}"] + inputs: ["{mix,.formatter}.exs", "{lib,test}/**/*.{ex,exs}"] ] diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index f4e893f..ba5dade 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -18,11 +18,9 @@ jobs: fail-fast: false matrix: os: [ubuntu-22.04, ubuntu-20.04] - elixir_version: [1.12.3, 1.13.3, 1.14.1] - otp_version: [24, 25] - exclude: - - otp_version: 25 - elixir_version: 1.12.3 + elixir_version: [1.13, 1.14, 1.15] + otp_version: [24, 25, 26] + steps: - uses: actions/checkout@v4 diff --git a/.tool-versions b/.tool-versions index b6551eb..328b6a4 100644 --- a/.tool-versions +++ b/.tool-versions @@ -1 +1 @@ -elixir 1.13.1 +elixir 1.15.7-otp-26 diff --git a/CHANGELOG.md b/CHANGELOG.md index c3b7535..bf4fd37 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,12 @@ # Changelog +## 7.0.0 - unreleased + +- remove `poolboy` dependency +- improve ETS backend efficiency ~15x +- fix ETS backend overload +- remove global supervisor + ## 6.1.0 - 2022-06-13 ### Changed diff --git a/README.md b/README.md index 6061224..0136ce5 100644 --- a/README.md +++ b/README.md @@ -21,7 +21,7 @@ by adding `:hammer` to your list of dependencies in `mix.exs`: ```elixir def deps do [ - {:hammer, "~> 6.1"} + {:hammer, "~> 7.0"} ] end ``` @@ -38,7 +38,6 @@ Example: ```elixir defmodule MyApp.VideoUpload do - def upload(video_data, user_id) do case Hammer.check_rate("upload_video:#{user_id}", 60_000, 5) do {:allow, _count} -> @@ -47,7 +46,6 @@ defmodule MyApp.VideoUpload do # deny the request end end - end ``` @@ -58,14 +56,6 @@ The `Hammer` module provides the following functions: - `inspect_bucket(id, scale_ms, limit)` - `delete_buckets(id)` -Backends are configured via `Mix.Config`: - -```elixir -config :hammer, - backend: {Hammer.Backend.ETS, [expiry_ms: 60_000 * 60 * 4, - cleanup_interval_ms: 60_000 * 10]} -``` - See the [Tutorial](https://hexdocs.pm/hammer/tutorial.html) for more. See the [Hammer Testbed](https://github.com/ExHammer/hammer-testbed) app for an example of diff --git a/config/config.exs b/config/config.exs deleted file mode 100644 index 86d7cb7..0000000 --- a/config/config.exs +++ /dev/null @@ -1,12 +0,0 @@ -# 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 - ]} diff --git a/lib/hammer.ex b/lib/hammer.ex index 6015872..f45dc7f 100644 --- a/lib/hammer.ex +++ b/lib/hammer.ex @@ -2,16 +2,20 @@ defmodule Hammer do @moduledoc """ Documentation for Hammer module. - 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. + This is the main API for the Hammer rate-limiter. This module assumes the + backend has been started. + + Example: + + def start(_, _) do + children = [{Hammer.Backend.ETS, cleanup_interval_ms: :timer.seconds(5)}] + Supervisor.init(children, []) + end + """ - alias Hammer.Utils + @backend Application.compile_env(:hammer, :backend, Hammer.Backend.ETS) - @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. @@ -22,7 +26,7 @@ defmodule Hammer do - `scale_ms`: 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}` + Returns either `{:allow, count}`, `{:deny, limit}` Example: @@ -34,85 +38,35 @@ defmodule Hammer do # render an error page or something end """ + @spec check_rate(id :: String.t(), scale_ms :: integer, limit :: integer) :: + {:allow, count :: integer} | {:deny, limit :: integer} | {:error, reason :: any} def check_rate(id, scale_ms, limit) do - check_rate(:single, id, scale_ms, limit) + check_rate_inc(id, scale_ms, 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. - """ - 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 - 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 - @spec check_rate_inc( - backend :: atom, id :: String.t(), scale_ms :: integer, limit :: integer, - increment :: integer + increment :: non_neg_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 + {:allow, count :: integer} | {:deny, limit :: integer} | {:error, reason :: any} + def check_rate_inc(id, scale_ms, limit, increment) do + now = System.system_time(:millisecond) + now_bucket = div(now, scale_ms) + key = {id, now_bucket} + expires_at = (now_bucket + 1) * scale_ms - {:error, reason} -> - {:error, reason} + with {:ok, count} <- @backend.count_hit(key, increment, expires_at) do + if count <= limit, do: {:allow, count}, else: {:deny, limit} 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 +91,29 @@ 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) :: + @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 """ - 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} + def inspect_bucket(id, scale_ms, limit) do + now = System.system_time(:millisecond) + now_bucket = div(now, scale_ms) + key = {id, now_bucket} + ms_to_next_bucket = now_bucket * scale_ms + scale_ms - now + + with {:ok, count} <- @backend.get_bucket(key) do + {:ok, + { + count, + max(limit - count, 0), + ms_to_next_bucket, + _created_at = nil, + _updated_at = nil + }} end end - @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. @@ -186,79 +131,8 @@ defmodule Hammer do {:ok, _count} = delete_buckets("file_uploads:\#{user_id}") """ + @spec delete_buckets(id :: String.t()) :: {:ok, count :: integer} | {:error, reason :: any} 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. - - 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 - ) + @backend.delete_buckets(id) end end diff --git a/lib/hammer/application.ex b/lib/hammer/application.ex deleted file mode 100644 index d406d96..0000000 --- a/lib/hammer/application.ex +++ /dev/null @@ -1,68 +0,0 @@ -defmodule Hammer.Application do - @moduledoc """ - Hammer application, responsible for starting the backend worker pools. - - Configured with the `:hammer` environment key: - - - `:backend`, Either a tuple of `{module, config}`, or a keyword-list - of separate, named backends. Examples: - `{Hammer.Backend.ETS, []}`, `[ets: {Hammer.Backend.ETS, []}, ...]` - - `:suppress_logs`, if set to `true`, stops all log messages from Hammer - - ### General Backend Options - - Different backends take different options, but all will accept the following - options, and with the same effect: - - - `:expiry_ms` (int): expiry time in milliseconds, after which a bucket will - be deleted. The exact mechanism for cleanup will vary by backend. This configuration - option is mandatory - - `:pool_size` (int): size of the backend worker pool (default=2) - - `:pool_max_overflow` int(): number of extra workers the pool is permitted - to spawn when under pressure. The worker pool (managed by the poolboy library) - will automatically create and destroy workers up to the max-overflow limit - (default=0) - - Example of a single backend: - - config :hammer, - backend: {Hammer.Backend.ETS, [expiry_ms: 60_000 * 60 * 2]} - - Example of config for multiple-backends: - - config :hammer, - backend: [ - ets: { - Hammer.Backend.ETS, - [ - ets_table_name: :hammer_backend_ets_buckets, - expiry_ms: 60_000 * 60 * 2, - cleanup_interval_ms: 60_000 * 2, - ] - }, - redis: { - Hammer.Backend.Redis, - [ - expiry_ms: 60_000 * 60 * 2, - redix_config: [host: "localhost", port: 6379], - pool_size: 4, - ] - } - ] - - """ - - use Application - require Logger - - def start(_type, _args) do - config = - Application.get_env( - :hammer, - :backend, - {Hammer.Backend.ETS, []} - ) - - Hammer.Supervisor.start_link(config, name: Hammer.Supervisor) - end -end diff --git a/lib/hammer/backend.ex b/lib/hammer/backend.ex index 81cf0f5..8f73070 100644 --- a/lib/hammer/backend.ex +++ b/lib/hammer/backend.ex @@ -3,39 +3,14 @@ defmodule Hammer.Backend do The backend Behaviour module. """ - @type bucket_key :: {bucket :: integer, id :: String.t()} - @type bucket_info :: - {key :: bucket_key, count :: integer, created :: integer, updated :: integer} + @type bucket_key :: {id :: String.t(), bucket :: integer} - @callback count_hit( - pid :: pid(), - key :: bucket_key, - now :: integer - ) :: - {:ok, count :: integer} - | {:error, reason :: any} + @callback count_hit(bucket_key, increment :: integer, expires_at :: integer) :: + {: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(bucket_key) :: + {:ok, count :: integer} | {: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()) :: + {:ok, count_deleted :: integer} | {:error, reason :: any} end diff --git a/lib/hammer/backend/ets.ex b/lib/hammer/backend/ets.ex index 4fed847..b21f5e8 100644 --- a/lib/hammer/backend/ets.ex +++ b/lib/hammer/backend/ets.ex @@ -1,204 +1,72 @@ defmodule Hammer.Backend.ETS do - @moduledoc """ - An ETS backend for Hammer. - - The public API of this module is used by Hammer to store information about - rate-limit 'buckets'. A bucket is identified by a `key`, which is a tuple - `{bucket_number, id}`. The essential schema of a bucket is: - `{key, count, created_at, updated_at}`, although backends are free to - store and retrieve this data in whichever way they wish. - - Use `start` or `start_link` to start the server: - - {:ok, pid} = Hammer.Backend.ETS.start_link(args) - - `args` is a keyword list: - - `expiry_ms`: (integer) time in ms before a bucket is auto-deleted, - should be larger than the expected largest size/duration of a bucket - - `cleanup_interval_ms`: (integer) time between cleanup runs, - - Example: - - Hammer.Backend.ETS.start_link( - expiry_ms: 1000 * 60 * 60, - cleanup_interval_ms: 1000 * 60 * 10 - ) - """ - - @behaviour Hammer.Backend + @moduledoc "TODO" use GenServer - alias Hammer.Utils - - @type bucket_key :: {bucket :: integer, id :: String.t()} - @type bucket_info :: - {key :: bucket_key, count :: integer, created :: integer, updated :: integer} - - @ets_table_name :hammer_ets_buckets - - ## Public API - - def start do - start([]) - end - - def start(args) do - GenServer.start(__MODULE__, args) - end - - def start_link do - start_link([]) - end + @behaviour Hammer.Backend + @table :hammer_ets_buckets - @doc """ - """ - def start_link(args) do - GenServer.start_link(__MODULE__, args) - end + @doc "TODO" + def start_link(opts) do + {gen_opts, opts} = + Keyword.split(opts, [:debug, :name, :timeout, :spawn_opt, :hibernate_after]) - def stop do - GenServer.call(__MODULE__, :stop) + GenServer.start_link(__MODULE__, opts, gen_opts) 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) + @impl Hammer.Backend + def count_hit(key, increment, expires_at) do + {:ok, :ets.update_counter(@table, key, increment, {key, 0, expires_at})} 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 - if :ets.member(@ets_table_name, key) do - [count, _] = - :ets.update_counter(@ets_table_name, key, [ - # Increment count field - {2, increment}, - # Set updated_at to now - {4, 1, 0, now} - ]) + @impl Hammer.Backend + def get_bucket(key) do + count = + case :ets.lookup(@table, key) do + [{^key, count, _expires_at}] -> count + [] -> 0 + end - {:ok, count} - else - # Insert {key, count, created_at, updated_at} - true = :ets.insert(@ets_table_name, {key, increment, now, now}) - {:ok, increment} - end - rescue - e -> - {:error, e} + {:ok, count} 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 - result = - case :ets.lookup(@ets_table_name, key) do - [] -> - {:ok, nil} - - [bucket] -> - {:ok, bucket} - end - - result - rescue - e -> - {:error, e} + @impl Hammer.Backend + def delete_buckets(id) do + ms = [{{{:"$1", :_}, :_, :_}, [], [{:==, :"$1", {:const, id}}]}] + {:ok, :ets.select_delete(@table, ms)} 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 - # Compiled from: - # fun do {{bucket_number, bid},_,_,_} when bid == ^id -> true end - count_deleted = - :ets.select_delete(@ets_table_name, [ - {{{:"$1", :"$2"}, :_, :_, :_}, [{:==, :"$2", id}], [true]} + @impl GenServer + def init(opts) do + cleanup_interval_ms = Keyword.fetch!(opts, :cleanup_interval_ms) + + @table = + :ets.new(@table, [ + :named_table, + :set, + :public, + {:read_concurrency, true}, + {:write_concurrency, true}, + {:decentralized_counters, true} ]) - {:ok, count_deleted} - rescue - e -> - {:error, e} + schedule(cleanup_interval_ms) + {:ok, %{cleanup_interval_ms: cleanup_interval_ms}} end - ## GenServer Callbacks - - def init(args) do - cleanup_interval_ms = Keyword.get(args, :cleanup_interval_ms) - expiry_ms = Keyword.get(args, :expiry_ms) - - if !expiry_ms do - raise RuntimeError, "Missing required config: expiry_ms" - end - - if !cleanup_interval_ms do - raise RuntimeError, "Missing required config: cleanup_interval_ms" - end - - case :ets.info(@ets_table_name) do - :undefined -> - :ets.new(@ets_table_name, [:named_table, :ordered_set, :public]) - :timer.send_interval(cleanup_interval_ms, :prune) - - _ -> - nil - end - - state = %{ - cleanup_interval_ms: cleanup_interval_ms, - expiry_ms: expiry_ms - } - - {:ok, state} + @impl GenServer + def handle_info(:clean, state) do + cleanup() + schedule(state.cleanup_interval_ms) + {:noreply, state} end - def handle_call(:stop, _from, state) do - {:stop, :normal, :ok, state} + defp cleanup do + now = System.system_time(:millisecond) + ms = [{{{:_, :_}, :_, :"$1"}, [], [{:<, :"$1", {:const, now}}]}] + :ets.select_delete(@table, ms) end - def handle_info(:prune, state) do - %{expiry_ms: expiry_ms} = state - now = Utils.timestamp() - expire_before = now - expiry_ms - - :ets.select_delete(@ets_table_name, [ - {{:_, :_, :_, :"$1"}, [{:<, :"$1", expire_before}], [true]} - ]) - - {:noreply, state} + defp schedule(period) do + Process.send_after(self(), :clean, period) end end diff --git a/lib/hammer/supervisor.ex b/lib/hammer/supervisor.ex deleted file mode 100644 index 60b88e3..0000000 --- a/lib/hammer/supervisor.ex +++ /dev/null @@ -1,51 +0,0 @@ -defmodule Hammer.Supervisor do - @moduledoc """ - Top-level Supervisor for the Hammer application. - - Starts a set of poolboy pools based on provided configuration, - which are latter called to by the `Hammer` module. - See the Application module for configuration examples. - """ - - use Supervisor - - def start_link(config, opts) do - Supervisor.start_link(__MODULE__, config, opts) - end - - # Single backend - def init(config) when is_tuple(config) do - children = [ - to_pool_spec(:hammer_backend_single_pool, 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_existing_atom() - |> to_pool_spec(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 deleted file mode 100644 index c4a7749..0000000 --- a/lib/hammer/utils.ex +++ /dev/null @@ -1,51 +0,0 @@ -defmodule Hammer.Utils do - @moduledoc false - - def pool_name do - pool_name(:single) - end - - def pool_name(name) do - String.to_existing_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) - end - - # Returns tuple of {timestamp, key}, where key is {bucket_number, id} - def stamp_key(id, scale_ms) do - stamp = timestamp() - # with scale_ms = 1 bucket changes every millisecond - bucket_number = trunc(stamp / scale_ms) - key = {bucket_number, id} - {stamp, key} - end - - def get_backend_module(:single) do - case Application.get_env(:hammer, :backend) do - {backend_module, _config} -> - backend_module - - nil -> - raise RuntimeError, "Hammer :backend not configured" - - _ -> - raise RuntimeError, "trying to get single backend, but multiple backends configured" - end - end - - def get_backend_module(which) do - case Application.get_env(:hammer, :backend)[which] do - {backend_module, _config} -> - backend_module - - nil -> - raise RuntimeError, "Hammer :backend not configured" - - _ -> - raise RuntimeError, "backend #{which} is not configured" - end - end -end diff --git a/mix.exs b/mix.exs index 1abe29f..689260d 100644 --- a/mix.exs +++ b/mix.exs @@ -1,17 +1,15 @@ -defmodule Hammer.Mixfile do +defmodule Hammer.MixProject do use Mix.Project @source_url "https://github.com/ExHammer/hammer" - @version "6.1.0" + @version "7.0.0" def project do [ app: :hammer, description: "A rate-limiter with plugable backends.", version: @version, - elixir: "~> 1.6", - build_embedded: Mix.env() == :prod, - start_permanent: Mix.env() == :prod, + elixir: "~> 1.15", deps: deps(), docs: docs(), package: package(), @@ -21,15 +19,14 @@ defmodule Hammer.Mixfile do def application do # Specify extra applications you'll use from Erlang/Elixir - [mod: {Hammer.Application, []}, extra_applications: [:logger, :runtime_tools]] + [extra_applications: [:logger]] end defp deps 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 11b1e52..43493bc 100644 --- a/mix.lock +++ b/mix.lock @@ -11,5 +11,4 @@ "makeup_elixir": {:hex, :makeup_elixir, "0.16.1", "cc9e3ca312f1cfeccc572b37a09980287e243648108384b97ff2b76e505c3555", [: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", "e127a341ad1b209bd80f7bd1620a15693a9908ed780c3b763bccf7d200c767c6"}, "makeup_erlang": {:hex, :makeup_erlang, "0.1.2", "ad87296a092a46e03b7e9b0be7631ddcf64c790fa68a9ef5323b6cbb36affc72", [:mix], [{:makeup, "~> 1.0", [hex: :makeup, repo: "hexpm", optional: false]}], "hexpm", "f3f5a1ca93ce6e092d92b6d9c049bcda58a3b617a8d888f8e7231c85630e8108"}, "nimble_parsec": {:hex, :nimble_parsec, "1.3.1", "2c54013ecf170e249e9291ed0a62e5832f70a476c61da16f6aac6dca0189f2af", [:mix], [], "hexpm", "2682e3c0b2eb58d90c6375fc0cc30bc7be06f365bf72608804fb9cffa5e1b167"}, - "poolboy": {:hex, :poolboy, "1.5.2", "392b007a1693a64540cead79830443abf5762f5d30cf50bc95cb2c1aaafa006b", [:rebar3], [], "hexpm", "dad79704ce5440f3d5a3681c8590b9dc25d1a561e8f5a9c995281012860901e3"}, } diff --git a/test/hammer/backend/ets_test.exs b/test/hammer/backend/ets_test.exs new file mode 100644 index 0000000..82f82be --- /dev/null +++ b/test/hammer/backend/ets_test.exs @@ -0,0 +1,18 @@ +defmodule Hammer.Backend.ETSTest do + use ExUnit.Case + + test "pruning" do + start_supervised!({Hammer.Backend.ETS, cleanup_interval_ms: 100}) + Hammer.check_rate("something-pruned", _scale_ms = 100, _limit = 10) + + assert [{{"something-pruned", _bucket}, _count = 1, expires_at}] = + :ets.tab2list(:hammer_ets_buckets) + + assert expires_at > System.system_time(:millisecond) + assert expires_at < System.system_time(:millisecond) + 100 + + :timer.sleep(_ms = 200) + + assert :ets.tab2list(:hammer_ets_buckets) == [] + end +end diff --git a/test/hammer_ets_test.exs b/test/hammer_ets_test.exs deleted file mode 100644 index 787763a..0000000 --- a/test/hammer_ets_test.exs +++ /dev/null @@ -1,65 +0,0 @@ -defmodule ETSTest do - use ExUnit.Case - - alias Hammer.Backend.ETS - alias Hammer.Utils - - @table_name :hammer_ets_buckets - - setup _context do - case :ets.info(@table_name) do - :undefined -> - nil - - _ -> - :ets.delete(@table_name) - end - - opts = [expiry_ms: 5, cleanup_interval_ms: 5] - {:ok, hammer_ets_pid} = start_supervised({ETS, opts}) - {:ok, Keyword.put(opts, :pid, hammer_ets_pid)} - end - - test "count_hit", context do - pid = context[:pid] - {stamp, key} = Utils.stamp_key("one", 200_000) - assert {:ok, 1} = ETS.count_hit(pid, key, stamp) - assert {:ok, 2} = ETS.count_hit(pid, key, stamp) - assert {:ok, 3} = ETS.count_hit(pid, key, stamp) - end - - test "get_bucket", context do - pid = context[:pid] - {stamp, key} = Utils.stamp_key("two", 200_000) - # With no hits - assert {:ok, nil} = ETS.get_bucket(pid, key) - # With one hit - assert {:ok, 1} = ETS.count_hit(pid, key, stamp) - assert {:ok, {{_, "two"}, 1, _, _}} = ETS.get_bucket(pid, key) - # With two hits - assert {:ok, 2} = ETS.count_hit(pid, key, stamp) - assert {:ok, {{_, "two"}, 2, _, _}} = ETS.get_bucket(pid, key) - end - - test "delete_buckets", context do - pid = context[:pid] - {stamp, key} = Utils.stamp_key("three", 200_000) - # With no hits - assert {:ok, 0} = ETS.delete_buckets(pid, "three") - # With three hits in same bucket - assert {:ok, 1} = ETS.count_hit(pid, key, stamp) - assert {:ok, 2} = ETS.count_hit(pid, key, stamp) - assert {:ok, 3} = ETS.count_hit(pid, key, stamp) - assert {:ok, 1} = ETS.delete_buckets(pid, "three") - end - - test "timeout pruning", context do - pid = context[:pid] - expiry_ms = context[:expiry_ms] - {stamp, key} = Utils.stamp_key("something-pruned", 200_000) - assert {:ok, 1} = ETS.count_hit(pid, key, stamp) - assert {:ok, {{_, "something-pruned"}, 1, _, _}} = ETS.get_bucket(pid, key) - :timer.sleep(expiry_ms * 5) - assert {:ok, nil} = ETS.get_bucket(pid, key) - end -end diff --git a/test/hammer_test.exs b/test/hammer_test.exs index d07177f..b8e7425 100644 --- a/test/hammer_test.exs +++ b/test/hammer_test.exs @@ -1,79 +1,35 @@ defmodule HammerTest do - use ExUnit.Case, async: false + use ExUnit.Case - setup _context do - pool = Hammer.Utils.pool_name() - - opts = [ - name: {:local, :toto}, - worker_module: Hammer.Backend.ETS, - size: 4, - max_overflow: 4 - ] - - worker_args = [ - expiry_ms: 5001, - cleanup_interval_ms: 6002 - ] - - child_spec = - {:local, pool} - |> :poolboy.child_spec(opts, worker_args) - |> tuple_spec_to_map_spec() - - {:ok, _pid} = start_supervised(child_spec) - {:ok, [pool: pool, bucket: "bucket_#{:rand.uniform()}"]} - end - - # Poolboy returns an old tuple-based child spec - # Supervisor empirically works with it, but child specs should be maps - # Poolboy has child_spec/4 to return maps, but it has not yet been published - # https://github.com/devinus/poolboy/blob/master/src/poolboy.erl#L109 - @spec tuple_spec_to_map_spec(:supervisor.child_spec()) :: Supervisor.child_spec() - defp tuple_spec_to_map_spec({id, start, restart, shutdown, type, modules}) do - %{ - id: id, - start: start, - restart: restart, - shutdown: shutdown, - type: type, - modules: modules - } + setup do + start_supervised!({Hammer.Backend.ETS, cleanup_interval_ms: :timer.seconds(60)}) + :ok end - test "make_rate_checker" do - check = Hammer.make_rate_checker("some-prefix:", 10_000, 2) - assert {:allow, 1} = check.("aaa") - assert {:allow, 2} = check.("aaa") - assert {:deny, 2} = check.("aaa") - assert {:deny, 2} = check.("aaa") - assert {:allow, 1} = check.("bbb") - assert {:allow, 2} = check.("bbb") - assert {:deny, 2} = check.("bbb") - assert {:deny, 2} = check.("bbb") - end + defp bucket, do: "bucket:#{System.unique_integer([:positive])}" - test "returns {:ok, 1} tuple on first access", %{bucket: bucket} do - assert {:allow, 1} = Hammer.check_rate(bucket, 10_000, 10) + test "returns {:ok, 1} tuple on first access" do + assert {:allow, 1} = Hammer.check_rate(bucket(), 10_000, 10) end - test "returns {:ok, 4} tuple on in-limit checks", %{bucket: bucket} do + test "returns {:ok, 4} tuple on in-limit checks" do + bucket = bucket() assert {:allow, 1} = Hammer.check_rate(bucket, 10_000, 10) assert {:allow, 2} = Hammer.check_rate(bucket, 10_000, 10) assert {:allow, 3} = Hammer.check_rate(bucket, 10_000, 10) assert {:allow, 4} = Hammer.check_rate(bucket, 10_000, 10) end - test "returns expected tuples on mix of in-limit and out-of-limit checks", %{bucket: bucket} do + test "returns expected tuples on mix of in-limit and out-of-limit checks" do + bucket = bucket() assert {:allow, 1} = Hammer.check_rate(bucket, 10_000, 2) assert {:allow, 2} = Hammer.check_rate(bucket, 10_000, 2) assert {:deny, 2} = Hammer.check_rate(bucket, 10_000, 2) assert {:deny, 2} = Hammer.check_rate(bucket, 10_000, 2) end - test "returns expected tuples on 1000ms bucket check with a sleep in the middle", %{ - bucket: bucket - } do + test "returns expected tuples on 1000ms bucket check with a sleep in the middle" do + bucket = bucket() assert {:allow, 1} = Hammer.check_rate(bucket, 1000, 2) assert {:allow, 2} = Hammer.check_rate(bucket, 1000, 2) assert {:deny, 2} = Hammer.check_rate(bucket, 1000, 2) @@ -84,46 +40,54 @@ defmodule HammerTest do end test "returns expected tuples on inspect_bucket" do - assert {:ok, {0, 2, _, nil, nil}} = Hammer.inspect_bucket("inspect_bucket11", 1000, 2) - assert {:allow, 1} = Hammer.check_rate("inspect_bucket11", 1000, 2) - assert {:ok, {1, 1, _, _, _}} = Hammer.inspect_bucket("inspect_bucket11", 1000, 2) - assert {:allow, 2} = Hammer.check_rate("inspect_bucket11", 1000, 2) - assert {:allow, 1} = Hammer.check_rate("inspect_bucket22", 1000, 2) - assert {:ok, {2, 0, _, _, _}} = Hammer.inspect_bucket("inspect_bucket11", 1000, 2) - assert {:deny, 2} = Hammer.check_rate("inspect_bucket11", 1000, 2) + bucket1 = bucket() + bucket2 = bucket() - assert {:ok, {3, 0, ms_to_next_bucket, _, _}} = - Hammer.inspect_bucket("inspect_bucket11", 1000, 2) + assert {:ok, {0, 2, _, nil, nil}} = Hammer.inspect_bucket(bucket1, 1000, 2) + assert {:allow, 1} = Hammer.check_rate(bucket1, 1000, 2) + assert {:ok, {1, 1, _, _, _}} = Hammer.inspect_bucket(bucket1, 1000, 2) + assert {:allow, 2} = Hammer.check_rate(bucket1, 1000, 2) + + assert {:allow, 1} = Hammer.check_rate(bucket2, 1000, 2) + assert {:ok, {2, 0, _, _, _}} = Hammer.inspect_bucket(bucket1, 1000, 2) + assert {:deny, 2} = Hammer.check_rate(bucket1, 1000, 2) + assert {:ok, {3, 0, ms_to_next_bucket, _, _}} = Hammer.inspect_bucket(bucket1, 1000, 2) assert ms_to_next_bucket < 1000 end test "returns expected tuples on delete_buckets" do - assert {:allow, 1} = Hammer.check_rate("my-bucket1", 1000, 2) - assert {:allow, 2} = Hammer.check_rate("my-bucket1", 1000, 2) - assert {:deny, 2} = Hammer.check_rate("my-bucket1", 1000, 2) - assert {:allow, 1} = Hammer.check_rate("my-bucket2", 1000, 2) - assert {:allow, 2} = Hammer.check_rate("my-bucket2", 1000, 2) - assert {:deny, 2} = Hammer.check_rate("my-bucket2", 1000, 2) - assert {:ok, 1} = Hammer.delete_buckets("my-bucket1") - assert {:allow, 1} = Hammer.check_rate("my-bucket1", 1000, 2) - assert {:deny, 2} = Hammer.check_rate("my-bucket2", 1000, 2) - - assert {:ok, 0} = Hammer.delete_buckets("unknown-bucket") + bucket1 = bucket() + bucket2 = bucket() + bucket3 = bucket() + + assert {:allow, 1} = Hammer.check_rate(bucket1, 1000, 2) + assert {:allow, 2} = Hammer.check_rate(bucket1, 1000, 2) + assert {:deny, 2} = Hammer.check_rate(bucket1, 1000, 2) + assert {:allow, 1} = Hammer.check_rate(bucket2, 1000, 2) + assert {:allow, 2} = Hammer.check_rate(bucket2, 1000, 2) + assert {:deny, 2} = Hammer.check_rate(bucket2, 1000, 2) + assert {:ok, 1} = Hammer.delete_buckets(bucket1) + assert {:allow, 1} = Hammer.check_rate(bucket1, 1000, 2) + assert {:deny, 2} = Hammer.check_rate(bucket2, 1000, 2) + + assert {:ok, 0} = Hammer.delete_buckets(bucket3) end test "count_hit_inc" do - assert {:allow, 4} = Hammer.check_rate_inc("cost-bucket1", 1000, 10, 4) - assert {:allow, 9} = Hammer.check_rate_inc("cost-bucket1", 1000, 10, 5) - assert {:deny, 10} = Hammer.check_rate_inc("cost-bucket1", 1000, 10, 3) + bucket = bucket() + assert {:allow, 4} = Hammer.check_rate_inc(bucket, 1000, 10, 4) + assert {:allow, 9} = Hammer.check_rate_inc(bucket, 1000, 10, 5) + assert {:deny, 10} = Hammer.check_rate_inc(bucket, 1000, 10, 3) end test "mixing count_hit with count_hit_inc" do - assert {:allow, 3} = Hammer.check_rate_inc("cost-bucket2", 1000, 10, 3) - assert {:allow, 4} = Hammer.check_rate("cost-bucket2", 1000, 10) - assert {:allow, 5} = Hammer.check_rate("cost-bucket2", 1000, 10) - assert {:allow, 9} = Hammer.check_rate_inc("cost-bucket2", 1000, 10, 4) - assert {:allow, 10} = Hammer.check_rate("cost-bucket2", 1000, 10) - assert {:deny, 10} = Hammer.check_rate_inc("cost-bucket2", 1000, 10, 2) + bucket = bucket() + assert {:allow, 3} = Hammer.check_rate_inc(bucket, 1000, 10, 3) + assert {:allow, 4} = Hammer.check_rate(bucket, 1000, 10) + assert {:allow, 5} = Hammer.check_rate(bucket, 1000, 10) + assert {:allow, 9} = Hammer.check_rate_inc(bucket, 1000, 10, 4) + assert {:allow, 10} = Hammer.check_rate(bucket, 1000, 10) + assert {:deny, 10} = Hammer.check_rate_inc(bucket, 1000, 10, 2) end end diff --git a/test/hammer_util_test.exs b/test/hammer_util_test.exs deleted file mode 100644 index e7939b4..0000000 --- a/test/hammer_util_test.exs +++ /dev/null @@ -1,37 +0,0 @@ -defmodule UtilsTest do - use ExUnit.Case - - test "timestamp" do - assert is_integer(Hammer.Utils.timestamp()) - end - - test "stamp_key" do - id = "test_one_two" - {stamp, key} = Hammer.Utils.stamp_key(id, 60_000) - assert is_integer(stamp) - assert is_tuple(key) - {bucket_number, b_id} = key - assert is_integer(bucket_number) - assert b_id == id - end - - test "get_backend_module" do - # With :single and default backend config - assert Hammer.Utils.get_backend_module(:single) == Hammer.Backend.ETS - - # With :single and configured backend config - Application.put_env(:hammer, :backend, {Hammer.Backend.SomeBackend, []}) - assert Hammer.Utils.get_backend_module(:single) == Hammer.Backend.SomeBackend - - # with a specific backend config - Application.put_env(:hammer, :backend, one: {Hammer.Backend.SomeBackend, []}) - assert Hammer.Utils.get_backend_module(:one) == Hammer.Backend.SomeBackend - - # with an erroneus backend key - Application.put_env(:hammer, :backend, one: {Hammer.Backend.SomeBackend, []}) - - assert_raise RuntimeError, fn -> - Hammer.Utils.get_backend_module(:no_not_real) == Hammer.Backend.SomeBackend - end - end -end