lockstep

defmodule Lockstep.RPC do
  @moduledoc """
  Mirror of Erlang's `:rpc` module under Lockstep's controlled
  scheduling. Cross-node `apply/4`-style calls are modeled by
  spawning a process on the target node, having it execute the
  function, and shipping the result back to the caller via a
  Lockstep-controlled message.

  Partition-aware: if the caller's node and the target are in
  opposing partition groups, the spawn message itself is dropped or
  deferred (per partition mode). Calls in `:drop` partitions error
  with `{:badrpc, :nodedown}`-equivalent semantics; in `:defer`
  partitions they hang until heal (and then complete).

  ## Supported

    * `call/4,5` -- synchronous remote call
    * `multicall/4,5` -- call same function on a list of nodes,
      gather successes and failures
    * `cast/4` -- fire-and-forget remote call
    * `abcast/2,3` -- async broadcast a message to a registered
      atom name on every node
  """

  @doc """
  Synchronous call: invoke `apply(module, fun, args)` on `node`,
  return its result.

  Returns the function's value, or `{:badrpc, reason}` on failure
  (timeout, unreachable node, etc.) -- mirroring `:rpc.call/4-5`.
  """
  @spec call(atom(), module(), atom(), [any()], timeout()) :: any()
  def call(node, module, fun, args, timeout \\ 5000)
      when is_atom(node) and is_atom(module) and is_atom(fun) and is_list(args) do
    parent = self()
    ref = make_ref()

    _ =
      Lockstep.Cluster.spawn(node, fn ->
        result =
          try do
            apply(module, fun, args)
          catch
            kind, reason -> {:badrpc, {kind, reason}}
          end

        Lockstep.send(parent, {__MODULE__, ref, result})
      end)

    # Bound the wait by virtual time using send_after.
    timer_ref = Lockstep.send_after(parent, {__MODULE__, ref, :__rpc_timeout__}, timeout)

    case Lockstep.recv_first(fn
           {__MODULE__, ^ref, _} -> true
           _ -> false
         end) do
      {__MODULE__, ^ref, :__rpc_timeout__} ->
        {:badrpc, :timeout}

      {__MODULE__, ^ref, value} ->
        _ = Lockstep.cancel_timer(timer_ref)
        value
    end
  end

  @doc """
  Call `apply(module, fun, args)` on each node in `nodes`. Returns
  `{successes, bad_nodes}` where `successes` is the list of return
  values (in node order) and `bad_nodes` is a list of nodes that
  errored or timed out.
  """
  @spec multicall([atom()], module(), atom(), [any()], timeout()) ::
          {[any()], [atom()]}
  def multicall(nodes, module, fun, args, timeout \\ 5000)
      when is_list(nodes) and is_atom(module) and is_atom(fun) and is_list(args) do
    results = Enum.map(nodes, &call(&1, module, fun, args, timeout))

    {successes, failures} =
      Enum.zip(nodes, results)
      |> Enum.split_with(fn {_node, result} ->
        not match?({:badrpc, _}, result)
      end)

    {Enum.map(successes, fn {_, v} -> v end), Enum.map(failures, fn {node, _} -> node end)}
  end

  @doc """
  Fire-and-forget call: `apply(module, fun, args)` on `node` with no
  reply. Returns `true` (matching `:rpc.cast/4`'s contract).
  """
  @spec cast(atom(), module(), atom(), [any()]) :: true
  def cast(node, module, fun, args)
      when is_atom(node) and is_atom(module) and is_atom(fun) and is_list(args) do
    _ =
      Lockstep.Cluster.spawn(node, fn ->
        try do
          apply(module, fun, args)
        catch
          _, _ -> :ok
        end
      end)

    true
  end

  @doc """
  Async-broadcast `msg` to a registered atom `name` on every node
  in `nodes`. Mirrors `:rpc.abcast/3`.
  """
  @spec abcast([atom()], atom(), any()) :: :abcast
  def abcast(nodes, name, msg) when is_list(nodes) and is_atom(name) do
    for node <- nodes do
      Lockstep.Cluster.spawn(node, fn ->
        case Process.whereis(name) do
          nil -> :ok
          pid -> Lockstep.send(pid, msg)
        end
      end)
    end

    :abcast
  end

  @doc "Single-arity abcast: broadcasts to all known nodes."
  @spec abcast(atom(), any()) :: :abcast
  def abcast(name, msg) when is_atom(name) do
    abcast(Lockstep.Cluster.list(), name, msg)
  end
end