leaky

defmodule Leaky do
  @moduledoc """
  Implements a token-based rate limiter using the leaky bucket algorithm, ideal for controlling access rates to resources in Elixir applications.
  This implementation leverages ETS for optimized performance, making it suitable for high-load environments.

  ## Overview

  The leaky bucket algorithm helps smooth out bursty traffic by limiting the rate at which actions are taken.
  It's like a bucket with a hole: tokens, representing permissions to take an action, drip out at a constant rate.
  If the bucket is full, new actions must wait, ensuring the overall rate does not exceed the desired threshold.
  Read more about the algorithm here: https://en.wikipedia.org/wiki/Leaky_bucket

  ## Usage

  The rate limiter process can be customized with several options:
  - `bucket_name`: Unique identifier for the bucket.
  - `max_accumulated`: The maximum tokens the bucket can hold. Once full, new tokens will not accumulate.
  - `refill`: The number of tokens added to the bucket on each refill cycle.
  - `interval`: The time (in ms) between each refill cycle. Defaults to 1000 ms.
  - `name`: The GenServer process name. Defaults to `Leaky`.

  ```elixir
    configuration = [bucket_name: :user_requests, max_accumulated: 100, refill: 10, interval: 1000]
    {:ok, _pid} = Leaky.start_link(configuration)
  ```

  To attempt an action, checking if it's allowed under the current rate:
  ```elixir
    case Leaky.acquire(:user_requests, 1) do
      {:allow, tokens_left} -> "Action allowed."
      :deny ->"Action denied."
    end
  ```
  """
  use GenServer

  @type bucket :: integer | binary | tuple | atom

  defmodule State do
    @moduledoc false

    @type t :: %__MODULE__{
            refill: integer,
            initial: integer,
            interval: integer,
            max_accumulated: integer,
            name: GenServer.name()
          }

    defstruct [:bucket_name, :refill, :max_accumulated, :initial, :interval, :name]
  end

  @doc """
  The function evaluates whether a particular action can proceed without violating the configured rate limit of the bucket.
  It is designed to ensure that the frequency of actions does not exceed the predetermined limits, thus preventing system overload or abuse.

  ## Examples

  ```elixir
  .  case Leaky.acquire(:user_requests, 1) do
      {:allow, tokens_left} -> "Action allowed."
      :deny ->"Action denied."
    end
  ```
  """
  @spec acquire(bucket :: bucket(), cost :: integer, name :: GenServer.name()) ::
          {:allow, integer} | :deny
  def acquire(bucket, cost, name \\ __MODULE__) do
    GenServer.call(name, {:acquire, bucket, cost})
  end

  @doc """
  The function retrieves the current number of available tokens in the specified bucket,
  providing insight into the bucket's current state without altering it.
  It accepts the bucket identifier and an optional name parameter for the GenServer managing the rate limiter's state,
  returning either the number of tokens left or nil if the bucket does not exist.
  This function is useful for monitoring and debugging purposes, allowing developers to assess the rate limiter's status at any given moment.

  ## Example

  ```
   iex> Leaky.inspect(:user_requests)
   iex> 3
  ```
  """
  @spec inspect(bucket :: bucket(), name :: GenServer.name()) :: integer | nil
  def inspect(bucket, name \\ __MODULE__) do
    case GenServer.call(name, {:acquire, bucket, 0}) do
      {:allow, tokens_left} -> tokens_left
      :deny -> 0
    end
  end

  @doc """
  The function adjusts the number of tokens in a specified bucket by either adding (incrementing) or removing (decrementing) tokens,
  directly influencing the bucket's current capacity without performing a rate check.
  """
  @spec adjust_tokens(bucket :: bucket(), amount :: integer, name :: GenServer.name()) :: :ok
  def adjust_tokens(bucket, amount, name \\ __MODULE__) do
    GenServer.cast(name, {:increment_tokens_left, bucket, amount})
  end

  @doc """
  Updates the configuration of the rate limiter process. Changes are applied immediately, affecting the rate limiter's behavior.
  It is useful for dynamically adjusting the rate limiter's settings without restarting the process.

  Options which can be updated: `max_accumulated`, `interval`, and `refill`.


  ## Example

  ```
  iex> Leaky.update_configuration(max_accumulated: 10, refill: 2, interval: 5)
  :ok
  iex> {:allow, 8} == Leaky.acquire(:user_requests, 2)
  true
  ```

  """
  @spec update_configuration(opts :: Keyword.t(), name :: GenServer.name()) :: :ok
  def update_configuration(opts, name \\ __MODULE__) do
    GenServer.cast(name, {:update_configuration, opts})
  end

  def start_link(opts) do
    bucket_name = Keyword.fetch!(opts, :bucket_name)
    max_accumulated = Keyword.fetch!(opts, :max_accumulated)
    refill = Keyword.fetch!(opts, :refill)
    interval = Keyword.get(opts, :interval, 1_000)
    name = Keyword.get(opts, :name, __MODULE__)

    GenServer.start_link(
      __MODULE__,
      %State{
        bucket_name: bucket_name,
        max_accumulated: max_accumulated,
        interval: interval,
        refill: refill
      },
      name: name
    )
  end

  @impl GenServer
  def init(%State{} = state) do
    create_ets_table(state)

    {:ok, state}
  end

  defp create_ets_table(%State{} = state) do
    :ets.new(state.bucket_name, [
      :named_table,
      :set
    ])
  end

  @impl GenServer
  def handle_cast({:increment_tokens_left, bucket, amount}, %State{} = state) do
    now = :erlang.system_time(:milli_seconds)
    table = state.bucket_name

    case :ets.lookup(table, bucket) do
      [] ->
        :ets.insert(table, {bucket, state.max_accumulated + amount, now})

        :ok

      [{bucket, tokens, last_time}] ->
        accumulated_tokens = calculate_accumulated_tokens(tokens, now, last_time, state)
        tokens_left = accumulated_tokens + amount

        :ets.update_element(table, bucket, [{2, tokens_left}, {3, now}])
    end

    {:noreply, state}
  end

  def handle_cast({:update_configuration, opts}, %State{} = state) do
    max_accumulated = Keyword.get(opts, :max_accumulated, state.max_accumulated)
    interval = Keyword.get(opts, :interval, state.interval)
    refill = Keyword.get(opts, :refill, state.refill)

    {:noreply, %{state | max_accumulated: max_accumulated, interval: interval, refill: refill}}
  end

  @impl GenServer
  def handle_call({:acquire, bucket, cost}, _from, %State{} = state) do
    now = :erlang.system_time(:milli_seconds)
    table = state.bucket_name

    response =
      case :ets.lookup(table, bucket) do
        [] ->
          tokens_left = state.max_accumulated - cost

          :ets.insert(table, {bucket, tokens_left, now})

          {:allow, tokens_left}

        [{bucket, tokens, last_time}] ->
          accumulated_tokens = calculate_accumulated_tokens(tokens, now, last_time, state)
          tokens_left = accumulated_tokens - cost

          if tokens_left < 0 do
            :deny
          else
            :ets.update_element(table, bucket, [{2, tokens_left}, {3, now}])

            {:allow, tokens_left}
          end
      end

    {:reply, response, state}
  end

  defp calculate_accumulated_tokens(tokens, now, last_inserted_at, %State{} = state) do
    accumulated_tokens = tokens + (now - last_inserted_at) / state.interval * state.refill

    if accumulated_tokens > state.max_accumulated do
      state.max_accumulated
    else
      accumulated_tokens
    end
  end
end