towerops/vendor/oban_pro/lib/oban/pro/partition.ex

defmodule Oban.Pro.Partition do
  @moduledoc false

  @behaviour Oban.Pro.Handler

  import Ecto.Query

  alias Ecto.Changeset
  alias Oban.{Config, Job, Repo}
  alias Oban.Pro.{Producer, Queue, Utils}

  @keys_cache_ttl Application.compile_env(:oban_pro, [__MODULE__, :keys_cache_ttl], 3_000)
  @conf_cache_ttl Application.compile_env(:oban_pro, [__MODULE__, :conf_cache_ttl], 60_000)
  @keys_limit_mlt Application.compile_env(:oban_pro, [__MODULE__, :keys_limit_multiplier], 3)

  @conf_tab Module.concat(__MODULE__, Conf)
  @keys_tab Module.concat(__MODULE__, Keys)

  defmacrop coalesce_partition(meta) do
    quote do
      fragment(
        "COALESCE(?->'global_limit'->'partition', ?->'rate_limit'->'partition')",
        unquote(meta),
        unquote(meta)
      )
    end
  end

  @impl Oban.Pro.Handler
  def on_start do
    :ets.new(@conf_tab, [:set, :named_table, :public, read_concurrency: true])
    :ets.new(@keys_tab, [:set, :named_table, :public, read_concurrency: true])

    # Clear the cache with a timeout, as not all queues are guaranteed to be running on a given
    # node and the cache won't clear naturally.
    :timer.apply_interval(@conf_cache_ttl, __MODULE__, :__clear_cache__, [
      @conf_tab,
      @conf_cache_ttl
    ])

    :timer.apply_interval(@keys_cache_ttl, __MODULE__, :__clear_cache__, [
      @keys_tab,
      @keys_cache_ttl
    ])

    :telemetry.attach(
      "oban.pro.partition",
      [:oban, :engine, :put_meta, :stop],
      &__MODULE__.handle_event/4,
      nil
    )
  end

  @impl Oban.Pro.Handler
  def on_stop do
    :telemetry.detach("oban.pro.partition")
  end

  def __clear_cache__(tab, ttl) do
    now = System.monotonic_time(:millisecond)
    spec = {{:_, :_, :"$3"}, [{:>=, {:-, now, :"$3"}, ttl}], [true]}

    :ets.select_delete(tab, [spec])
  end

  def handle_event(_event, _measure, %{conf: conf, key: key, meta: meta}, nil) do
    # Bust the partition cache on any change to a partitionable limit. This is rare enough an
    # event that it's better to be safe.
    if key in ~w(global_limit rate_limit)a do
      :ets.delete(@conf_tab, {conf.name, meta.queue})
    end
  end

  @spec get_key(Job.t() | Job.changeset(), any()) :: any()
  def get_key(changeset, default \\ nil)

  def get_key(%{changes: %{meta: %{partition_key: key}}}, _default), do: key
  def get_key(%{meta: %{"partition_key" => key}}, _default), do: key
  def get_key(_changeset, default), do: default

  @spec with_partition_meta(Job.changeset(), Config.t()) :: Job.changeset()
  def with_partition_meta(changeset, conf) do
    queue = Changeset.get_field(changeset, :queue, "default")

    case config_for_queue(conf, queue) do
      :missing ->
        changeset

      partition ->
        part_meta = %{partition: true, partition_key: gen_key(changeset, partition)}

        meta =
          changeset
          |> Changeset.get_change(:meta, %{})
          |> Map.merge(part_meta)

        Changeset.put_change(changeset, :meta, meta)
    end
  end

  @spec gen_key(Job.changeset(), %{fields: [atom()], keys: [atom()]}) :: String.t()
  def gen_key(changeset, %{fields: fields, keys: keys}) do
    fields
    |> Enum.sort()
    |> Enum.map(fn
      :args -> Utils.take_keys(changeset, :args, keys)
      :meta -> Utils.take_keys(changeset, :meta, keys)
      :worker -> Changeset.get_field(changeset, :worker)
    end)
    |> Utils.hash64()
  end

  @spec available_keys(Config.t(), String.t(), pos_integer()) :: [String.t()]
  def available_keys(conf, queue, limit) do
    Utils.ets_cache(@keys_tab, {conf.name, queue}, @keys_cache_ttl, fn ->
      query_limit = @keys_limit_mlt * limit

      subquery =
        Job
        |> where([j], j.state == "available" and j.queue == ^queue)
        |> where([j], fragment("meta \\? 'partition_key'"))
        |> group_by([j], fragment("meta->>'partition_key'"))
        |> order_by([j], asc: min(j.priority), asc: min(j.scheduled_at))
        |> select([j], %{partition_key: fragment("meta->>'partition_key'")})
        |> limit(^query_limit)

      partitioned_keys =
        subquery
        |> subquery()
        |> select([s], s.partition_key)
        |> then(&Repo.all(conf, &1))

      ["none" | partitioned_keys]
    end)
  end

  # Rescuing

  @doc false
  def repair_partitions(conf, opts) do
    limit = Keyword.fetch!(opts, :limit)
    timeout = Keyword.fetch!(opts, :timeout)

    for {queue, partition} <- partitioned_queues(conf, timeout) do
      query =
        Job
        |> where([j], j.queue == ^queue)
        |> where([j], j.state in ~w(available scheduled retryable))
        |> where([j], fragment("NOT ? \\? 'partition_key'", j.meta))
        |> limit(^limit)

      conf
      |> Repo.all(query, timeout: timeout)
      |> Enum.group_by(fn job -> gen_partition_key(job, partition) end)
      |> Enum.each(fn {key, grouped_jobs} ->
        ids = Enum.map(grouped_jobs, & &1.id)

        Job
        |> where([j], j.id in ^ids)
        |> update([j], set: [meta: fragment("jsonb_set(?, '{partition_key}', ?)", j.meta, ^key)])
        |> then(&Repo.update_all(conf, &1, [], timeout: timeout))
      end)
    end

    :ok
  end

  defp gen_partition_key(job, partition) do
    job
    |> Ecto.Changeset.change()
    |> gen_key(partition)
  end

  defp partitioned_queues(conf, timeout) do
    normalize_partition = fn partition ->
      fields = partition |> Map.get("fields", []) |> Enum.map(&String.to_existing_atom/1)

      %{fields: fields, keys: Map.get(partition, "keys", [])}
    end

    Producer
    |> where([p], fragment("jsonb_typeof(?) = 'object'", coalesce_partition(p.meta)))
    |> select([p], {p.queue, coalesce_partition(p.meta)})
    |> distinct(true)
    |> then(&Repo.all(conf, &1, timeout: timeout))
    |> Enum.map(fn {queue, partition} -> {queue, normalize_partition.(partition)} end)
  end

  @doc """
  We can't guarantee that the queue is running locally, running on a current node, or running at
  all. This will check the following locations:

  - From a locally running producer
  - From another node's producer in the database
  - From a DynamicQueues entry in the database
  """
  @spec config_for_queue(Config.t(), String.t()) :: :missing | %{fields: list(), keys: list()}
  def config_for_queue(conf, queue) do
    # Partitioning isn't used for draining and there's no point in hitting the database just to
    # find out there's no config available.
    if Process.get(:oban_draining) do
      :missing
    else
      Utils.ets_cache(@conf_tab, {conf.name, queue}, @conf_cache_ttl, fn ->
        with :missing <- reg_config(conf, queue),
             :missing <- pro_config(conf, queue),
             do: dyn_config(conf, queue)
      end)
    end
  end

  defp reg_config(conf, queue) do
    case Registry.meta(Oban.Registry, {conf.name, {:producer, queue}}) do
      {:ok, %{meta: %{global_limit: %{partition: %{} = partition}}}} -> partition
      {:ok, %{meta: %{rate_limit: %{partition: %{} = partition}}}} -> partition
      _other -> :missing
    end
  end

  defp pro_config(conf, queue) do
    query =
      Producer
      |> where(queue: ^queue)
      |> order_by(desc: :started_at)
      |> limit(1)

    case Repo.one(conf, query) do
      %{meta: %{global_limit: %{partition: %{} = partition}}} -> partition
      %{meta: %{rate_limit: %{partition: %{} = partition}}} -> partition
      _other -> :missing
    end
  end

  defp dyn_config(conf, queue) do
    case Repo.get_by(conf, Queue, name: queue) do
      %{opts: %{global_limit: %{partition: %{} = partition}}} -> partition
      %{opts: %{rate_limit: %{partition: %{} = partition}}} -> partition
      _other -> :missing
    end
  end
end