prop/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)
  @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)

  @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(@keys_cache_ttl, __MODULE__, :__clear_cache__, [@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__(ttl) do
    now = System.monotonic_time(:millisecond)

    :ets.select_delete(@keys_tab, [
      {{:_, :_, :"$3"}, [{:>=, {:-, now, :"$3"}, ttl}], [true]}
    ])
  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
    timed_cache(conf.name, queue, 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

  defp timed_cache(name, queue, fun) do
    cache_key = {name, queue}
    now = System.monotonic_time(:millisecond)

    case :ets.lookup(@keys_tab, cache_key) do
      [{^cache_key, keys, timestamp}] when now - timestamp < @keys_cache_ttl ->
        keys

      _ ->
        keys = fun.()
        :ets.insert(@keys_tab, {cache_key, keys, now})
        keys
    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
      case :ets.lookup(@conf_tab, {conf.name, queue}) do
        [{_key, value}] ->
          value

        [] ->
          value =
            with :missing <- reg_config(conf, queue),
                 :missing <- pro_config(conf, queue),
                 do: dyn_config(conf, queue)

          :ets.insert(@conf_tab, {{conf.name, queue}, value})

          value
      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