defmodule Oban.Queue.Producer do @moduledoc false use GenServer alias Oban.{Backoff, Engine, Notifier, PerformError, TimeoutError, Worker} alias Oban.Queue.Executor alias __MODULE__, as: State require Logger defstruct [ :conf, :foreman, :meta, :name, :dispatch_timer, :refresh_timer, dispatch_cooldown: 5, running: %{} ] @spec start_link([Keyword.t()]) :: GenServer.on_start() def start_link(opts) do GenServer.start_link(__MODULE__, opts, name: opts[:name]) end @spec check(GenServer.server()) :: Oban.queue_state() def check(producer) do GenServer.call(producer, :check) end @spec shutdown(GenServer.server()) :: :ok def shutdown(producer) do GenServer.call(producer, :shutdown) end # Callbacks @impl GenServer def init(opts) do Process.flag(:trap_exit, true) {base_opts, meta_opts} = Keyword.split(opts, [:conf, :foreman, :name, :dispatch_cooldown]) state = struct!(State, base_opts) {:ok, state, {:continue, {:start, meta_opts}}} end @impl GenServer def terminate(_reason, %State{} = state) do if is_reference(state.refresh_timer), do: Process.cancel_timer(state.refresh_timer) if is_reference(state.dispatch_timer), do: Process.cancel_timer(state.dispatch_timer) :ok end @impl GenServer def handle_continue({:start, meta_opts}, %State{} = state) do {:ok, meta} = Engine.init(state.conf, meta_opts) :ok = Notifier.listen(state.conf.name, [:insert, :signal]) {:noreply, schedule_refresh(%{state | meta: meta})} end @impl GenServer def handle_info({ref, _val}, %State{} = state) when is_reference(ref) do state = state |> release_ref(ref) |> schedule_dispatch() {:noreply, state} end def handle_info({:DOWN, ref, :process, pid, reason}, %State{} = state) do {^pid, exec} = Map.get(state.running, ref) # The worker module is resolved after storing the exec struct. We try to resolve it again in # order to get a worker's custom `backoff/1`. exec = case Worker.from_string(exec.job.worker) do {:ok, worker} -> %{exec | worker: worker} {:error, _error} -> exec end exec = case reason do %TimeoutError{} -> %{exec | kind: :error, error: reason, state: :failure} :shutdown -> result = {:cancel, :shutdown} reason = PerformError.exception({exec.job.worker, result}) %{exec | result: result, error: reason, state: :cancelled} {error, stack} -> %{exec | kind: {:EXIT, pid}, error: error, stacktrace: stack, state: :failure} _ -> %{exec | kind: {:EXIT, pid}, error: reason, state: :failure} end Task.Supervisor.async_nolink(state.foreman, fn -> exec = exec |> Executor.normalize_state() |> Executor.record_finished() |> Executor.cancel_timeout() Backoff.with_retry(fn -> Executor.report_finished(exec) end) end) state = state |> release_ref(ref) |> schedule_dispatch() {:noreply, state} end def handle_info({:notification, :insert, %{"queue" => queue}}, %State{} = state) do state = if state.meta.queue == queue, do: schedule_dispatch(state), else: state {:noreply, state} end def handle_info({:notification, :signal, payload}, %State{} = state) do queue = state.meta.queue meta = case payload do %{"action" => "pause", "queue" => "*"} -> Engine.put_meta(state.conf, state.meta, :paused, true) %{"action" => "resume", "queue" => "*"} -> Engine.put_meta(state.conf, state.meta, :paused, false) %{"action" => "pause", "queue" => ^queue} -> Engine.put_meta(state.conf, state.meta, :paused, true) %{"action" => "resume", "queue" => ^queue} -> Engine.put_meta(state.conf, state.meta, :paused, false) %{"action" => "scale", "queue" => ^queue} -> payload |> Map.drop(~w(action ident node queue)) |> Enum.reduce(state.meta, fn {key, val}, meta -> Engine.put_meta(state.conf, meta, String.to_existing_atom(key), val) end) %{"action" => "pkill", "job_id" => jid} -> for {ref, {pid, exec}} <- state.running, exec.job.id == jid do pkill(ref, pid, state) end state.meta _ -> state.meta end {:noreply, schedule_dispatch(%{state | meta: meta})} end def handle_info(:dispatch, %State{} = state) do {:noreply, dispatch(%{state | dispatch_timer: nil})} end def handle_info(:refresh, %State{} = state) do meta = Engine.refresh(state.conf, state.meta) {:noreply, schedule_refresh(%{state | meta: meta})} end def handle_info(message, state) do Logger.warning( message: "Received unexpected message: #{inspect(message)}", source: :oban, module: __MODULE__ ) {:noreply, state} end @impl GenServer def handle_call(:check, _from, %State{} = state) do meta = Engine.check_meta(state.conf, state.meta, state.running) {:reply, meta, state} end def handle_call({:put_meta, key, value}, _from, %State{} = state) do meta = Engine.put_meta(state.conf, state.meta, key, value) {:reply, meta, %{state | meta: meta}} end def handle_call(:shutdown, _from, %State{} = state) do meta = Engine.shutdown(state.conf, state.meta) {:reply, :ok, %{state | meta: meta}} end # Killing defp pkill(ref, pid, %State{} = state) do case Task.Supervisor.terminate_child(state.foreman, pid) do :ok -> state {:error, :not_found} -> release_ref(state, ref) end end # Dispatching defp schedule_dispatch(%State{} = state) do if is_reference(state.dispatch_timer) do state else timer = Process.send_after(self(), :dispatch, state.dispatch_cooldown) %{state | dispatch_timer: timer} end end defp schedule_refresh(%State{} = state) do cooldown = Backoff.jitter(state.meta.refresh_interval, mode: :dec, mult: 0.5) timer = Process.send_after(self(), :refresh, cooldown) %{state | refresh_timer: timer} end defp dispatch(%State{} = state) do tele_meta = %{conf: state.conf, queue: state.meta.queue} {meta, dispatched} = :telemetry.span([:oban, :producer], tele_meta, fn -> {meta, dispatched} = start_jobs(state) {{meta, dispatched}, Map.put(tele_meta, :dispatched_count, map_size(dispatched))} end) %{state | meta: meta, running: Map.merge(state.running, dispatched)} end defp start_jobs(%State{} = state) do {:ok, {meta, jobs}} = Engine.fetch_jobs(state.conf, state.meta, state.running) dispatched = for job <- jobs, into: %{} do exec = Executor.new(state.conf, job) task = Task.Supervisor.async_nolink(state.foreman, Executor, :call, [exec]) {task.ref, {task.pid, exec}} end {meta, dispatched} end defp release_ref(%State{} = state, ref) do Process.demonitor(ref, [:flush]) %{state | running: Map.delete(state.running, ref)} end end