rpc_load_balancer

defmodule RpcLoadBalancer do
  @moduledoc """
  Distributed RPC load balancer built on `:pg`.

  Acts as a per-instance Supervisor that starts the caches and GenServer
  needed for a single load balancer. Also provides the public API for
  node selection, RPC calls/casts, and low-level `:erpc` wrappers.

  ## Starting a load balancer

      RpcLoadBalancer.start_link(
        name: :my_lb,
        selection_algorithm: RpcLoadBalancer.LoadBalancer.SelectionAlgorithm.RoundRobin,
        algorithm_opts: [weights: %{node() => 1}]
      )
  """

  use Supervisor

  alias RpcLoadBalancer.LoadBalancer.Drainer
  alias RpcLoadBalancer.LoadBalancer.SelectionAlgorithm
  alias RpcLoadBalancer.Retry

  @type name :: atom()

  @pg_group_name RpcLoadBalancer.LoadBalancer.Pg.pg_group_name()
  @telemetry_prefix [:rpc_load_balancer, :rpc]

  # -------------------------------------------------------------------
  # Supervisor
  # -------------------------------------------------------------------

  @spec start_link(keyword()) :: Supervisor.on_start()
  def start_link(opts) do
    name = Keyword.fetch!(opts, :name)
    Supervisor.start_link(__MODULE__, opts, name: name)
  end

  @impl true
  def init(opts) do
    algorithm = Keyword.get(opts, :selection_algorithm, SelectionAlgorithm.Random)
    algorithm_opts = Keyword.get(opts, :algorithm_opts, [])
    name = Keyword.fetch!(opts, :name)

    algorithm_children = SelectionAlgorithm.child_specs(algorithm, name, algorithm_opts)

    # Algorithm children (pollers, counters) must be registered BEFORE the
    # LoadBalancer GenServer so that callers who invoke `select_node/1` as
    # soon as `start_link/1` returns can reach these processes. They do not
    # need to have produced data yet — selection tolerates a cold cache and
    # falls back to per-algorithm defaults.
    #
    # Shared caches (`{Cache, [...]}`) are owned by `RpcLoadBalancer.Application`,
    # not this per-LB supervisor — this keeps cache agent lifetime bound to
    # the VM instead of any individual load balancer.
    children = algorithm_children ++ [{RpcLoadBalancer.LoadBalancer, opts}]

    Supervisor.init(children, strategy: :one_for_all)
  end

  # -------------------------------------------------------------------
  # Public API
  # -------------------------------------------------------------------

  @spec get_members(name()) :: {:ok, [node()]} | {:error, ErrorMessage.t()}
  def get_members(load_balancer_name) do
    case :pg.get_members(@pg_group_name, load_balancer_name) do
      [] ->
        {:error,
         ErrorMessage.service_unavailable(
           "no members registered",
           %{load_balancer: load_balancer_name}
         )}

      pids ->
        nodes =
          pids
          |> Enum.map(&node/1)
          |> Enum.uniq()

        {:ok, nodes}
    end
  end

  @spec select_node(name(), keyword()) :: {:ok, node()} | {:error, ErrorMessage.t()}
  def select_node(load_balancer_name, opts \\ []) do
    with {:ok, algorithm} <- SelectionAlgorithm.get_algorithm(load_balancer_name),
         {:ok, members} <- get_members(load_balancer_name) do
      node = SelectionAlgorithm.choose_from_nodes(algorithm, load_balancer_name, members, opts)
      {:ok, node}
    end
  end

  # -------------------------------------------------------------------
  # :erpc wrappers
  # -------------------------------------------------------------------

  @spec call(node(), module(), atom(), [any()], keyword()) :: ErrorMessage.t_res(any())
  def call(node, module, fun, args, opts \\ [])

  def call(node, module, fun, args, opts) when is_atom(node) do
    metadata = build_span_metadata(:call, node, module, fun, opts)

    :telemetry.span(@telemetry_prefix, metadata, fn ->
      result =
        case Keyword.get(opts, :load_balancer) do
          nil -> erpc_call(node, module, fun, args, opts)
          name -> lb_call(name, module, fun, args, opts)
        end

      {result, Map.put(metadata, :status, status_from_result(result))}
    end)
  end

  @spec cast(node(), module(), atom(), [term()], keyword()) :: :ok | {:error, ErrorMessage.t()}
  def cast(node, module, fun, args, opts \\ [])

  def cast(node, module, fun, args, opts) when is_atom(node) do
    metadata = build_span_metadata(:cast, node, module, fun, opts)

    :telemetry.span(@telemetry_prefix, metadata, fn ->
      result =
        case Keyword.get(opts, :load_balancer) do
          nil -> erpc_cast(node, module, fun, args)
          name -> lb_cast(name, module, fun, args, opts)
        end

      {result, Map.put(metadata, :status, status_from_result(result))}
    end)
  end

  defp build_span_metadata(type, node, module, fun, opts) do
    %{
      type: type,
      node: node,
      module: inspect(module),
      function: fun,
      load_balancer: Keyword.get(opts, :load_balancer)
    }
  end

  defp status_from_result(:ok), do: :ok
  defp status_from_result({:ok, _}), do: :ok
  defp status_from_result({:error, %ErrorMessage{code: code}}), do: code
  defp status_from_result({:error, _}), do: :error

  defp lb_call(load_balancer_name, module, fun, args, opts) do
    call_directly? = Keyword.get(opts, :call_directly?, RpcLoadBalancer.Config.call_directly?())

    if call_directly? do
      {:ok, apply(module, fun, args)}
    else
      with {:ok, algorithm} <- SelectionAlgorithm.get_algorithm(load_balancer_name) do
        if SelectionAlgorithm.local?(algorithm) do
          {:ok, apply(module, fun, args)}
        else
          with_no_route_retry(load_balancer_name, opts, fn ->
            route_and_run(load_balancer_name, algorithm, opts, fn selected_node ->
              erpc_call(selected_node, module, fun, args, Keyword.take(opts, [:timeout]))
            end)
          end)
        end
      end
    end
  end

  defp lb_cast(load_balancer_name, module, fun, args, opts) do
    call_directly? = Keyword.get(opts, :call_directly?, RpcLoadBalancer.Config.call_directly?())

    if call_directly? do
      spawn(module, fun, args)
      :ok
    else
      with {:ok, algorithm} <- SelectionAlgorithm.get_algorithm(load_balancer_name) do
        if SelectionAlgorithm.local?(algorithm) do
          spawn(module, fun, args)
          :ok
        else
          with_no_route_retry(load_balancer_name, opts, fn ->
            route_and_run(load_balancer_name, algorithm, opts, fn selected_node ->
              erpc_cast(selected_node, module, fun, args)
            end)
          end)
        end
      end
    end
  end

  # Selects a member for the load balancer and runs `run_fun.(selected_node)`,
  # releasing the node afterward. Returns `:retry` when no member is registered
  # so the surrounding `with_no_route_retry/3` can back off and try again
  # (cluster boot, rolling restart). Other selection errors pass straight back.
  defp route_and_run(load_balancer_name, algorithm, opts, run_fun) do
    case select_node(load_balancer_name, Keyword.take(opts, [:key, :call_directly?, :load_balancer])) do
      {:error, %ErrorMessage{code: :service_unavailable}} ->
        :retry

      {:error, _reason} = error ->
        error

      {:ok, selected_node} ->
        with_drainer(load_balancer_name, fn ->
          result = run_fun.(selected_node)
          SelectionAlgorithm.release_node(algorithm, load_balancer_name, selected_node)
          result
        end)
    end
  end

  defp with_no_route_retry(load_balancer_name, opts, dispatch_fun) do
    case Retry.with_retry(opts, dispatch_fun) do
      :error -> no_members_error(load_balancer_name)
      result -> result
    end
  end

  defp no_members_error(load_balancer_name) do
    {:error,
     ErrorMessage.service_unavailable("no members registered", %{
       load_balancer: load_balancer_name
     })}
  end

  defp erpc_call(node, module, fun, args, opts) do
    timeout = Keyword.get(opts, :timeout, :timer.seconds(10))

    try do
      {:ok, :erpc.call(node, module, fun, args, timeout)}
    rescue
      e in ErlangError ->
        {:error, erlang_error_to_error_message(e, node)}

      e ->
        {:error, ErrorMessage.service_unavailable("unavailable", %{node: node, details: e})}
    end
  end

  defp erpc_cast(node, module, fun, args) do
    :erpc.cast(node, module, fun, args)
  rescue
    e in ErlangError ->
      {:error, erlang_error_to_error_message(e, node)}

    e ->
      {:error, ErrorMessage.service_unavailable("unavailable", %{node: node, details: e})}
  end

  # -------------------------------------------------------------------
  # Random-node helpers
  # -------------------------------------------------------------------

  @spec call_on_random_node(String.t(), module(), atom(), [any()], keyword()) ::
          ErrorMessage.t_res(any())
  def call_on_random_node(node_filter, module, fun, args, opts \\ []) do
    call_directly? = Keyword.get(opts, :call_directly?, RpcLoadBalancer.Config.call_directly?())
    load_balancer_name = Keyword.get(opts, :load_balancer)

    if call_directly? or current_node_matches_filter?(node_filter) do
      {:ok, apply(module, fun, args)}
    else
      no_nodes_error(Retry.with_retry(opts, fn ->
        case filter_nodes(node_filter) do
          [] ->
            :retry

          node_list ->
            selected_node = Enum.random(node_list)
            with_drainer(load_balancer_name, fn ->
              call(selected_node, module, fun, args, Keyword.take(opts, [:timeout]))
            end)
        end
      end), node_filter)
    end
  end

  @spec cast_on_random_node(String.t(), module(), atom(), [any()], keyword()) ::
          :ok | {:error, ErrorMessage.t()}
  def cast_on_random_node(node_filter, module, fun, args, opts \\ []) do
    call_directly? = Keyword.get(opts, :call_directly?, RpcLoadBalancer.Config.call_directly?())
    load_balancer_name = Keyword.get(opts, :load_balancer)

    if call_directly? or current_node_matches_filter?(node_filter) do
      spawn(module, fun, args)
      :ok
    else
      no_nodes_error(Retry.with_retry(opts, fn ->
        case filter_nodes(node_filter) do
          [] ->
            :retry

          node_list ->
            selected_node = Enum.random(node_list)
            with_drainer(load_balancer_name, fn ->
              cast(selected_node, module, fun, args)
            end)
        end
      end), node_filter)
    end
  end

  # -------------------------------------------------------------------
  # Private
  # -------------------------------------------------------------------

  defp erlang_error_to_error_message(%ErlangError{original: {:erpc, :timeout}}, node) do
    ErrorMessage.request_timeout("timeout", %{node: node})
  end

  defp erlang_error_to_error_message(%ErlangError{original: {:erpc, :noconnection}}, node) do
    ErrorMessage.service_unavailable("noconnection", %{node: node})
  end

  defp erlang_error_to_error_message(%ErlangError{original: {:erpc, :badarg}}, node) do
    ErrorMessage.bad_request("bad request", %{node: node})
  end

  defp erlang_error_to_error_message(%ErlangError{} = error, _node) do
    ErrorMessage.service_unavailable("unavailable", %{details: error})
  end

  defp with_drainer(nil, fun), do: fun.()

  defp with_drainer(load_balancer_name, fun) do
    drainer_index = Drainer.register(load_balancer_name)
    Drainer.track_call(drainer_index)

    try do
      fun.()
    after
      Drainer.release_call(drainer_index)
    end
  end

  defp no_nodes_error(:error, node_filter) do
    {:error,
     ErrorMessage.service_unavailable(
       "no nodes in cluster found with that filter",
       %{node_filter: node_filter}
     )}
  end

  defp no_nodes_error(result, _node_filter), do: result

  defp filter_nodes(node_filter) do
    Enum.filter(Node.list(), &(to_string(&1) =~ node_filter))
  end

  defp current_node_matches_filter?(node_filter) do
    to_string(node()) =~ node_filter
  end
end