hammer_backend_mnesia

defmodule Hammer.Mnesia do
  @moduledoc """
  This backend uses the `:mnesia`.

      defmodule MyApp.RateLimit do
        # the default table is the current's module name
        use Hammer, backend: Hammer.Mnesia, table: __MODULE__
      end

      # attempts to create a distributed in-memory table and run the clean up every 10 minutes
      MyApp.RateLimit.start_link(clean_period: :timer.minutes(10))

      {:allow, _count} = MyApp.RateLimit.hit(key, scale, limit)

  """
  use GenServer

  @type mnesia_option :: {:table, atom()}
  @type mnesia_options :: [mnesia_option()]

  # credo:disable-for-next-line Credo.Check.Refactor.CyclomaticComplexity
  defmacro __before_compile__(%{module: module} = _env) do
    hammer_opts = Module.get_attribute(module, :hammer_opts)
    table = Keyword.get(hammer_opts, :table, module)

    algorithm =
      case Keyword.get(hammer_opts, :algorithm) do
        nil ->
          Hammer.Mnesia.FixWindow

        :fix_window ->
          Hammer.Mnesia.FixWindow

        :leaky_bucket ->
          Hammer.Mnesia.LeakyBucket

        :token_bucket ->
          Hammer.Mnesia.TokenBucket

        module ->
          case Code.ensure_compiled(module) do
            {:module, _} ->
              module

            {:error, _} ->
              raise ArgumentError, """
                Hammer requires a valid backend to be specified. Must be one of: :fix_window, :leaky_bucket, :token_bucket.
                If none is specified, :fix_window is used.

              Example:

                use Hammer, backend: Hammer.Mnesia, algorithm: Hammer.Mnesia.FixWindow
              """
          end
      end

    Code.ensure_loaded!(algorithm)

    quote do
      @table unquote(table)
      @algorithm unquote(algorithm)

      def child_spec(opts) do
        %{
          id: __MODULE__,
          start: {__MODULE__, :start_link, [opts]},
          type: :worker
        }
      end

      def start_link(opts) do
        opts = Keyword.put(opts, :table, @table)
        opts = Keyword.put_new(opts, :algorithm, @algorithm)

        Hammer.Mnesia.start_link(opts)
      end

      def hit(key, scale, limit, increment \\ 1) do
        @algorithm.hit(@table, key, scale, limit, increment)
      end

      if function_exported?(@algorithm, :inc, 4) do
        def inc(key, scale, increment \\ 1) do
          @algorithm.inc(@table, key, scale, increment)
        end
      end

      if function_exported?(@algorithm, :set, 4) do
        def set(key, scale, count) do
          @algorithm.set(@table, key, scale, count)
        end
      end

      if function_exported?(@algorithm, :get, 2) do
        def get(key) do
          @algorithm.get(@table, key)
        end
      end

      if function_exported?(@algorithm, :get, 3) do
        def get(key, scale) do
          @algorithm.get(@table, key, scale)
        end
      end
    end
  end

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

  Accepts the following options:
    - `:clean_period` for how often to perform garbage collection
    - optional mnesia options from `:mnesia.create_table/2`
    - `:key_older_than` - How old a key can be before it is removed from the table (in milliseconds). Defaults to 10 minutes.
    - `:algorithm` - The rate limiting algorithm to use. Can be `:fixed_window`, `:token_bucket`, or `:leaky_bucket`. Defaults to `:fixed_window`.
    - optional `:debug`, `:spawn_opts`, and `:hibernate_after` `GenServer.options()`
  """
  @spec start_link(Hammer.Mnesia.mnesia_options()) ::
          {:ok, pid()} | :ignore | {:error, term()}
  def start_link(opts) do
    {gen_opts, opts} = Keyword.split(opts, [:debug, :spawn_opt, :hibernate_after])
    GenServer.start_link(__MODULE__, opts, gen_opts)
  end

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

  @impl GenServer
  def init(opts) do
    {:ok, opts, {:continue, :init}}
  end

  # TODO listen for cluster changes
  # TODO attempt unsplit
  @impl GenServer
  def handle_continue(:init, opts) do
    {clean_period, opts} = Keyword.pop!(opts, :clean_period)
    {algorithm, opts} = Keyword.pop!(opts, :algorithm)
    {key_older_than, opts} = Keyword.pop(opts, :key_older_than, :timer.minutes(10))
    {table, mnesia_opts} = Keyword.pop!(opts, :table)
    mnesia_opts = Keyword.merge(mnesia_opts, algorithm.mnesia_opts())

    case :mnesia.create_table(table, mnesia_opts) do
      {:atomic, :ok} -> :ok
      {:aborted, {:already_exists, _}} -> :ok
      {:aborted, reason} -> :erlang.error(reason)
    end

    schedule(clean_period)

    {:noreply,
     %{
       table: table,
       algorithm: algorithm,
       clean_period: clean_period,
       key_older_than: key_older_than
     }}
  end

  @impl GenServer
  def handle_info(:clean, state) do
    algorithm = state.algorithm
    algorithm.clean(state)
    schedule(state.clean_period)
    {:noreply, state}
  end

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