Graham McIntire
e99bf06eb4
Previous commit (3c988a5f) switched oban_pro to the licensed hex repo at deploy time. Reverting that — keep all three Pro packages vendored so prod image builds don't depend on oban.pro reachability or auth on every CI run. oban_pro 1.7.0 dropped into vendor/oban_pro via `mix hex.package fetch oban_pro 1.7.0 --repo=oban --unpack`. mix.exs goes back to `path: "vendor/oban_pro"` (matching oban_met / oban_web, which were already vendored — and stay vendored since the licensed repo only has older versions of those: 0.1.11 / 2.10.6 vs the 1.1.0 / 2.12.1 we vendor). Schema migration from3c988a5fstays — 1.7.0 tables/indexes are already applied. No code changes.
376 lines
11 KiB
Elixir
376 lines
11 KiB
Elixir
require Protocol
|
|
|
|
defmodule Oban.Pro.Producer do
|
|
@moduledoc false
|
|
|
|
use Ecto.Schema
|
|
|
|
import Ecto.Changeset
|
|
|
|
alias Ecto.Changeset
|
|
alias Oban.Pro.Utils
|
|
|
|
defmodule ListOrMap do
|
|
@moduledoc false
|
|
|
|
# This is a custom type to allow a graceful transition from the legacy rate limit windows
|
|
# structure to the new, flattened map.
|
|
|
|
@behaviour Ecto.Type
|
|
|
|
@impl true
|
|
def type, do: :any
|
|
|
|
@impl true
|
|
def cast(data) when is_map(data), do: {:ok, data}
|
|
def cast([data | _]), do: {:ok, data}
|
|
def cast(_data), do: :error
|
|
|
|
@impl true
|
|
def load(data) when is_map(data), do: {:ok, data}
|
|
def load([data | _]), do: {:ok, data}
|
|
def load([]), do: {:ok, %{}}
|
|
def load(_data), do: :error
|
|
|
|
@impl true
|
|
def dump(data) when is_map(data), do: {:ok, data}
|
|
def dump([data | _]), do: {:ok, data}
|
|
def dump([]), do: {:ok, %{}}
|
|
def dump(_data), do: :error
|
|
|
|
@impl true
|
|
def equal?(a, b), do: a == b
|
|
|
|
@impl true
|
|
def embed_as(_data), do: :dump
|
|
end
|
|
|
|
@type t :: %__MODULE__{
|
|
uuid: Ecto.UUID.t(),
|
|
name: binary(),
|
|
node: binary(),
|
|
queue: binary(),
|
|
meta: map(),
|
|
started_at: DateTime.t(),
|
|
updated_at: DateTime.t()
|
|
}
|
|
|
|
@type meta :: %__MODULE__.Meta{local_limit: pos_integer(), paused: boolean()}
|
|
|
|
@primary_key {:uuid, :binary_id, autogenerate: false}
|
|
schema "oban_producers" do
|
|
field :name, :string
|
|
field :node, :string
|
|
field :queue, :string
|
|
field :started_at, :utc_datetime_usec
|
|
field :updated_at, :utc_datetime_usec
|
|
|
|
field :ack_async, :boolean, virtual: true, default: true
|
|
field :ack_tab, :any, virtual: true
|
|
field :refresh_interval, :integer, virtual: true, default: :timer.seconds(30)
|
|
field :xact_delay, :integer, default: :timer.seconds(1), virtual: true
|
|
field :xact_retry, :integer, default: 5, virtual: true
|
|
field :xact_timeout, :integer, default: :timer.seconds(30), virtual: true
|
|
|
|
embeds_one :meta, Meta, on_replace: :update, primary_key: false do
|
|
@moduledoc false
|
|
|
|
field :local_limit, :integer, default: 10
|
|
field :paused, :boolean, default: false
|
|
field :shutdown_started_at, :utc_datetime_usec
|
|
|
|
embeds_one :global_limit, GlobalLimit, on_replace: :update, primary_key: false do
|
|
@moduledoc false
|
|
|
|
field :allowed, :integer
|
|
field :burst, :boolean, default: false
|
|
field :tracked, ListOrMap, default: %{}
|
|
|
|
embeds_one :partition, Partition, on_replace: :update, primary_key: false do
|
|
@moduledoc false
|
|
|
|
field :fields, {:array, Ecto.Enum}, values: [:args, :meta, :worker]
|
|
field :keys, {:array, :string}
|
|
end
|
|
end
|
|
|
|
embeds_one :rate_limit, RateLimit, on_replace: :update, primary_key: false do
|
|
@moduledoc false
|
|
|
|
field :algorithm, Ecto.Enum,
|
|
values: ~w(sliding_window fixed_window token_bucket)a,
|
|
default: :sliding_window
|
|
|
|
field :allowed, :integer
|
|
field :period, :integer
|
|
field :window_time, :integer, skip_default_validation: true
|
|
field :windows, ListOrMap, default: %{}
|
|
|
|
embeds_one :partition, Partition, on_replace: :update, primary_key: false do
|
|
@moduledoc false
|
|
|
|
field :fields, {:array, Ecto.Enum}, values: [:args, :meta, :worker]
|
|
field :keys, {:array, :string}
|
|
end
|
|
end
|
|
end
|
|
end
|
|
|
|
@spec new(map() | Keyword.t()) :: Changeset.t(t())
|
|
def new(params) when is_list(params) or is_map(params) do
|
|
params =
|
|
params
|
|
|> Map.new()
|
|
|> Map.put_new_lazy(:uuid, &Oban.Pro.UUIDv7.generate/0)
|
|
|> Map.update!(:name, &to_clean_string/1)
|
|
|> Map.update!(:node, &to_clean_string/1)
|
|
|> Map.update!(:queue, &to_clean_string/1)
|
|
|
|
%__MODULE__{}
|
|
|> cast(params, ~w(name node queue refresh_interval started_at updated_at uuid)a)
|
|
|> cast_embed(:meta, required: true, with: &meta_changeset/2)
|
|
|> validate_required(~w(name node queue)a)
|
|
|> validate_length(:name, min: 1)
|
|
|> validate_length(:node, min: 1)
|
|
|> validate_length(:queue, min: 1)
|
|
|> Utils.enforce_keys(params, __MODULE__)
|
|
end
|
|
|
|
@spec update_meta(t(), atom(), term()) :: Changeset.t()
|
|
def update_meta(schema, key, value) do
|
|
update_meta(schema, %{key => value})
|
|
end
|
|
|
|
@spec update_meta(t(), map()) :: Changeset.t()
|
|
def update_meta(schema, changes) do
|
|
meta =
|
|
schema.meta
|
|
|> meta_changeset(changes)
|
|
|> apply_changes()
|
|
|> Ecto.embedded_dump(:json)
|
|
|
|
schema
|
|
|> cast(%{meta: meta}, [])
|
|
|> cast_embed(:meta, with: &meta_changeset/2)
|
|
end
|
|
|
|
@allowed ~w(local_limit paused shutdown_started_at)a
|
|
|
|
@doc false
|
|
def meta_changeset(schema, params, allowed \\ @allowed) do
|
|
params =
|
|
params
|
|
|> Map.new()
|
|
|> Map.delete(:limit)
|
|
|> Map.put_new_lazy(:local_limit, fn -> default_local_limit(params, schema) end)
|
|
|
|
params =
|
|
if Map.get(params, :global_limit) do
|
|
Map.update!(params, :global_limit, &cast_global_limit/1)
|
|
else
|
|
params
|
|
end
|
|
|
|
params =
|
|
if Map.get(params, :rate_limit) do
|
|
Map.update!(params, :rate_limit, &cast_rate_limit/1)
|
|
else
|
|
params
|
|
end
|
|
|
|
schema
|
|
|> cast(params, allowed)
|
|
|> cast_embed(:global_limit, with: &global_changeset/2)
|
|
|> cast_embed(:rate_limit, with: &rate_changeset/2)
|
|
|> validate_required(~w(local_limit)a)
|
|
|> validate_number(:local_limit, greater_than: 0)
|
|
|> validate_single_partitioner()
|
|
|> Utils.enforce_keys(params, schema.__struct__)
|
|
end
|
|
|
|
@doc false
|
|
@spec default_local_limit(map(), Ecto.Schema.t()) :: non_neg_integer()
|
|
def default_local_limit(params, schema) do
|
|
cond do
|
|
is_integer(params[:limit]) ->
|
|
params[:limit]
|
|
|
|
is_integer(params[:global_limit]) ->
|
|
params[:global_limit]
|
|
|
|
is_integer(get_in(params, [:global_limit, :allowed])) ->
|
|
get_in(params, [:global_limit, :allowed])
|
|
|
|
true ->
|
|
schema.local_limit
|
|
end
|
|
end
|
|
|
|
@doc false
|
|
@spec validate_single_partitioner(Changeset.t()) :: Changeset.t()
|
|
def validate_single_partitioner(changeset) do
|
|
case {get_field(changeset, :global_limit), get_field(changeset, :rate_limit)} do
|
|
{%{partition: %{}}, %{partition: %{}}} ->
|
|
add_error(changeset, :global_limit, "only one limiter may have partitioning")
|
|
|
|
_ ->
|
|
changeset
|
|
end
|
|
end
|
|
|
|
@doc false
|
|
@spec global_changeset(Ecto.Schema.t(), integer() | map() | Keyword.t()) :: Changeset.t()
|
|
def global_changeset(schema, params) do
|
|
params = Map.update(params, :partition, nil, &normalize_partition/1)
|
|
|
|
schema
|
|
|> cast(params, ~w(allowed burst)a)
|
|
|> cast_embed(:partition, with: &partition_changeset/2)
|
|
|> validate_number(:allowed, greater_than: 0)
|
|
|> Utils.enforce_keys(params, schema.__struct__)
|
|
end
|
|
|
|
@doc false
|
|
@spec rate_changeset(Ecto.Schema.t(), map() | Keyword.t()) :: Changeset.t()
|
|
def rate_changeset(schema, params) do
|
|
params =
|
|
params
|
|
|> Map.update(:period, nil, &Utils.cast_period/1)
|
|
|> Map.update(:partition, nil, &normalize_partition/1)
|
|
|> Map.put_new_lazy(:window_time, fn -> DateTime.to_unix(DateTime.utc_now(), :second) end)
|
|
|
|
fields = schema.__struct__.__schema__(:fields) -- [:partition]
|
|
|
|
schema
|
|
|> cast(params, fields)
|
|
|> cast_embed(:partition, with: &partition_changeset/2)
|
|
|> validate_required(~w(allowed period)a)
|
|
|> validate_number(:allowed, greater_than: 0)
|
|
|> validate_number(:period, greater_than: 0)
|
|
|> clear_windows_on_algorithm_change()
|
|
|> Utils.enforce_keys(params, schema.__struct__)
|
|
end
|
|
|
|
defp clear_windows_on_algorithm_change(changeset) do
|
|
case get_change(changeset, :algorithm) do
|
|
nil -> changeset
|
|
_new_algorithm -> force_change(changeset, :windows, %{})
|
|
end
|
|
end
|
|
|
|
@doc false
|
|
@spec partition_changeset(Ecto.Schema.t(), map() | Keyword.t()) :: Changeset.t()
|
|
def partition_changeset(schema, params) do
|
|
params =
|
|
params
|
|
|> Map.new()
|
|
|> Map.update(:keys, [], &for(key <- &1, do: to_string(key)))
|
|
|
|
schema
|
|
|> cast(params, ~w(fields keys)a)
|
|
|> validate_required(~w(fields)a)
|
|
|> Utils.enforce_keys(params, schema.__struct__)
|
|
end
|
|
|
|
# Global Limit Helpers
|
|
|
|
@doc false
|
|
@spec cast_global_limit(pos_integer() | list() | map()) :: map()
|
|
def cast_global_limit(limit) when is_integer(limit) do
|
|
%{allowed: limit}
|
|
end
|
|
|
|
def cast_global_limit(opts), do: cast_rate_limit(opts)
|
|
|
|
# Rate Limit Helpers
|
|
|
|
@doc false
|
|
@spec cast_rate_limit(list() | map()) :: map()
|
|
def cast_rate_limit([[key | _] | _] = opts) when is_binary(key) do
|
|
opts
|
|
|> Enum.map(&List.to_tuple/1)
|
|
|> Map.new()
|
|
|> cast_rate_limit()
|
|
end
|
|
|
|
def cast_rate_limit(opts)
|
|
when is_map_key(opts, "algorithm") or
|
|
is_map_key(opts, "allowed") or
|
|
is_map_key(opts, "burst") or
|
|
is_map_key(opts, "period") or
|
|
is_map_key(opts, "partition") do
|
|
Map.new(opts, fn
|
|
{"algorithm", algorithm} ->
|
|
{:algorithm, String.to_existing_atom(algorithm)}
|
|
|
|
{"allowed", allowed} ->
|
|
{:allowed, allowed}
|
|
|
|
{"burst", allowed} ->
|
|
{:burst, allowed}
|
|
|
|
{"period", [time, unit]} ->
|
|
{:period, {time, unit}}
|
|
|
|
{"period", period} ->
|
|
{:period, period}
|
|
|
|
{"partition", [["fields", fields]]} ->
|
|
{:partition, fields: fields}
|
|
|
|
{"partition", [["fields", fields], ["keys", keys]]} ->
|
|
{:partition, fields: fields, keys: keys}
|
|
|
|
{"partition", field} when is_binary(field) ->
|
|
{:partition, fields: [field]}
|
|
end)
|
|
end
|
|
|
|
def cast_rate_limit(opts), do: opts
|
|
|
|
def normalize_partition(partition) do
|
|
cond do
|
|
is_nil(partition) ->
|
|
nil
|
|
|
|
is_map(partition) ->
|
|
partition
|
|
|
|
Keyword.keyword?(partition) and Keyword.has_key?(partition, :fields) ->
|
|
partition
|
|
|
|
true ->
|
|
partition
|
|
|> List.wrap()
|
|
|> Enum.reduce([fields: [], keys: []], fn
|
|
{field, keys}, opts when field in ~w(args meta)a ->
|
|
opts
|
|
|> Keyword.put(:keys, List.wrap(keys))
|
|
|> Keyword.update!(:fields, &[field | &1])
|
|
|
|
field, opts ->
|
|
Keyword.update!(opts, :fields, &[field | &1])
|
|
end)
|
|
end
|
|
end
|
|
|
|
# Helpers
|
|
|
|
defp to_clean_string(value) when is_reference(value) do
|
|
inspect(value)
|
|
end
|
|
|
|
defp to_clean_string(value) do
|
|
value
|
|
|> to_string()
|
|
|> String.trim_leading("Elixir.")
|
|
end
|
|
end
|
|
|
|
encoder = if Code.ensure_loaded?(JSON.Encoder), do: JSON.Encoder, else: Jason.Encoder
|
|
|
|
Protocol.derive(encoder, Oban.Pro.Producer, except: [:__meta__])
|
|
Protocol.derive(encoder, Oban.Pro.Producer.Meta.GlobalLimit)
|
|
Protocol.derive(encoder, Oban.Pro.Producer.Meta.GlobalLimit.Partition)
|
|
Protocol.derive(encoder, Oban.Pro.Producer.Meta.RateLimit)
|
|
Protocol.derive(encoder, Oban.Pro.Producer.Meta.RateLimit.Partition)
|