Skip to content

Instantly share code, notes, and snippets.

@schrockwell
Last active July 16, 2019 12:16
Show Gist options
  • Save schrockwell/feaf6fc4826ccea50991ea04b9bebdd7 to your computer and use it in GitHub Desktop.
Save schrockwell/feaf6fc4826ccea50991ea04b9bebdd7 to your computer and use it in GitHub Desktop.
Simple Elixir rate limiter with Plug and ETS
defmodule RHRLiveWeb.RateLimiter.Plug do
@behaviour Plug
import Plug.Conn
@doc """
Example configuration:
plug RHRLiveWeb.RateLimiter.Plug, event: :graphql, count: 10, per: 5_000
"""
def init(opts) do
%{
event: Keyword.fetch!(opts, :event),
count: Keyword.fetch!(opts, :count),
per: Keyword.fetch!(opts, :per),
status: Keyword.get(opts, :status, :too_many_requests),
body: Keyword.get(opts, :body, "Throttled")
}
end
def call(conn, opts) do
if RHRLiveWeb.RateLimiter.throttle?(conn, opts.event, opts.count, opts.per) do
conn
|> send_resp(opts.status, opts.body)
|> halt()
else
conn
end
end
end
defmodule RHRLiveWeb.RateLimiter do
@moduledoc """
Rate-limits Plug requests for a particular number of event hits over a given interval. Utilizes
the first three octets of the `remote_ip` and the event as the uniqueness key.
`initialize/0` must be called before it can be used, to create the ETS table.
Utilize `RHRLiveWeb.RateLimiter.Server` to set up automatic ETS cleanup to conserve memory.
"""
@table __MODULE__
@doc """
Required before this module can be used. Creates a public ETS table for persistence.
"""
@spec initialize :: no_return
def initialize do
:ets.new(@table, [:bag, :public, :named_table])
end
@doc """
Records a "hit" and determines if it should be throttled according to the specified hits per
interval. Purges all hits older than the interval.
"""
@spec throttle?(Plug.Conn.t(), term, non_neg_integer, non_neg_integer) :: boolean
def throttle?(conn, event, max_count, interval_ms) do
now_ms = :erlang.monotonic_time(:millisecond)
threshold_ms = now_ms - interval_ms
key = client_event_key(conn, event)
insert(key, now_ms)
delete_before(key, threshold_ms)
count = @table |> :ets.lookup(key) |> length()
count > max_count
end
@doc """
Deletes all hits more than `ms` milliseconds in the past.
Returns the number of hits deleted.
"""
@spec purge(non_neg_integer) :: non_neg_integer
def purge(ms \\ 60_000) do
threshold_ms = :erlang.monotonic_time(:millisecond) - ms
delete_before(:_, threshold_ms)
end
# ETS object typespec:
# {{{ip1 :: non_neg_integer, ip2 ::non_neg_integer, ip3 ::non_neg_integer}, event :: term}, timestamp_ms ::integer}
defp insert(key, timestamp) do
:ets.insert(@table, {key, timestamp})
end
defp delete_before(key_spec, timestamp) do
# :ets.fun2ms(fn
# {"key", timestamp} -> timestamp < "timestamp"
# _ -> false
# end)
match_spec = [{{key_spec, :"$1"}, [], [{:<, :"$1", timestamp}]}, {:_, [], [false]}]
:ets.select_delete(@table, match_spec)
end
# Utilize the remote_ip as the client identifier
defp client_event_key(%Plug.Conn{remote_ip: remote_ip}, event) do
{remote_ip, event}
end
end
defmodule RHRLiveWeb.RateLimiter.Server do
@moduledoc """
Periodically purges the RateLimiter ETS table to conserve memory
"""
use GenServer
@purge_ms 60_000
def start_link(_) do
GenServer.start_link(__MODULE__, [], name: __MODULE__)
end
def init(_) do
RHRLiveWeb.RateLimiter.initialize()
Process.send_after(self(), :purge, @purge_ms)
{:ok, []}
end
def handle_info(:purge, _) do
RHRLiveWeb.RateLimiter.purge(@purge_ms)
Process.send_after(self(), :purge, @purge_ms)
{:noreply, []}
end
end
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment