lib/paraxial/rate_limit.ex

defmodule Paraxial.RateLimit do
  @moduledoc """
  Thin wrapper around `:ets.update_counter/4` and a
  clean-up process to act as a rate limiter.

  """

  use GenServer

  @doc """
  Starts the process that creates and cleans the ETS table.

  Accepts the following options:
    - `GenServer.options()`
    - `:table` for the ETS table name, defaults to `#{__MODULE__}`
    - `:clean_period` for how often to perform garbage collection
  """
  @spec start_link([GenServer.option() | {:table, atom} | {:clean_period, pos_integer}]) ::
          GenServer.on_start()
  def start_link(opts) do
    {gen_opts, opts} =
      Keyword.split(opts, [:debug, :name, :timeout, :spawn_opt, :hibernate_after])

    GenServer.start_link(__MODULE__, opts, gen_opts)
  end

  @doc """
  Checks the rate-limit for a key.
  """
  @spec check_rate(:ets.table(), key, scale, limit, increment) :: {:allow, count} | {:deny, limit}
        when key: term,
             scale: pos_integer,
             limit: pos_integer,
             increment: pos_integer,
             count: pos_integer
  def check_rate(table \\ __MODULE__, key, scale, limit, increment \\ 1) do
    bucket = div(now(), scale)
    full_key = {key, bucket}
    expires_at = (bucket + 1) * scale
    count = :ets.update_counter(table, full_key, increment, {full_key, 0, expires_at})
    if count <= limit, do: {:allow, count}, else: {:deny, limit}
  end

  @impl true
  def init(opts) do
    clean_period = Keyword.fetch!(opts, :clean_period)
    table = Keyword.get(opts, :table, __MODULE__)

    ^table =
      :ets.new(table, [
        :named_table,
        :set,
        :public,
        {:read_concurrency, true},
        {:write_concurrency, true},
        {:decentralized_counters, true}
      ])

    schedule(clean_period)
    {:ok, %{table: table, clean_period: clean_period}}
  end

  @impl true
  def handle_info(:clean, state) do
    clean(state.table)
    schedule(state.clean_period)
    {:noreply, state}
  end

  defp schedule(clean_period) do
    Process.send_after(self(), :clean, clean_period)
  end

  defp clean(table) do
    ms = [{{{:_, :_}, :_, :"$1"}, [], [{:<, :"$1", {:const, now()}}]}]
    :ets.select_delete(table, ms)
  end

  @compile inline: [now: 0]
  defp now do
    System.system_time(:millisecond)
  end
end