towerops/lib/towerops/rate_limit.ex

defmodule Towerops.RateLimit do
  @moduledoc """
  In-memory fix-window rate limiter backed by ETS.

  Each call to `hit/3` (or `hit/4`) increments a counter for the bucket
  `{key, current_window}` where the window index is `div(now_ms, scale_ms)`.
  Counters are stored in a single ETS table owned by this GenServer and are
  swept periodically by the supervisor process.

  This module is the in-tree replacement for the previous Hammer + ETS
  backend dependency. It exposes only what the rest of the codebase actually
  uses: `start_link/1`, `child_spec/1`, `hit/3`/`hit/4`, `get/3`, and
  `reset/3`.

  ## Usage

      children = [
        {Towerops.RateLimit, clean_period: :timer.minutes(1)}
      ]

      # 5 requests per second per key
      case Towerops.RateLimit.hit("user:123", 1_000, 5) do
        {:allow, _count} -> :ok
        {:deny, retry_after_ms} -> {:error, retry_after_ms}
      end
  """

  use GenServer

  @type key :: term()
  @type scale_ms :: pos_integer()
  @type limit :: pos_integer()
  @type increment :: non_neg_integer()
  @type count :: non_neg_integer()
  @type retry_after_ms :: non_neg_integer()

  @default_table __MODULE__
  @default_clean_period 60_000

  # ---------------------------------------------------------------------------
  # Supervisor API
  # ---------------------------------------------------------------------------

  @doc """
  Returns the child spec used to start this rate limiter under a supervisor.

  Accepts the same options as `start_link/1`. The `:id` option overrides the
  default supervisor id.
  """
  @spec child_spec(keyword()) :: Supervisor.child_spec()
  def child_spec(opts) do
    {id, opts} = Keyword.pop(opts, :id, __MODULE__)

    %{
      id: id,
      start: {__MODULE__, :start_link, [opts]},
      type: :worker
    }
  end

  @doc """
  Starts the rate limiter process and creates its backing ETS table.

  ## Options

    * `:name` - the GenServer name. Defaults to `Towerops.RateLimit`.
    * `:table` - the ETS table name. Defaults to `Towerops.RateLimit`.
    * `:clean_period` - milliseconds between cleanup ticks. Defaults to
      `60_000` (one minute).

  """
  @spec start_link(keyword()) :: GenServer.on_start()
  def start_link(opts) do
    name = Keyword.get(opts, :name, __MODULE__)
    GenServer.start_link(__MODULE__, opts, name: name)
  end

  # ---------------------------------------------------------------------------
  # Public API
  # ---------------------------------------------------------------------------

  @doc """
  Records a hit against `key`, returning `{:allow, count}` if the bucket is
  still under `limit`, or `{:deny, retry_after_ms}` if it has been exhausted.

  When `table` is omitted the default table (`Towerops.RateLimit`) is used.
  """
  @spec hit(key(), scale_ms(), limit()) :: {:allow, count()} | {:deny, retry_after_ms()}
  def hit(key, scale, limit), do: hit(@default_table, key, scale, limit, 1)

  @spec hit(atom() | key(), key() | scale_ms(), scale_ms() | limit(), limit() | increment()) ::
          {:allow, count()} | {:deny, retry_after_ms()}
  def hit(table, key, scale, limit) when is_atom(table) and is_integer(scale) do
    hit(table, key, scale, limit, 1)
  end

  def hit(key, scale, limit, increment) when is_integer(scale) and is_integer(increment) do
    hit(@default_table, key, scale, limit, increment)
  end

  @doc """
  Records a hit with a custom `increment`. See `hit/3` for the return shape.
  """
  @spec hit(atom(), key(), scale_ms(), limit(), increment()) ::
          {:allow, count()} | {:deny, retry_after_ms()}
  def hit(table, key, scale, limit, increment)
      when is_atom(table) and is_integer(scale) and scale > 0 and is_integer(limit) and limit > 0 and
             is_integer(increment) and increment >= 0 do
    GenServer.call(table, {:hit, key, scale, limit, increment})
  end

  @doc """
  Returns the current bucket count for `key` in the active window, or `0` if
  the bucket has not yet been created.
  """
  @spec get(atom(), key(), scale_ms()) :: count()
  def get(table \\ @default_table, key, scale) when is_integer(scale) and scale > 0 do
    GenServer.call(table, {:get, key, scale})
  end

  @doc """
  Removes the bucket entry for `key` in the current window.
  """
  @spec reset(atom(), key(), scale_ms()) :: :ok
  def reset(table \\ @default_table, key, scale) when is_integer(scale) and scale > 0 do
    GenServer.call(table, {:reset, key, scale})
  end

  # ---------------------------------------------------------------------------
  # GenServer callbacks
  # ---------------------------------------------------------------------------

  @impl GenServer
  def init(opts) do
    table = Keyword.get(opts, :table, @default_table)
    clean_period = Keyword.get(opts, :clean_period, @default_clean_period)

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

    _ = schedule_clean(clean_period)
    {:ok, %{table: table, clean_period: clean_period}}
  end

  @impl GenServer
  def handle_call({:hit, key, scale, limit, increment}, _from, state) do
    now = now_ms()
    window = div(now, scale)
    expires_at = (window + 1) * scale
    full_key = {key, window}

    count =
      :ets.update_counter(state.table, full_key, increment, {full_key, 0, expires_at})

    result =
      if count <= limit do
        {:allow, count}
      else
        {:deny, expires_at - now}
      end

    {:reply, result, state}
  end

  @impl GenServer
  def handle_call({:get, key, scale}, _from, state) do
    window = div(now_ms(), scale)

    count =
      case :ets.lookup(state.table, {key, window}) do
        [{_full_key, c, _expires_at}] -> c
        [] -> 0
      end

    {:reply, count, state}
  end

  @impl GenServer
  def handle_call({:reset, key, scale}, _from, state) do
    window = div(now_ms(), scale)
    :ets.delete(state.table, {key, window})
    {:reply, :ok, state}
  end

  @impl GenServer
  def handle_info(:clean, state) do
    sweep(state.table)
    schedule_clean(state.clean_period)
    {:noreply, state}
  end

  # ---------------------------------------------------------------------------
  # Internals
  # ---------------------------------------------------------------------------

  defp sweep(table) do
    match_spec = [{{:_, :_, :"$1"}, [{:<, :"$1", {:const, now_ms()}}], [true]}]
    :ets.select_delete(table, match_spec)
  end

  defp schedule_clean(period), do: Process.send_after(self(), :clean, period)

  defp now_ms, do: System.system_time(:millisecond)
end