eb444cb012
- hackney: 4.5.2 → 4.6.0 (constraint ~> 4.4 → ~> 4.6) - stream_data: 1.3.0 → 1.4.0 - mint: 1.9.1 → 1.9.2 (fixes CVE-2026-58229, CVE-2026-59246) - oban_pro: 1.7.6 → 1.7.8 (vendor sync from Oban Pro private repo) - oban_web: lock synced to 2.12.6 (vendor already had it)
393 lines
12 KiB
Elixir
393 lines
12 KiB
Elixir
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
|
|
|
|
{producer, producer_capacity(callback, rl, partition, curr_time)}
|
|
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 producer_capacity(callback, rate_limit, partition, curr_time) do
|
|
%{allowed: allowed, period: period, window_time: window_time, windows: windows} = rate_limit
|
|
|
|
if is_nil(rate_limit.partition) do
|
|
callback.demand(windows, allowed, period, curr_time, window_time)
|
|
else
|
|
windows
|
|
|> callback.partition_demands(allowed, period, curr_time, window_time, [partition])
|
|
|> Map.get(partition, 0)
|
|
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
|