├── .formatter.exs ├── .gitignore ├── .tool-versions ├── LICENSE ├── README.md ├── config └── config.exs ├── lib ├── ex_limiter.ex └── ex_limiter │ ├── base.ex │ ├── bucket.ex │ ├── plug.ex │ ├── storage.ex │ ├── storage │ ├── memcache.ex │ ├── pg2_shard.ex │ └── pg2_shard │ │ ├── pruner.ex │ │ ├── router.ex │ │ ├── shutdown.ex │ │ ├── supervisor.ex │ │ └── worker.ex │ └── utils.ex ├── mix.exs ├── mix.lock └── test ├── ex_limiter ├── plug_test.exs └── storage │ ├── pg2_shard_test.exs │ └── pruner_test.exs ├── ex_limiter_test.exs ├── support ├── pg2_limiter.ex └── test_utils.ex └── test_helper.exs /.formatter.exs: -------------------------------------------------------------------------------- 1 | [ 2 | inputs: [ 3 | "mix.exs", 4 | ".formatter.exs", 5 | "config/*.exs", 6 | "lib/**/*.ex" 7 | ], 8 | line_length: 120, 9 | plugins: [Styler], 10 | import_deps: [:plug] 11 | ] 12 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # The directory Mix will write compiled artifacts to. 2 | /_build/ 3 | 4 | # If you run "mix test --cover", coverage assets end up here. 5 | /cover/ 6 | 7 | # The directory Mix downloads your dependencies sources to. 8 | /deps/ 9 | 10 | # Where 3rd-party dependencies like ExDoc output generated docs. 11 | /doc/ 12 | 13 | # Ignore .fetch files in case you like to edit your project deps locally. 14 | /.fetch 15 | 16 | # If the VM crashes, it generates a dump, let's ignore it too. 17 | erl_crash.dump 18 | 19 | # Also ignore archive artifacts (built via "mix archive.build"). 20 | *.ez 21 | 22 | /tags 23 | /.elixir_ls/ 24 | -------------------------------------------------------------------------------- /.tool-versions: -------------------------------------------------------------------------------- 1 | elixir 1.18.2-otp-27 2 | erlang 27.2.3 3 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2018 Frame.io 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # ex_limiter 2 | Rate Limiter written in elixir with configurable backends 3 | 4 | Implements leaky bucket rate limiting ([wiki](https://en.wikipedia.org/wiki/Leaky_bucket)), which is superior to most naive approaches by handling bursts even around time windows. You can define your own storage backend by implementing the `ExLimiter.Storage` behaviour, and configuring it with 5 | 6 | ```elixir 7 | config :ex_limiter, :storage, MyStorage 8 | ``` 9 | 10 | usage once configured is: 11 | 12 | ```elixir 13 | case ExLimiter.consume(bucket, 1, scale: 1000, limit: 5) do 14 | {:ok, bucket} -> #do some work 15 | {:error, :rate_limited} -> #fail 16 | end 17 | ``` 18 | 19 | Additionally, if you want to have multiple rate limiters with diverse backend implementations you can use the `ExLimiter.Base` macro, like so: 20 | 21 | ```elixir 22 | defmodule MyLimiter do 23 | use ExLimiter.Base, storage: MyStorage 24 | end 25 | ``` 26 | 27 | ## ExLimiter.Plug 28 | 29 | ExLimiter also ships with a simple plug implementation. Usage is 30 | 31 | ```elixir 32 | plug ExLimiter.Plug, scale: 5000, limit: 20 33 | ``` 34 | 35 | You can also configure how the bucket is inferred from the given conn, how many tokens to consume and what limiter to use. -------------------------------------------------------------------------------- /config/config.exs: -------------------------------------------------------------------------------- 1 | import Config 2 | 3 | config :ex_limiter, ExLimiter.Plug, 4 | limiter: ExLimiter, 5 | fallback: ExLimiter.Plug, 6 | limit: 10, 7 | scale: 1000 8 | 9 | config :ex_limiter, ExLimiter.Storage.PG2Shard, shard_count: 20 10 | config :ex_limiter, :storage, ExLimiter.Storage.Memcache 11 | -------------------------------------------------------------------------------- /lib/ex_limiter.ex: -------------------------------------------------------------------------------- 1 | defmodule ExLimiter do 2 | @moduledoc """ 3 | Configurable, leaky bucket rate limiting. You can define your own storage backend by 4 | implementing the `ExLimiter.Storage` behaviour, and configuring it with 5 | 6 | ``` 7 | config :ex_limiter, :storage, MyStorage 8 | ``` 9 | 10 | usage once configured is: 11 | 12 | ``` 13 | case ExLimiter.consume(bucket, 1, scale: 1000, limit: 5) do 14 | {:ok, bucket} -> #do some work 15 | {:error, :rate_limited} -> #fail 16 | end 17 | ``` 18 | 19 | Additionally, if you want to have multiple rate limiters with diverse backend implementations, 20 | you can use the `ExLimiter.Base` macro, like so: 21 | 22 | ``` 23 | defmodule MyLimiter do 24 | use ExLimiter.Base, storage: MyStorage 25 | end 26 | ``` 27 | """ 28 | use ExLimiter.Base, storage: Application.get_env(:ex_limiter, :storage) 29 | end 30 | -------------------------------------------------------------------------------- /lib/ex_limiter/base.ex: -------------------------------------------------------------------------------- 1 | defmodule ExLimiter.Base do 2 | @moduledoc """ 3 | Base module for arbitrary rate limiter implementations. Usage is: 4 | 5 | ``` 6 | defmodule MyLimiter do 7 | use ExLimiterBase, storage: MyCustomStorage 8 | end 9 | ``` 10 | """ 11 | alias ExLimiter.Bucket 12 | alias ExLimiter.Utils 13 | 14 | defmacro __using__(storage: storage) do 15 | quote do 16 | import ExLimiter.Base 17 | 18 | @storage unquote(storage) 19 | 20 | def remaining(%Bucket{value: val}, opts \\ []) do 21 | limit = Keyword.get(opts, :limit, 10) 22 | scale = Keyword.get(opts, :scale, 1000) 23 | 24 | round(max(scale - val, 0) / (scale / limit)) 25 | end 26 | 27 | @doc """ 28 | Consumes `amount` from the rate limiter aliased by bucket. 29 | 30 | `opts` params are: 31 | * `:limit` - the maximum amount for the rate limiter (default 10) 32 | * `:scale` - the duration under which `:limit` applies in milliseconds 33 | """ 34 | @spec consume(bucket :: binary, amount :: integer, opts :: keyword) :: {:ok, Bucket.t()} | {:error, :rate_limited} 35 | def consume(bucket, amount \\ 1, opts \\ []), do: consume(@storage, bucket, amount, opts) 36 | 37 | def delete(bucket), do: @storage.delete(%Bucket{key: bucket}) 38 | end 39 | end 40 | 41 | @doc """ 42 | Delegate function for rate limiter implementations 43 | """ 44 | @spec consume(atom, binary, integer, keyword) :: {:ok, Bucket.t()} | {:error, :rate_limited} 45 | def consume(storage, bucket, amount, opts) do 46 | limit = Keyword.get(opts, :limit, 10) 47 | scale = Keyword.get(opts, :scale, 1000) 48 | 49 | mult = scale / limit 50 | incr = round(amount * mult) 51 | 52 | storage.leak_and_consume( 53 | bucket, 54 | fn %Bucket{value: value, last: time} = b -> 55 | now = Utils.now() 56 | amount = max(value - (now - time), 0) 57 | 58 | %{b | last: now, value: amount} 59 | end, 60 | fn 61 | %Bucket{value: v} = b when v + incr <= scale -> b 62 | _ -> {:error, :rate_limited} 63 | end, 64 | incr 65 | ) 66 | end 67 | end 68 | -------------------------------------------------------------------------------- /lib/ex_limiter/bucket.ex: -------------------------------------------------------------------------------- 1 | defmodule ExLimiter.Bucket do 2 | @moduledoc false 3 | alias ExLimiter.Utils 4 | 5 | @type t :: %__MODULE__{} 6 | 7 | defstruct key: nil, 8 | value: 0, 9 | last: nil, 10 | version: %{} 11 | 12 | def new(key), do: %__MODULE__{key: key, last: Utils.now()} 13 | 14 | def new(contents, key) when is_map(contents) do 15 | struct(__MODULE__, Map.put(contents, :key, key)) 16 | end 17 | end 18 | -------------------------------------------------------------------------------- /lib/ex_limiter/plug.ex: -------------------------------------------------------------------------------- 1 | defmodule ExLimiter.Plug do 2 | @moduledoc """ 3 | Plug for enforcing rate limits. The usage should be something like 4 | 5 | ``` 6 | plug ExLimiter.Plug, scale: 1000, limit: 5 7 | ``` 8 | 9 | Additionally, you can pass the following options: 10 | 11 | - `:bucket`, a 1-arity function of a `Plug.Conn.t` which determines 12 | the bucket for the rate limit. Defaults to the phoenix controller, 13 | action and remote_ip. 14 | 15 | - `:consumes`, a 1-arity function of a `Plug.Conn.t` which determines 16 | the amount to consume. Defaults to 1 respectively. 17 | 18 | - `:decorate`, a 2-arity function which can return an updated conn 19 | based on the outcome of the limiter call. The first argument is the 20 | `Plug.Conn.t`, and the second can be: 21 | 22 | - `{:ok, Bucket.t}` 23 | - `{:rate_limited, binary}` Where the second element is the bucket 24 | name that triggered the rate limit. 25 | 26 | Additionally, you can configure a custom limiter with 27 | 28 | ``` 29 | config :ex_limiter, ExLimiter.Plug, limiter: MyLimiter 30 | ``` 31 | 32 | and you can also configure the rate limited response with 33 | 34 | ``` 35 | config :ex_limiter, ExLimiter.Plug, fallback: MyFallback 36 | ``` 37 | 38 | `MyFallback` needs to implement a function `render_error(conn, :rate_limited)` 39 | """ 40 | import Plug.Conn 41 | 42 | @limiter Application.get_env(:ex_limiter, __MODULE__)[:limiter] 43 | 44 | defmodule Config do 45 | @moduledoc false 46 | @limit Application.get_env(:ex_limiter, ExLimiter.Plug)[:limit] 47 | @scale Application.get_env(:ex_limiter, ExLimiter.Plug)[:scale] 48 | @fallback Application.get_env(:ex_limiter, ExLimiter.Plug)[:fallback] 49 | 50 | defstruct scale: @scale, 51 | limit: @limit, 52 | bucket: &ExLimiter.Plug.get_bucket/1, 53 | consumes: nil, 54 | decorate: nil, 55 | fallback: @fallback 56 | 57 | def new(opts) do 58 | contents = 59 | opts 60 | |> Map.new() 61 | |> Map.put_new(:consumes, fn _ -> 1 end) 62 | |> Map.put_new(:decorate, &ExLimiter.Plug.decorate/2) 63 | 64 | struct(__MODULE__, contents) 65 | end 66 | end 67 | 68 | def get_bucket(%{private: %{phoenix_controller: contr, phoenix_action: ac}} = conn) do 69 | "#{contr}.#{ac}.#{ip(conn)}" 70 | end 71 | 72 | def render_error(conn, :rate_limited) do 73 | conn 74 | |> resp(429, "Rate Limit Exceeded") 75 | |> halt() 76 | end 77 | 78 | @spec decorate(Plug.Conn.t(), {:ok, Bucket.t()} | {:rate_limited, bucket_name :: binary}) :: Plug.Conn.t() 79 | def decorate(conn, _), do: conn 80 | 81 | def init(opts), do: Config.new(opts) 82 | 83 | def call(conn, %Config{ 84 | bucket: bucket_fun, 85 | scale: scale, 86 | limit: limit, 87 | consumes: consume_fun, 88 | decorate: decorate_fun, 89 | fallback: fallback 90 | }) do 91 | bucket_name = bucket_fun.(conn) 92 | 93 | bucket_name 94 | |> @limiter.consume(consume_fun.(conn), scale: scale, limit: limit) 95 | |> case do 96 | {:ok, bucket} = response -> 97 | remaining = @limiter.remaining(bucket, scale: scale, limit: limit) 98 | 99 | conn 100 | |> put_resp_header("x-ratelimit-limit", to_string(limit)) 101 | |> put_resp_header("x-ratelimit-window", to_string(scale)) 102 | |> put_resp_header("x-ratelimit-remaining", to_string(remaining)) 103 | |> decorate_fun.(response) 104 | 105 | {:error, :rate_limited} -> 106 | conn 107 | |> decorate_fun.({:rate_limited, bucket_name}) 108 | |> fallback.render_error(:rate_limited) 109 | end 110 | end 111 | 112 | defp ip(conn), do: conn.remote_ip |> Tuple.to_list() |> Enum.join(".") 113 | end 114 | -------------------------------------------------------------------------------- /lib/ex_limiter/storage.ex: -------------------------------------------------------------------------------- 1 | defmodule ExLimiter.Storage do 2 | @moduledoc false 3 | alias ExLimiter.Bucket 4 | 5 | defmacro __using__(_) do 6 | quote do 7 | @behaviour ExLimiter.Storage 8 | 9 | alias ExLimiter.Bucket 10 | 11 | def leak_and_consume(bucket, update_fn, boundary_fn, incr) do 12 | with %Bucket{} = bucket <- update(bucket, update_fn), 13 | %Bucket{} = bucket <- boundary_fn.(bucket), 14 | do: consume(bucket, incr) 15 | end 16 | 17 | defoverridable leak_and_consume: 4 18 | end 19 | end 20 | 21 | @type response :: {:ok, Bucket.t()} | {:error, any} 22 | 23 | @doc """ 24 | Fetch the current state of the given bucket 25 | """ 26 | @callback fetch(bucket :: Bucket.t()) :: Bucket.t() 27 | 28 | @doc """ 29 | Set the current state of the given bucket. Specify hard if you want to 30 | force a write 31 | """ 32 | @callback refresh(bucket :: Bucket.t()) :: response 33 | @callback refresh(bucket :: Bucket.t(), type :: :hard | :soft) :: response 34 | 35 | @doc """ 36 | Atomically update the bucket denoted by `key` with `fun`. Leverage whatever 37 | concurrency controls are available in the given storage mechanism (eg cas for memcached) 38 | """ 39 | @callback update(key :: binary, fun :: (Bucket.t() -> Bucket.t())) :: Bucket.t() 40 | 41 | @doc """ 42 | Consumes n elements from the bucket (atomically) 43 | """ 44 | @callback consume(bucket :: Bucket.t(), incr :: integer) :: {:ok, Bucket.t()} 45 | 46 | @callback delete(bucket :: Bucket.t()) :: Bucket.t() 47 | end 48 | -------------------------------------------------------------------------------- /lib/ex_limiter/storage/memcache.ex: -------------------------------------------------------------------------------- 1 | defmodule ExLimiter.Storage.Memcache do 2 | @moduledoc """ 3 | Token bucket backend written for memcache. Stores the last timestamp 4 | and amount in separate keys, and utilizes memcache increments for consumption 5 | """ 6 | use ExLimiter.Storage 7 | 8 | alias ExLimiter.Utils 9 | 10 | def fetch(%Bucket{key: key}) do 11 | key_map = keys(key) 12 | 13 | try do 14 | key_map 15 | |> Map.keys() 16 | |> Memcachir.mget(cas: true) 17 | |> case do 18 | {:ok, result} -> from_memcached(result, key, key_map) 19 | _ -> Bucket.new(key) 20 | end 21 | catch 22 | :exit, _ -> Bucket.new(key) 23 | end 24 | end 25 | 26 | def refresh(%Bucket{key: key} = bucket, _type \\ :soft) do 27 | key 28 | |> keys() 29 | |> Enum.map(&mset_command(&1, bucket)) 30 | |> Memcachir.mset_cas() 31 | |> case do 32 | {:ok, _} -> {:ok, bucket} 33 | {:error, error} -> {:error, error} 34 | end 35 | catch 36 | :exit, _ -> {:error, :memcached} 37 | end 38 | 39 | def delete(%Bucket{key: key}) do 40 | key 41 | |> keys() 42 | |> Map.keys() 43 | |> Enum.map(&Memcachir.delete/1) 44 | 45 | Bucket.new(key) 46 | catch 47 | :exit, _ -> Bucket.new(key) 48 | end 49 | 50 | def update(key, update_fun) do 51 | bucket = %Bucket{key: key} 52 | 53 | bucket 54 | |> fetch() 55 | |> update_fun.() 56 | |> refresh() 57 | |> case do 58 | {:ok, b} -> b 59 | # memcached latency is short enough that there probably wasn't much leaked here 60 | {:error, _} -> fetch(bucket) 61 | end 62 | end 63 | 64 | def consume(%Bucket{key: key} = bucket, inc) do 65 | case Memcachir.incr("amount_#{key}", inc) do 66 | {:ok, result} -> {:ok, %{bucket | value: result}} 67 | _ -> {:ok, Bucket.new(key)} 68 | end 69 | catch 70 | :exit, _ -> {:ok, Bucket.new(key)} 71 | end 72 | 73 | defp keys(key), do: %{"amount_#{key}" => :value, "last_#{key}" => :last} 74 | 75 | defp from_memcached(map, key, key_map) when is_map(map) do 76 | key_map 77 | |> Enum.reduce(%{version: %{}}, &reduce_memcached(&1, &2, map)) 78 | |> Map.new() 79 | |> Bucket.new(key) 80 | end 81 | 82 | defp mset_command({key, bucket_key}, %Bucket{version: versions} = b) do 83 | value = b |> Map.get(bucket_key) |> to_string() 84 | {key, value, Map.get(versions, bucket_key, 0)} 85 | end 86 | 87 | defp reduce_memcached({key, bucket_key}, acc, map), do: add_result(acc, bucket_key, map[key]) 88 | 89 | defp add_result(%{version: versions} = acc, bucket_key, {val, cas}) do 90 | acc 91 | |> Map.put(bucket_key, Utils.parse_integer(val)) 92 | |> Map.put(:version, Map.put(versions, bucket_key, cas)) 93 | end 94 | 95 | defp add_result(acc, bucket_key, _), do: add_result(acc, bucket_key, default(bucket_key)) 96 | 97 | defp default(:value), do: {0, 0} 98 | defp default(:last), do: {Utils.now(), 0} 99 | end 100 | -------------------------------------------------------------------------------- /lib/ex_limiter/storage/pg2_shard.ex: -------------------------------------------------------------------------------- 1 | defmodule ExLimiter.Storage.PG2Shard do 2 | @moduledoc """ 3 | Implements the leaky bucket using a fleet of GenServers discoverable via a 4 | pg2 group. 5 | 6 | To configure the pool size, do: 7 | 8 | ``` 9 | config :ex_limit, ExLimiter.Storage.PG2Shard, 10 | shard_count: 20 11 | ``` 12 | 13 | You must also include the shard supervisor in your app supervision tree, with 14 | something like: 15 | 16 | ``` 17 | ... 18 | supervise(ExLimiter.Storage.PG2Shard.Supervisor, []) 19 | ``` 20 | """ 21 | use ExLimiter.Storage 22 | 23 | alias ExLimiter.Storage.PG2Shard.Router 24 | 25 | def fetch(%Bucket{key: key}), do: with_worker(key, &call(&1, {:fetch, key})) 26 | 27 | def refresh(%Bucket{key: key} = bucket, _type \\ :soft), do: {:ok, with_worker(key, &call(&1, {:set, bucket}))} 28 | 29 | def delete(%Bucket{key: key} = bucket) do 30 | with_worker(key, &call(&1, {:delete, key})) 31 | bucket 32 | end 33 | 34 | def update(key, update_fun), do: with_worker(key, &call(&1, {:update, key, update_fun})) 35 | 36 | def consume(%Bucket{key: key}, incr), do: {:ok, with_worker(key, &call(&1, {:consume, key, incr}))} 37 | 38 | def leak_and_consume(key, update_fn, boundary_fn, incr) do 39 | with_worker( 40 | key, 41 | fn pid -> 42 | call(pid, {:leak_and_consume, key, update_fn, boundary_fn, incr}) 43 | end, 44 | {:ok, Bucket.new(key)} 45 | ) 46 | end 47 | 48 | defp call(pid, operation), do: GenServer.call(pid, operation) 49 | 50 | defp with_worker(key, fun, fallback \\ nil) do 51 | case Router.shard(key) do 52 | pid when is_pid(pid) -> fun.(pid) 53 | _ -> fallback || Bucket.new(key) 54 | end 55 | rescue 56 | _error -> fallback || Bucket.new(key) 57 | catch 58 | :exit, _ -> fallback || Bucket.new(key) 59 | end 60 | end 61 | -------------------------------------------------------------------------------- /lib/ex_limiter/storage/pg2_shard/pruner.ex: -------------------------------------------------------------------------------- 1 | defmodule ExLimiter.Storage.PG2Shard.Pruner do 2 | @moduledoc """ 3 | Responsible for creating and pruning a write optimized ets table for 4 | bucket state 5 | """ 6 | use GenServer 7 | 8 | import Ex2ms 9 | 10 | alias ExLimiter.Storage.PG2Shard 11 | alias ExLimiter.Utils 12 | 13 | @table_name :exlimiter_buckets 14 | @expiry Application.get_env(:ex_limiter, PG2Shard)[:expiry] || 10 * 60_000 15 | @eviction_count Application.get_env(:ex_limiter, PG2Shard)[:eviction_count] || 1000 16 | @max_size Application.get_env(:ex_limiter, PG2Shard)[:max_size] || 50_000 17 | @prune_interval Application.get_env(:ex_limiter, PG2Shard)[:prune_interval] || 5_000 18 | @eviction_interval Application.get_env(:ex_limiter, PG2Shard)[:eviction_interval] || 30_000 19 | 20 | def start_link(_args \\ :ok) do 21 | GenServer.start_link(__MODULE__, [], name: __MODULE__) 22 | end 23 | 24 | def init(_) do 25 | table = :ets.new(@table_name, [:set, :public, read_concurrency: true, write_concurrency: true]) 26 | prune() 27 | expire() 28 | 29 | {:ok, table} 30 | end 31 | 32 | def table, do: GenServer.call(__MODULE__, :fetch) 33 | 34 | def handle_call(:fetch, _from, table), do: {:reply, table, table} 35 | 36 | def handle_info(:expire, table) do 37 | expire() 38 | now = Utils.now() 39 | 40 | count = 41 | :ets.select_delete( 42 | table, 43 | fun do 44 | {_, updated_at, _} when updated_at < ^now - ^@expiry -> true 45 | end 46 | ) 47 | 48 | :telemetry.execute([:ex_limiter, :shards, :expirations], %{value: count}) 49 | {:noreply, table} 50 | end 51 | 52 | def handle_info(:prune, table) do 53 | prune() 54 | size = :ets.info(table, :size) 55 | 56 | if size >= @max_size do 57 | count = remove(table, @eviction_count) 58 | :telemetry.execute([:ex_limiter, :shards, :evictions], %{value: count}) 59 | end 60 | 61 | :telemetry.execute([:ex_limiter, :shards, :size], %{value: size}) 62 | {:noreply, table, :hibernate} 63 | end 64 | 65 | def remove(table, count) do 66 | Utils.batched_ets(table, {:"$1", :_, :_}, 1000, count, fn keys -> 67 | for [key] <- keys, 68 | do: :ets.delete(table, key) 69 | end) 70 | end 71 | 72 | defp prune, do: Process.send_after(self(), :prune, @prune_interval) 73 | 74 | defp expire, do: Process.send_after(self(), :expire, @eviction_interval) 75 | end 76 | -------------------------------------------------------------------------------- /lib/ex_limiter/storage/pg2_shard/router.ex: -------------------------------------------------------------------------------- 1 | defmodule ExLimiter.Storage.PG2Shard.Router do 2 | @moduledoc """ 3 | The routing mechanism for pg2 shard instances. Currently 4 | it resyncs (by calling `:pg2.get_members/1`) on nodeup/nodedown and 5 | on a fixed poll interval. 6 | 7 | Ideally we could implement a subscription mechanism in a process grouper like 8 | swarm that could make this more snappy, but since the workers are statically 9 | configured anyways, node connect/reconnect is actually a fairly reliable mechanism. 10 | """ 11 | use GenServer 12 | 13 | @process_group :ex_limiter_shards 14 | @table_name :ex_limiter_router 15 | 16 | def start_link(_args \\ :ok) do 17 | GenServer.start_link(__MODULE__, :ok, name: __MODULE__) 18 | end 19 | 20 | def init(_) do 21 | :ok = :net_kernel.monitor_nodes(true, node_type: :all) 22 | table = :ets.new(@table_name, [:set, :protected, :named_table, {:read_concurrency, true}]) 23 | :ets.insert(table, {:ring, shard_ring()}) 24 | :timer.send_interval(1000, :resync) 25 | send(self(), :resync) 26 | {:ok, table} 27 | end 28 | 29 | def shard(key) do 30 | case :ets.lookup(@table_name, :ring) do 31 | [{:ring, ring}] -> HashRing.key_to_node(ring, key) 32 | _ -> {:error, :noring} 33 | end 34 | end 35 | 36 | def handle_cast(:refresh, table) do 37 | {:noreply, regen(table)} 38 | end 39 | 40 | def handle_info({:nodeup, _, _}, table) do 41 | {:noreply, regen(table)} 42 | end 43 | 44 | def handle_info({:nodedown, _, _}, table) do 45 | {:noreply, regen(table)} 46 | end 47 | 48 | def handle_info(:resync, table) do 49 | {:noreply, regen(table)} 50 | end 51 | 52 | def shards do 53 | :pg.get_members(@process_group) 54 | end 55 | 56 | defp regen(table) do 57 | :ets.insert(table, {:ring, shard_ring()}) 58 | table 59 | end 60 | 61 | defp shard_ring, do: HashRing.add_nodes(HashRing.new(), shards()) 62 | end 63 | -------------------------------------------------------------------------------- /lib/ex_limiter/storage/pg2_shard/shutdown.ex: -------------------------------------------------------------------------------- 1 | defmodule ExLimiter.Storage.PG2Shard.Shutdown do 2 | @moduledoc """ 3 | Traps exits and notifies other nodes to resync on shutdowns. 4 | """ 5 | use GenServer 6 | 7 | alias ExLimiter.Storage.PG2Shard.Router 8 | alias ExLimiter.Storage.PG2Shard.Worker 9 | 10 | def start_link(_ \\ :ok) do 11 | GenServer.start_link(__MODULE__, [], name: __MODULE__) 12 | end 13 | 14 | def init(_) do 15 | Process.flag(:trap_exit, true) 16 | {:ok, []} 17 | end 18 | 19 | def register(pid), do: GenServer.cast(__MODULE__, {:register, pid}) 20 | 21 | def handle_cast({:register, pid}, pids) do 22 | {:noreply, [pid | pids]} 23 | end 24 | 25 | def terminate(_, pids) do 26 | Enum.each(pids, &:pg.leave(Worker.group(), &1)) 27 | Enum.each(Node.list(), &GenServer.cast({Router, &1}, :refresh)) 28 | :timer.sleep(5_000) 29 | end 30 | end 31 | -------------------------------------------------------------------------------- /lib/ex_limiter/storage/pg2_shard/supervisor.ex: -------------------------------------------------------------------------------- 1 | defmodule ExLimiter.Storage.PG2Shard.Supervisor do 2 | @moduledoc """ 3 | Supervisor for the workers and shard router for the PG2Shard 4 | backend. 5 | 6 | This *must* be manually specified in your supervision tree for this 7 | storage backend to work. 8 | """ 9 | use Supervisor 10 | 11 | alias ExLimiter.Storage.PG2Shard 12 | alias ExLimiter.Storage.PG2Shard.Pruner 13 | alias ExLimiter.Storage.PG2Shard.Router 14 | alias ExLimiter.Storage.PG2Shard.Shutdown 15 | alias ExLimiter.Storage.PG2Shard.Worker 16 | 17 | @telemetry Application.get_env(:ex_limiter, PG2Shard)[:telemetry] || Worker 18 | 19 | def start_link(_args \\ :ok) do 20 | Supervisor.start_link(__MODULE__, [], name: __MODULE__) 21 | end 22 | 23 | def init(_) do 24 | shards = [{Worker, []}] |> Stream.cycle() |> Enum.take(shard_count()) 25 | children = Enum.reverse([{Router, []} | shards]) 26 | children = [pg_spec(), {Pruner, []}, {Shutdown, []} | children] 27 | 28 | :telemetry.attach_many("exlimiter-metrics-handler", Worker.telemetry_events(), &@telemetry.handle_event/4, nil) 29 | 30 | Supervisor.init(children, strategy: :one_for_one) 31 | end 32 | 33 | def handle_event(_, _, _, _), do: :ok 34 | 35 | defp pg_spec do 36 | %{ 37 | id: :pg, 38 | start: {:pg, :start_link, []} 39 | } 40 | end 41 | 42 | defp shard_count do 43 | :ex_limiter 44 | |> Application.get_env(PG2Shard) 45 | |> Keyword.get(:shard_count, 0) 46 | end 47 | end 48 | -------------------------------------------------------------------------------- /lib/ex_limiter/storage/pg2_shard/worker.ex: -------------------------------------------------------------------------------- 1 | defmodule ExLimiter.Storage.PG2Shard.Worker do 2 | @moduledoc """ 3 | Simple Genserver for implementing the `ExLimiter.Storage` behavior for a set 4 | of buckets. 5 | 6 | Buckets are pruned after 10 minutes of inactivity, and buckets will be evicted 7 | if a maximum threshold is reached. To tune these values, use: 8 | 9 | ``` 10 | config :ex_limiter, ExLimiter.Storage.PG2Shard, 11 | max_size: 50_000, 12 | eviction_count: 1000 13 | ``` 14 | 15 | It will also publish these metrics via telemetry: 16 | 17 | ``` 18 | [:ex_limiter, :shards, :map_size], 19 | [:ex_limiter, :shards, :evictions], 20 | [:ex_limiter, :shards, :expirations] 21 | ``` 22 | 23 | You can auto-configure a telemetry handler via: 24 | 25 | ``` 26 | config :ex_limiter, ExLimiter.Storage.PG2Shard, 27 | telemetry: MyTelemetryHandler 28 | ``` 29 | """ 30 | use GenServer 31 | 32 | alias ExLimiter.Bucket 33 | alias ExLimiter.Storage.PG2Shard.Pruner 34 | 35 | @process_group :ex_limiter_shards 36 | @telemetry_events [ 37 | [:ex_limiter, :shards, :size], 38 | [:ex_limiter, :shards, :evictions], 39 | [:ex_limiter, :shards, :expirations] 40 | ] 41 | 42 | def start_link do 43 | GenServer.start_link(__MODULE__, []) 44 | end 45 | 46 | def group, do: @process_group 47 | 48 | def init(_) do 49 | :pg.join(@process_group, self()) 50 | 51 | {:ok, Pruner.table()} 52 | end 53 | 54 | def handle_call({:update, key, fun}, _from, table) do 55 | bucket = table |> fetch(key) |> fun.() 56 | {:reply, bucket, upsert(table, key, bucket)} 57 | end 58 | 59 | def handle_call({:consume, key, amount}, _from, table) do 60 | %{value: val} = bucket = fetch(table, key) 61 | bucket = %{bucket | value: val + amount} 62 | {:reply, bucket, upsert(table, key, bucket)} 63 | end 64 | 65 | def handle_call({:leak_and_consume, key, update_fn, boundary_fn, incr}, _from, table) do 66 | with %{value: val} = bucket <- table |> fetch(key) |> update_fn.(), 67 | {_old_bucket, %{} = bucket} <- {bucket, boundary_fn.(bucket)} do 68 | bucket = %{bucket | value: val + incr} 69 | {:reply, {:ok, bucket}, upsert(table, key, bucket)} 70 | else 71 | {bucket, {:error, _} = error} -> {:reply, error, upsert(table, key, bucket)} 72 | end 73 | end 74 | 75 | def handle_call({:fetch, key}, _from, table) do 76 | {:reply, fetch(table, key), table} 77 | end 78 | 79 | def handle_call({:set, %Bucket{key: k} = bucket}, _from, table) do 80 | {:reply, bucket, upsert(table, k, bucket)} 81 | end 82 | 83 | def handle_call({:delete, key}, _from, table) do 84 | :ets.delete(table, key) 85 | {:reply, :ok, table} 86 | end 87 | 88 | def child_spec(_args) do 89 | %{ 90 | id: make_ref(), 91 | start: {__MODULE__, :start_link, []} 92 | } 93 | end 94 | 95 | def handle_event(_, _, _, _), do: :ok 96 | 97 | def telemetry_events, do: @telemetry_events 98 | 99 | defp upsert(table, key, bucket) do 100 | :ets.insert(table, {key, bucket.last, bucket}) 101 | table 102 | end 103 | 104 | defp fetch(table, key) do 105 | case :ets.lookup(table, key) do 106 | [{_, _, bucket}] -> bucket 107 | _ -> Bucket.new(key) 108 | end 109 | end 110 | end 111 | -------------------------------------------------------------------------------- /lib/ex_limiter/utils.ex: -------------------------------------------------------------------------------- 1 | defmodule ExLimiter.Utils do 2 | @moduledoc false 3 | def now, do: :os.system_time(:millisecond) 4 | 5 | def batched_ets(table, match_spec \\ {:"$1", :_, :_}, batch_size \\ 1000, total \\ 100_000, fnc) do 6 | table 7 | |> :ets.match(match_spec, batch_size) 8 | |> process_batch(0, total, fnc) 9 | end 10 | 11 | defp process_batch(_, count, total, _) when count >= total, do: count 12 | 13 | defp process_batch({elem, cnt}, count, total, fnc) do 14 | fnc.(elem) 15 | 16 | cnt 17 | |> :ets.match() 18 | |> process_batch(length(elem) + count, total, fnc) 19 | end 20 | 21 | defp process_batch(:"$end_of_table", count, _, _), do: count 22 | 23 | def ets_stream(table) do 24 | Stream.resource( 25 | fn -> :ets.first(table) end, 26 | fn 27 | :"$end_of_table" -> {:halt, nil} 28 | previous_key -> {[previous_key], :ets.next(table, previous_key)} 29 | end, 30 | fn _ -> :ok end 31 | ) 32 | end 33 | 34 | def parse_integer(val) when is_binary(val), do: val |> Integer.parse() |> parse_integer() 35 | def parse_integer(val) when is_integer(val), do: val 36 | def parse_integer(:error), do: :error 37 | def parse_integer({val, _}), do: val 38 | end 39 | -------------------------------------------------------------------------------- /mix.exs: -------------------------------------------------------------------------------- 1 | defmodule ExLimiter.Mixfile do 2 | use Mix.Project 3 | 4 | @version "1.4.0" 5 | 6 | def project do 7 | [ 8 | app: :ex_limiter, 9 | version: @version, 10 | elixir: "~> 1.7", 11 | start_permanent: Mix.env() == :prod, 12 | elixirc_paths: elixirc_paths(Mix.env()), 13 | deps: deps(), 14 | description: description(), 15 | package: package(), 16 | docs: docs() 17 | ] 18 | end 19 | 20 | defp elixirc_paths(:test), do: ["lib", "test/support"] 21 | defp elixirc_paths(_), do: ["lib"] 22 | 23 | def application do 24 | [ 25 | extra_applications: [:logger] 26 | ] 27 | end 28 | 29 | defp deps do 30 | [ 31 | {:memcachir, "~> 3.3.1", [optional: true] ++ run_in_test()}, 32 | {:plug, "~> 1.4"}, 33 | {:libring, "~> 1.0"}, 34 | {:telemetry, "~> 0.4 or ~> 1.0"}, 35 | {:ex2ms, "~> 1.5"}, 36 | {:ex_doc, "~> 0.19", only: :dev}, 37 | {:styler, ">= 0.0.0", only: [:dev, :test], runtime: false} 38 | ] 39 | end 40 | 41 | defp docs do 42 | [ 43 | main: "readme", 44 | extras: ["README.md"], 45 | source_ref: "v#{@version}", 46 | source_url: "https://github.com/Frameio/ex_limiter" 47 | ] 48 | end 49 | 50 | defp run_in_test do 51 | case Mix.env() do 52 | :test -> [] 53 | _ -> [runtime: false] 54 | end 55 | end 56 | 57 | defp description do 58 | "Token bucket rate limiter written in elixir with configurable backends" 59 | end 60 | 61 | defp package do 62 | [ 63 | maintainers: ["Michael Guarino"], 64 | licenses: ["MIT"], 65 | links: %{"GitHub" => "https://github.com/Frameio/ex_limiter"} 66 | ] 67 | end 68 | end 69 | -------------------------------------------------------------------------------- /mix.lock: -------------------------------------------------------------------------------- 1 | %{ 2 | "connection": {:hex, :connection, "1.1.0", "ff2a49c4b75b6fb3e674bfc5536451607270aac754ffd1bdfe175abe4a6d7a68", [:mix], [], "hexpm", "722c1eb0a418fbe91ba7bd59a47e28008a189d47e37e0e7bb85585a016b2869c"}, 3 | "earmark_parser": {:hex, :earmark_parser, "1.4.43", "34b2f401fe473080e39ff2b90feb8ddfeef7639f8ee0bbf71bb41911831d77c5", [:mix], [], "hexpm", "970a3cd19503f5e8e527a190662be2cee5d98eed1ff72ed9b3d1a3d466692de8"}, 4 | "elasticachex": {:hex, :elasticachex, "1.1.3", "c5cc1255b3f25c53df16206959816824cc65e65be5be8462af069be59af63013", [:mix], [{:socket, "~> 0.3", [hex: :socket, repo: "hexpm", optional: false]}], "hexpm", "425814b1406729f2f037ff3b90755162b1d8b7fef23b3c23deac295e05cec2fc"}, 5 | "ex2ms": {:hex, :ex2ms, "1.7.0", "45b9f523d0b777667ded60070d82d871a37e294f0b6c5b8eca86771f00f82ee1", [:mix], [], "hexpm", "2589eee51f81f1b1caa6d08c990b1ad409215fe6f64c73f73c67d36ed10be827"}, 6 | "ex_doc": {:hex, :ex_doc, "0.37.2", "2a3aa7014094f0e4e286a82aa5194a34dd17057160988b8509b15aa6c292720c", [:mix], [{:earmark_parser, "~> 1.4.42", [hex: :earmark_parser, repo: "hexpm", optional: false]}, {:makeup_c, ">= 0.1.0", [hex: :makeup_c, repo: "hexpm", optional: true]}, {:makeup_elixir, "~> 0.14 or ~> 1.0", [hex: :makeup_elixir, repo: "hexpm", optional: false]}, {:makeup_erlang, "~> 0.1 or ~> 1.0", [hex: :makeup_erlang, repo: "hexpm", optional: false]}, {:makeup_html, ">= 0.1.0", [hex: :makeup_html, repo: "hexpm", optional: true]}], "hexpm", "4dfa56075ce4887e4e8b1dcc121cd5fcb0f02b00391fd367ff5336d98fa49049"}, 7 | "herd": {:hex, :herd, "0.4.3", "97469cf289c1e89a4f2b356da486ae5a354751f91c10cd3749af6aedebd9a775", [:mix], [{:libring, "~> 1.1", [hex: :libring, repo: "hexpm", optional: false]}, {:poolboy, "~> 1.5", [hex: :poolboy, repo: "hexpm", optional: false]}], "hexpm", "44bfd2c42a206431495d5103a77f52a992f4f1391a13459a9e8fd7b143cd99c9"}, 8 | "libring": {:hex, :libring, "1.7.0", "4f245d2f1476cd7ed8f03740f6431acba815401e40299208c7f5c640e1883bda", [:mix], [], "hexpm", "070e3593cb572e04f2c8470dd0c119bc1817a7a0a7f88229f43cf0345268ec42"}, 9 | "makeup": {:hex, :makeup, "1.2.1", "e90ac1c65589ef354378def3ba19d401e739ee7ee06fb47f94c687016e3713d1", [:mix], [{:nimble_parsec, "~> 1.4", [hex: :nimble_parsec, repo: "hexpm", optional: false]}], "hexpm", "d36484867b0bae0fea568d10131197a4c2e47056a6fbe84922bf6ba71c8d17ce"}, 10 | "makeup_elixir": {:hex, :makeup_elixir, "1.0.1", "e928a4f984e795e41e3abd27bfc09f51db16ab8ba1aebdba2b3a575437efafc2", [: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", "7284900d412a3e5cfd97fdaed4f5ed389b8f2b4cb49efc0eb3bd10e2febf9507"}, 11 | "makeup_erlang": {:hex, :makeup_erlang, "1.0.2", "03e1804074b3aa64d5fad7aa64601ed0fb395337b982d9bcf04029d68d51b6a7", [:mix], [{:makeup, "~> 1.0", [hex: :makeup, repo: "hexpm", optional: false]}], "hexpm", "af33ff7ef368d5893e4a267933e7744e46ce3cf1f61e2dccf53a111ed3aa3727"}, 12 | "memcachex": {:hex, :memcachex, "0.5.7", "00dc47d926eba11dfc1f2db606c37caff34d9fafd57f0adc10b49418931e77af", [:mix], [{:connection, "~> 1.0", [hex: :connection, repo: "hexpm", optional: false]}, {:poison, "~> 2.1 or ~> 3.0 or ~> 4.0 or ~> 5.0 or ~> 6.0", [hex: :poison, repo: "hexpm", optional: true]}, {:telemetry, "~> 0.4.0 or ~> 1.0", [hex: :telemetry, repo: "hexpm", optional: false]}], "hexpm", "a75e5b712c17bc25326c03c4c1b45fae39816b6fdacce3242724c7131f210ec1"}, 13 | "memcachir": {:hex, :memcachir, "3.3.1", "2099fd7d518b58172b417f932f9590e3fd3e11deb811c2a3a252605751af6aa7", [:mix], [{:elasticachex, "~> 1.1", [hex: :elasticachex, repo: "hexpm", optional: false]}, {:herd, "~> 0.4.3", [hex: :herd, repo: "hexpm", optional: false]}, {:memcachex, "~> 0.5", [hex: :memcachex, repo: "hexpm", optional: false]}], "hexpm", "526536e9585820894381a643aa3042f7f84e25e986e3c90f413fe0896f8bb513"}, 14 | "mime": {:hex, :mime, "2.0.6", "8f18486773d9b15f95f4f4f1e39b710045fa1de891fada4516559967276e4dc2", [:mix], [], "hexpm", "c9945363a6b26d747389aac3643f8e0e09d30499a138ad64fe8fd1d13d9b153e"}, 15 | "nimble_parsec": {:hex, :nimble_parsec, "1.4.2", "8efba0122db06df95bfaa78f791344a89352ba04baedd3849593bfce4d0dc1c6", [:mix], [], "hexpm", "4b21398942dda052b403bbe1da991ccd03a053668d147d53fb8c4e0efe09c973"}, 16 | "plug": {:hex, :plug, "1.16.1", "40c74619c12f82736d2214557dedec2e9762029b2438d6d175c5074c933edc9d", [:mix], [{:mime, "~> 1.0 or ~> 2.0", [hex: :mime, repo: "hexpm", optional: false]}, {:plug_crypto, "~> 1.1.1 or ~> 1.2 or ~> 2.0", [hex: :plug_crypto, repo: "hexpm", optional: false]}, {:telemetry, "~> 0.4.3 or ~> 1.0", [hex: :telemetry, repo: "hexpm", optional: false]}], "hexpm", "a13ff6b9006b03d7e33874945b2755253841b238c34071ed85b0e86057f8cddc"}, 17 | "plug_crypto": {:hex, :plug_crypto, "2.1.0", "f44309c2b06d249c27c8d3f65cfe08158ade08418cf540fd4f72d4d6863abb7b", [:mix], [], "hexpm", "131216a4b030b8f8ce0f26038bc4421ae60e4bb95c5cf5395e1421437824c4fa"}, 18 | "poolboy": {:hex, :poolboy, "1.5.2", "392b007a1693a64540cead79830443abf5762f5d30cf50bc95cb2c1aaafa006b", [:rebar3], [], "hexpm", "dad79704ce5440f3d5a3681c8590b9dc25d1a561e8f5a9c995281012860901e3"}, 19 | "socket": {:hex, :socket, "0.3.13", "98a2ab20ce17f95fb512c5cadddba32b57273e0d2dba2d2e5f976c5969d0c632", [:mix], [], "hexpm", "f82ea9833ef49dde272e6568ab8aac657a636acb4cf44a7de8a935acb8957c2e"}, 20 | "styler": {:hex, :styler, "1.4.0", "5944723d08afe4d38210b674d7e97dd1137a75968a85a633983cc308e86dc5f2", [:mix], [], "hexpm", "07de0e89c27490c8e469bb814d77ddaaa3283d7d8038501021d80a7705cf13e9"}, 21 | "telemetry": {:hex, :telemetry, "1.3.0", "fedebbae410d715cf8e7062c96a1ef32ec22e764197f70cda73d82778d61e7a2", [:rebar3], [], "hexpm", "7015fc8919dbe63764f4b4b87a95b7c0996bd539e0d499be6ec9d7f3875b79e6"}, 22 | } 23 | -------------------------------------------------------------------------------- /test/ex_limiter/plug_test.exs: -------------------------------------------------------------------------------- 1 | defmodule ExLimiter.PlugTest do 2 | use ExUnit.Case 3 | use Plug.Test 4 | alias ExLimiter.TestUtils 5 | 6 | describe "#call/2" do 7 | setup [:setup_limiter, :setup_conn] 8 | 9 | test "It will supply rate limiting headers if it passes", %{limiter: config, conn: conn} do 10 | conn = ExLimiter.Plug.call(conn, config) 11 | 12 | refute Enum.empty?(get_resp_header(conn, "x-ratelimit-limit")) 13 | refute Enum.empty?(get_resp_header(conn, "x-ratelimit-window")) 14 | refute Enum.empty?(get_resp_header(conn, "x-ratelimit-remaining")) 15 | end 16 | 17 | test "It will reject if the rate limit has been exceeded", %{limiter: config, conn: conn} do 18 | conn = 19 | %{conn | params: %{"count" => 11}} 20 | |> ExLimiter.Plug.call(config) 21 | 22 | assert conn.status == 429 23 | end 24 | 25 | test "it will respect scaling params", %{limiter: config, conn: conn} do 26 | config = %{config | limit: 1} 27 | conn = ExLimiter.Plug.call(conn, config) 28 | 29 | refute conn.status == 429 30 | 31 | conn = ExLimiter.Plug.call(conn, config) 32 | 33 | assert conn.status == 429 34 | end 35 | 36 | test "it will decorate a connection on ok", %{limiter: config, conn: conn} do 37 | config = %{config | decorate: &decorate/2} 38 | conn = ExLimiter.Plug.call(conn, config) 39 | 40 | refute conn.status == 429 41 | 42 | %{ex_limiter: %{bucket_name: bucket_name, bucket_version: bucket_version}} = conn.assigns 43 | 44 | assert String.ends_with?(bucket_name, "127.0.0.1") 45 | assert bucket_version == %{last: 0, value: 0} 46 | end 47 | 48 | test "it will decorate a connection on error", %{limiter: config, conn: conn} do 49 | config = %{config | decorate: &decorate/2, limit: 1} 50 | conn = ExLimiter.Plug.call(conn, config) 51 | conn = ExLimiter.Plug.call(conn, config) 52 | 53 | assert conn.status == 429 54 | 55 | %{ex_limiter: %{bucket_name: bucket_name}} = conn.assigns 56 | 57 | assert String.ends_with?(bucket_name, "127.0.0.1") 58 | end 59 | end 60 | 61 | defp decorate(conn, {:ok, %{key: bucket_name, version: bucket_version}}) do 62 | assign(conn, :ex_limiter, %{bucket_name: bucket_name, bucket_version: bucket_version}) 63 | end 64 | defp decorate(conn, {:rate_limited, bucket_name}) do 65 | assign(conn, :ex_limiter, %{bucket_name: bucket_name}) 66 | end 67 | 68 | defp setup_conn(_) do 69 | random = TestUtils.rand_string() 70 | conn = 71 | conn(:get, "/") 72 | |> merge_private(phoenix_controller: random, phoenix_action: random) 73 | 74 | [conn: conn] 75 | end 76 | 77 | defp setup_limiter(_) do 78 | [limiter: ExLimiter.Plug.Config.new(consumes: &consumes/1)] 79 | end 80 | 81 | defp consumes(%{params: %{"count" => count}}), do: count 82 | defp consumes(_), do: 1 83 | end 84 | -------------------------------------------------------------------------------- /test/ex_limiter/storage/pg2_shard_test.exs: -------------------------------------------------------------------------------- 1 | defmodule ExLimiter.Storage.PG2ShardTest do 2 | use ExUnit.Case, async: false 3 | alias ExLimiter.PG2Limiter 4 | alias ExLimiter.TestUtils 5 | 6 | describe "#consume" do 7 | test "it will rate limit" do 8 | bucket_name = bucket() 9 | {:ok, bucket} = PG2Limiter.consume(bucket_name, 1) 10 | 11 | assert bucket.key == bucket_name 12 | assert bucket.value >= 100 13 | 14 | {:ok, bucket} = PG2Limiter.consume(bucket_name, 5) 15 | 16 | assert bucket.value >= 500 17 | 18 | {:error, :rate_limited} = PG2Limiter.consume(bucket_name, 6) 19 | end 20 | end 21 | 22 | describe "#delete" do 23 | test "It will wipe a bucket" do 24 | bucket_name = bucket() 25 | 26 | {:ok, bucket} = PG2Limiter.consume(bucket_name, 5) 27 | 28 | assert bucket.value >= 500 29 | 30 | PG2Limiter.delete(bucket_name) 31 | 32 | {:ok, bucket} = PG2Limiter.consume(bucket_name, 1) 33 | 34 | assert bucket.value <= 500 35 | end 36 | end 37 | 38 | describe "#remaining" do 39 | test "It will properly deconvert the remaining capacity in a bucket" do 40 | bucket_name = bucket() 41 | 42 | {:ok, bucket} = PG2Limiter.consume(bucket_name, 5) 43 | 44 | assert PG2Limiter.remaining(bucket) == 5 45 | end 46 | end 47 | 48 | defp bucket() do 49 | "test_bucket_#{TestUtils.rand_string()}" 50 | end 51 | end 52 | -------------------------------------------------------------------------------- /test/ex_limiter/storage/pruner_test.exs: -------------------------------------------------------------------------------- 1 | defmodule ExLimiter.Storage.PG2Shard.PrunerTest do 2 | use ExUnit.Case, async: false 3 | alias ExLimiter.Storage.PG2Shard.Pruner 4 | 5 | @table_name :pruner_test 6 | describe "#remove" do 7 | test "it will remove stale bucket entries" do 8 | table = :ets.new(@table_name, [:set, :public, read_concurrency: true, write_concurrency: true]) 9 | 10 | :ets.insert(table, {"bucket", :os.system_time(:millisecond), %ExLimiter.Bucket{}}) 11 | assert Pruner.remove(table, 1) == 1 12 | assert :ets.lookup(table, "bucket") == [] 13 | end 14 | end 15 | end 16 | -------------------------------------------------------------------------------- /test/ex_limiter_test.exs: -------------------------------------------------------------------------------- 1 | defmodule ExLimiterTest do 2 | use ExUnit.Case 3 | alias ExLimiter.TestUtils 4 | doctest ExLimiter 5 | 6 | describe "#consume" do 7 | test "it will rate limit" do 8 | bucket_name = bucket() 9 | {:ok, bucket} = ExLimiter.consume(bucket_name, 1) 10 | 11 | assert bucket.key == bucket_name 12 | assert bucket.value >= 100 13 | 14 | {:ok, bucket} = ExLimiter.consume(bucket_name, 5) 15 | 16 | assert bucket.value >= 500 17 | 18 | {:error, :rate_limited} = ExLimiter.consume(bucket_name, 6) 19 | end 20 | 21 | test "it will rate limit for custom scale/limits" do 22 | bucket_name = bucket() 23 | args = [scale: 60_000, limit: 50] 24 | {:ok, bucket} = ExLimiter.consume(bucket_name, 1, args) 25 | 26 | assert bucket.key == bucket_name 27 | assert bucket.value >= 100 28 | 29 | for _ <- 0..10, do: {:ok, _} = ExLimiter.consume(bucket_name, 1, args) 30 | 31 | assert bucket.value >= 500 32 | 33 | {:error, :rate_limited} = ExLimiter.consume(bucket_name, 40, args) 34 | end 35 | end 36 | 37 | describe "#delete" do 38 | test "It will wipe a bucket" do 39 | bucket_name = bucket() 40 | 41 | {:ok, bucket} = ExLimiter.consume(bucket_name, 5) 42 | 43 | assert bucket.value >= 500 44 | 45 | ExLimiter.delete(bucket_name) 46 | 47 | {:ok, bucket} = ExLimiter.consume(bucket_name, 1) 48 | 49 | assert bucket.value <= 500 50 | end 51 | end 52 | 53 | describe "#remaining" do 54 | test "It will properly deconvert the remaining capacity in a bucket" do 55 | bucket_name = bucket() 56 | 57 | {:ok, bucket} = ExLimiter.consume(bucket_name, 5) 58 | 59 | assert ExLimiter.remaining(bucket) == 5 60 | end 61 | end 62 | 63 | defp bucket() do 64 | "test_bucket_#{TestUtils.rand_string()}" 65 | end 66 | end 67 | -------------------------------------------------------------------------------- /test/support/pg2_limiter.ex: -------------------------------------------------------------------------------- 1 | defmodule ExLimiter.PG2Limiter do 2 | use ExLimiter.Base, storage: ExLimiter.Storage.PG2Shard 3 | end 4 | -------------------------------------------------------------------------------- /test/support/test_utils.ex: -------------------------------------------------------------------------------- 1 | defmodule ExLimiter.TestUtils do 2 | def rand_string() do 3 | :crypto.strong_rand_bytes(8) 4 | |> Base.encode64() 5 | end 6 | end -------------------------------------------------------------------------------- /test/test_helper.exs: -------------------------------------------------------------------------------- 1 | ExUnit.start() 2 | {:ok, _pid} = ExLimiter.Storage.PG2Shard.Supervisor.start_link() 3 | --------------------------------------------------------------------------------