prop/vendor/oban_pro/lib/oban/pro/rate_limit.ex

defmodule Oban.Pro.RateLimit do
  @moduledoc """
  Programmatic API for inspecting and manipulating rate limit state.

  Rate limits are normally managed automatically by the Smart engine as jobs execute. This module
  provides functions to interact with rate limits outside of normal job execution: checking
  available capacity, manually consuming tokens, and resetting state.

  All functions in this module require the [Smart engine](`Oban.Pro.Engines.Smart`) and a queue
  configured with `rate_limit` options.

  ## Checking Available Capacity

  Use `available/2` to check how much capacity remains before performing operations:

  ```elixir
  case Oban.Pro.RateLimit.available(:some_api) do
    {:ok, capacity} when capacity >= 10 ->
      make_api_calls(10)

    {:ok, capacity} when capacity > 0 ->
      make_api_calls(capacity)

    {:ok, 0} ->
      {:snooze, 60}
  end
  ```

  For partitioned rate limits, check a specific partition with a job's computed `partition_key`:

  ```elixir
  Oban.Pro.RateLimit.available(:some_api, partition: job.meta["partition_key"])
  ```

  ## Manual Consumption

  Use `consume/3` when you need to track rate-limited operations that happen outside of job
  execution. This ensures the rate limit reflects all usage, not just job execution. Consumption
  can even track operations entirely outside of Oban:

  ```elixir
  def handle_webhook(conn, params) do
    :ok = Oban.Pro.RateLimit.consume(:api_calls, 1)
    result = ExternalAPI.process(params)
    json(conn, result)
  end
  ```

  ## Resetting State

  Use `reset/2` to clear all rate limit tracking. This is primarily for recovery after
  configuration changes when limits should be cleared:

  ```elixir
  :ok = Oban.Pro.RateLimit.reset(:api_calls)
  ```

  ## See Also

  * `Oban.Pro.Engines.Smart` — Rate limit configuration and algorithms
  * `Oban.Pro.Worker` — Weighted jobs with the `weight/1` callback
  """
  @moduledoc since: "1.7.0"

  import Ecto.Query
  import DateTime, only: [utc_now: 0]

  alias Oban.Pro.Limiters.Rate.Algorithm
  alias Oban.Pro.Producer
  alias Oban.Repo

  @type queue :: atom() | String.t()
  @type option :: {:oban, Oban.name()} | {:partition, String.t()}
  @type wait_option :: {:timeout, timeout()} | {:interval, pos_integer()}

  @doc """
  Manually consume rate limit capacity for a queue.

  Consumption is applied to the producer with the most available capacity. If the requested
  amount exceeds a single producer's capacity, consumption is spread across multiple producers.

  ## Options

    * `:oban` - The Oban instance name. Defaults to `Oban`.
    * `:partition` - The partition key to consume from. Defaults to `"*"` (the global partition).
    * `:require_full` - When `true`, returns `{:error, :insufficient_capacity}` if the full
      amount can't be consumed. Defaults to `false`, which consumes as much as available.

  ## Examples

  Consume 5 units from the default partition

      Oban.Pro.RateLimit.consume(:my_queue, 5)

  Consume from a specific partition

      Oban.Pro.RateLimit.consume(:my_queue, 3, partition: job.meta["partition_key"])

  Use a named Oban instance

      Oban.Pro.RateLimit.consume(:my_queue, 1, oban: MyApp.Oban)
  """
  @spec consume(queue(), pos_integer(), [option()]) ::
          :ok | {:error, :insufficient_capacity | :no_rate_limit | :queue_not_found}
  def consume(queue, amount, opts \\ []) when is_integer(amount) and amount > 0 do
    conf = fetch_conf(opts)
    partition = Keyword.get(opts, :partition, "*")
    require_full = Keyword.get(opts, :require_full, false)

    func = fn ->
      with {:ok, producers} <- fetch_producers(conf, queue, take_lock: true),
           {:ok, rate_limit} <- fetch_rate_limit(hd(producers)) do
        curr_time = unix_now()
        callback = Algorithm.callback(rate_limit)

        sorted =
          producers
          |> Enum.map(fn producer ->
            rl = producer.meta.rate_limit

            capacity =
              callback.demand(rl.windows, rl.allowed, rl.period, curr_time, rl.window_time)

            {producer, capacity}
          end)
          |> Enum.sort_by(&elem(&1, 1), :desc)

        total_capacity = Enum.reduce(sorted, 0, fn {_, cap}, acc -> acc + cap end)

        if require_full and total_capacity < amount do
          {:error, :insufficient_capacity}
        else
          consume(sorted, conf, partition, amount, curr_time, callback)
          :ok
        end
      end
    end

    with {:ok, result} <- Repo.transaction(conf, func), do: result
  end

  @doc """
  Check the available rate limit capacity for a queue.

  Returns the total available capacity across all producers for the queue, calculated by merging
  window states from all active producers.

  ## Options

    * `:oban` - The Oban instance name. Defaults to `Oban`.
    * `:partition` - The partition key to check. Defaults to `"*"` (the global partition).

  ## Examples

  Check capacity for a queue:

      {:ok, capacity} = Oban.Pro.RateLimit.available(:my_queue)

  Check capacity for a specific partition:

      {:ok, capacity} = Oban.Pro.RateLimit.available(:my_queue, partition: job.meta["partition_key"])
  """
  @spec available(queue(), [option()]) ::
          {:ok, non_neg_integer()} | {:error, :no_rate_limit | :queue_not_found}
  def available(queue, opts \\ []) do
    conf = fetch_conf(opts)
    partition = Keyword.get(opts, :partition, "*")

    with {:ok, producers} <- fetch_producers(conf, queue),
         {:ok, rate_limit} <- fetch_rate_limit(hd(producers)) do
      curr_time = unix_now()
      callback = Algorithm.callback(rate_limit)

      windows =
        producers
        |> Enum.map(& &1.meta.rate_limit)
        |> Enum.filter(&(&1.window_time >= curr_time - &1.period))
        |> Enum.reduce(%{}, fn rl, acc -> callback.merge(rl.windows, acc) end)

      demand =
        if is_nil(rate_limit.partition) do
          callback.demand(
            windows,
            rate_limit.allowed,
            rate_limit.period,
            curr_time,
            rate_limit.window_time
          )
        else
          demands =
            callback.partition_demands(
              windows,
              rate_limit.allowed,
              rate_limit.period,
              curr_time,
              rate_limit.window_time,
              [partition]
            )

          Map.get(demands, partition, 0)
        end

      {:ok, demand}
    end
  end

  @doc """
  Execute a function after atomically reserving rate limit capacity.

  This function waits for capacity to become available, atomically consumes the requested amount,
  then executes the provided function. This prevents race conditions where multiple callers might
  consume the same quota.

  ## Options

    * `:oban` - The Oban instance name. Defaults to `Oban`.
    * `:partition` - The partition key to check. Defaults to `"*"` (the global partition).
    * `:timeout` - Maximum time to wait in milliseconds. Defaults to `5_000` (5 seconds).
    * `:interval` - Polling interval in milliseconds. Defaults to `100`.

  ## Examples

  Execute a function after reserving 5 units of capacity:

      {:ok, result} = Oban.Pro.RateLimit.with_quota(:my_queue, 5, fn ->
        ExternalAPI.batch_request(items)
      end)

  Handle timeout when capacity isn't available:

      case Oban.Pro.RateLimit.with_quota(:my_queue, 5, &make_api_calls/0, timeout: 10_000) do
        {:ok, result} -> handle_result(result)
        {:error, :timeout} -> handle_timeout()
      end

  Reserve capacity on a specific partition:

      {:ok, result} = Oban.Pro.RateLimit.with_quota(:my_queue, 1, fun, partition: "user_123")
  """
  @spec with_quota(queue(), pos_integer(), (-> result), [option() | wait_option()]) ::
          {:ok, result} | {:error, :timeout | :no_rate_limit | :queue_not_found}
        when result: term()
  def with_quota(queue, amount, fun, opts \\ []) when is_integer(amount) and amount > 0 do
    {timeout, opts} = Keyword.pop(opts, :timeout, 5_000)
    {interval, opts} = Keyword.pop(opts, :interval, 100)

    deadline = System.monotonic_time(:millisecond) + timeout

    with_quota_loop(queue, amount, fun, opts, deadline, interval)
  end

  defp with_quota_loop(queue, amount, fun, opts, deadline, interval) do
    case consume(queue, amount, Keyword.put(opts, :require_full, true)) do
      :ok ->
        {:ok, fun.()}

      {:error, :insufficient_capacity} ->
        remaining = deadline - System.monotonic_time(:millisecond)

        if remaining <= 0 do
          {:error, :timeout}
        else
          Process.sleep(min(interval, remaining))

          with_quota_loop(queue, amount, fun, opts, deadline, interval)
        end

      {:error, _reason} = error ->
        error
    end
  end

  @doc """
  Reset the rate limit state for a queue.

  This clears all window data and resets the window time for all producers on the queue. Tracking
  is cleared across all partitions for partitioned queues.

  ## Options

    * `:oban` - The Oban instance name. Defaults to `Oban`.

  ## Examples

  Reset the rate limit for a queue:

      :ok = Oban.Pro.RateLimit.reset(:my_queue)

  Reset the rate limit for a queue:

      :ok = Oban.Pro.RateLimit.reset(:my_queue, oban: MyApp.Oban)
  """
  @spec reset(queue(), [option()]) :: :ok | {:error, :no_rate_limit | :queue_not_found}
  def reset(queue, opts \\ []) do
    conf = fetch_conf(opts)

    func = fn ->
      with {:ok, producers} <- fetch_producers(conf, queue, take_lock: true),
           {:ok, _rate_lim} <- fetch_rate_limit(hd(producers)) do
        curr_time = unix_now()

        Enum.each(producers, fn producer ->
          rate_limit = producer.meta.rate_limit
          updated = %{rate_limit | windows: %{}, window_time: curr_time}

          update_rate_limit(conf, producer, updated)
        end)
      end
    end

    with {:ok, result} <- Repo.transaction(conf, func), do: result
  end

  # Private Helpers

  defp fetch_conf(opts) do
    opts
    |> Keyword.get(:oban, Oban)
    |> Oban.config()
  end

  defp fetch_producers(conf, queue, opts \\ []) do
    query =
      if opts[:take_lock] do
        Producer
        |> where(queue: ^to_string(queue))
        |> lock("FOR UPDATE")
      else
        where(Producer, queue: ^to_string(queue))
      end

    case Repo.all(conf, query) do
      [_ | _] = producers -> {:ok, producers}
      _ -> {:error, :queue_not_found}
    end
  end

  defp fetch_rate_limit(producer) do
    case producer.meta.rate_limit do
      %{} = rate_limit -> {:ok, rate_limit}
      nil -> {:error, :no_rate_limit}
    end
  end

  defp consume([], _, _, _, _, _), do: :ok
  defp consume(_, _, _, 0, _, _), do: :ok

  defp consume([{producer, capacity} | rest], conf, partition, remaining, curr_time, callback) do
    to_consume = min(capacity, remaining)

    if to_consume > 0 do
      rate_limit = producer.meta.rate_limit
      new_counts = %{partition => to_consume}

      {all_windows, next_time} =
        callback.track(
          rate_limit.windows,
          new_counts,
          rate_limit.period,
          rate_limit.window_time,
          curr_time,
          rate_limit.allowed
        )

      updated = %{rate_limit | windows: all_windows, window_time: next_time}

      update_rate_limit(conf, producer, updated)
    end

    consume(rest, conf, partition, remaining - to_consume, curr_time, callback)
  end

  defp update_rate_limit(conf, producer, rate_limit) do
    meta = %{producer.meta | rate_limit: rate_limit}

    Producer
    |> where(uuid: ^producer.uuid)
    |> then(&Repo.update_all(conf, &1, set: [meta: meta, updated_at: utc_now()]))
  end

  defp unix_now do
    DateTime.to_unix(utc_now(), :second)
  end
end