hammer

defmodule Hammer do
  @moduledoc """
  Documentation for Hammer module.

  This is the main API for the Hammer rate-limiter. This module assumes an appropriate
  backend has been configured.
  """

  alias Hammer.Utils

  @spec check_rate(id::String.t, scale_ms::integer, limit::integer)
        :: {:allow, count::integer}
         | {:deny,  limit::integer}
         | {:error, reason::any}
  @doc """
  Check if the action you wish to perform is within the bounds of the rate-limit.

  Args:
  - `id`: String name of the bucket. Usually the bucket name is comprised of some fixed prefix,
  with some dynamic string appended, such as an IP address or user id.
  - `scale_ms`: Integer indicating size of bucket in milliseconds
  - `limit`: Integer maximum count of actions within the bucket

  Returns either `{:allow,  count}`, `{:deny,   limit}` or `{:error,  reason}`

  Example:

      user_id = 42076
      case  check_rate("file_upload:\#{user_id}", 60_000, 5) do
        {:allow, _count} ->
          # do the file upload
        {:deny, _limit} ->
          # render an error page or something
      end
  """
  def check_rate(id, scale_ms, limit) do
    check_rate(:single, id, scale_ms, limit)
  end

  @spec check_rate(backend::atom,
                   id::String.t,
                   scale_ms::integer,
                   limit::integer)
        :: {:allow, count::integer}
         | {:deny,  limit::integer}
         | {:error, reason::any}
  @doc """
  Same as check_rate/3, but allows specifying a backend
  """
  def check_rate(backend, id, scale_ms, limit) do
    {stamp, key} = Utils.stamp_key(id, scale_ms)
    case call_backend(backend, :count_hit, [key, stamp]) do
      {:ok, count} ->
        if count > limit do
          {:deny, limit}
        else
          {:allow, count}
        end
      {:error, reason} ->
        {:error, reason}
    end
  end

  @spec inspect_bucket(id::String.t, scale_ms::integer, limit::integer)
        :: {:ok, {count::integer,
                  count_remaining::integer,
                  ms_to_next_bucket::integer,
                  created_at :: integer | nil,
                  updated_at :: integer | nil}}
         | {:error, reason::any}
  @doc """
  Inspect bucket to get count, count_remaining, ms_to_next_bucket, created_at, updated_at.
  This function is free of side-effects and should be called with the same arguments you
  would use for `check_rate` if you intended to increment and check the bucket counter.

  Arguments:

  - `id`: String name of the bucket. Usually the bucket name is comprised of some fixed prefix,
  with some dynamic string appended, such as an IP address or user id.
  - `scale_ms`: Integer indicating size of bucket in milliseconds
  - `limit`: Integer maximum count of actions within the bucket

  Returns either
  `{:ok, {count, count_remaining, ms_to_next_bucket, created_at, updated_at}`,
  or `{:error, reason}`.

  Example:

      inspect_bucket("file_upload:2042", 60_000, 5)
      {:ok, {1, 2499, 29381612, 1450281014468, 1450281014468}}

  """
  def inspect_bucket(id, scale_ms, limit) do
    inspect_bucket(:single, id, scale_ms, limit)
  end

  @spec inspect_bucket(backend::atom,
                       id::String.t,
                       scale_ms::integer,
                       limit::integer)
        :: {:ok, {count::integer,
                  count_remaining::integer,
                  ms_to_next_bucket::integer,
                  created_at :: integer | nil,
                  updated_at :: integer | nil}}
         | {:error, reason::any}
  @doc """
  Same as inspect_bucket/3, but allows specifying a backend
  """
  def inspect_bucket(backend, id, scale_ms, limit) do
    {stamp, key} = Utils.stamp_key(id, scale_ms)
    ms_to_next_bucket = (elem(key, 0) * scale_ms) + scale_ms - stamp
    case call_backend(backend, :get_bucket, [key]) do
      {:ok, nil} ->
        {:ok, {0, limit, ms_to_next_bucket, nil, nil}}
      {:ok, {_, count, created_at, updated_at}} ->
        count_remaining = if limit > count, do: limit - count, else: 0
        {:ok, {count, count_remaining, ms_to_next_bucket,
          created_at, updated_at}}
      {:error, reason} ->
        {:error, reason}
    end
  end

  @spec delete_buckets(id::String.t)
        :: {:ok, count::integer}
          | {:error, reason::any}
  @doc """
  Delete all buckets belonging to the provided id, including the current one.
  Effectively resets the rate-limit for the id.

  Arguments:

  - `id`: String name of the bucket

  Returns either `{:ok, count}` where count is the number of buckets deleted,
  or `{:error, reason}`.

  Example:

      user_id = 2406
      {:ok, _count} = delete_buckets("file_uploads:\#{user_id}")

  """
  def delete_buckets(id) do
    delete_buckets(:single, id)
  end

  @spec delete_buckets(backend::atom, id::String.t)
        :: {:ok, count::integer}
          | {:error, reason::any}
  @doc """
  Same as delete_buckets/1, but allows specifying a backend
  """
  def delete_buckets(backend, id) do
    call_backend(backend, :delete_buckets, [id])
  end

  @spec make_rate_checker(id_prefix::String.t,
                          scale_ms::integer,
                          limit::integer)
        :: ((id::String.t) -> {:allow, count::integer}
                            | {:deny,  limit::integer}
                            | {:error, reason::any})
  @doc """
  Make a rate-checker function, with the given `id` prefix, scale_ms and limit.

  Arguments:

  - `id_prefix`: String prefix to the `id`
  - `scale_ms`: Integer indicating size of bucket in milliseconds
  - `limit`: Integer maximum count of actions within the bucket

  Returns a function which accepts an `id` suffix, which will be combined with the `id_prefix`.
  Calling this returned function is equivalent to:
  `Hammer.check_rate("\#{id_prefix}\#{id}", scale_ms, limit)`

  Example:

      chat_rate_limiter = make_rate_checker("send_chat_message:", 60_000, 20)
      user_id = 203517
      case chat_rate_limiter.(user_id) do
        {:allow, _count} ->
          # allow chat message
        {:deny, _limit} ->
          # deny
      end
  """
  def make_rate_checker(id_prefix, scale_ms, limit) do
    make_rate_checker(:single, id_prefix, scale_ms, limit)
  end

  @spec make_rate_checker(backend::atom,
                          id_prefix::String.t,
                          scale_ms::integer,
                          limit::integer)
        :: ((id::String.t) -> {:allow, count::integer}
                            | {:deny,  limit::integer}
                            | {:error, reason::any})
  @doc """
  """
  def make_rate_checker(backend, id_prefix, scale_ms, limit) do
     fn (id) ->
       check_rate(backend, "#{id_prefix}#{id}", scale_ms, limit)
     end
   end

  defp call_backend(which, function, args) do
    backend = Utils.get_backend_module(which)
    apply(backend, function, args)
  end

end