lockstep

defmodule Lockstep.Registry do
  @moduledoc """
  Lockstep-aware process registry. Models the most common subset of
  OTP `Registry` under Lockstep's controller, so libraries that depend
  on key-based process lookup work under controlled scheduling.

  ## Supported

    * Unique keys (`keys: :unique`) and duplicate keys (`keys: :duplicate`).
    * `register/3`, `unregister/2`, `lookup/2`, `dispatch/3`,
      `count/1`, `keys/2`.
    * Automatic cleanup on registering-process death (via
      `Lockstep.monitor`).
    * `:via` tuple integration with `Lockstep.GenServer`:

          via = {:via, Lockstep.Registry, {reg, :my_name}}
          {:ok, _} = Lockstep.GenServer.start_link(MyMod, [], name: via)
          Lockstep.GenServer.call(via, :req)

  ## Not modeled

    * Partitioned registries (`partitions:` option). Single-partition only.
    * Listeners.
    * `match/4` patterns -- only `:_` lookup via `lookup/2`.
    * Meta key/value store (`meta/2`, `put_meta/3`).

  Most production callers don't depend on the omitted bits; if you need
  them, file an issue.
  """

  alias Lockstep.GenServer, as: LGS

  @type registry :: pid()

  # ============================================================
  # Public API
  # ============================================================

  @doc """
  Start a registry. Options:

    * `:keys` -- `:unique` (default) or `:duplicate`.
    * `:name` -- accepted for shape compatibility with `Registry`,
      but Lockstep does not register the registry itself by name
      (pass the returned pid around or use `:via` only on regular
      Lockstep.GenServers).
  """
  @spec start_link(keyword()) :: {:ok, pid()}
  def start_link(opts \\ []) do
    keys = Keyword.get(opts, :keys, :unique)
    name = Keyword.get(opts, :name)
    initial_meta = Keyword.get(opts, :meta, [])

    unless keys in [:unique, :duplicate] do
      raise ArgumentError, "Lockstep.Registry: keys must be :unique or :duplicate"
    end

    init = %{keys: keys, initial_meta: initial_meta}

    if is_atom(name) and name != nil do
      LGS.start_link(__MODULE__.Server, init, name: name)
    else
      LGS.start_link(__MODULE__.Server, init)
    end
  end

  @doc """
  Register the calling process under `key` with `value`.

  For `:unique` registries returns `{:error, {:already_registered, pid}}`
  when the key is taken. `:duplicate` always succeeds.
  """
  @spec register(registry(), any(), any()) ::
          {:ok, pid()} | {:error, {:already_registered, pid()}}
  def register(reg, key, value) do
    LGS.call(reg, {:register, self(), key, value})
  end

  @doc """
  Remove the calling process's registration under `key`. Always returns
  `:ok` (no-op if no such registration).
  """
  @spec unregister(registry(), any()) :: :ok
  def unregister(reg, key) do
    LGS.call(reg, {:unregister, self(), key})
  end

  @doc """
  Look up entries under `key`. Returns `[{pid, value}]` (one entry for
  unique keys, possibly many for duplicate keys), or `[]` if no
  registration.
  """
  @spec lookup(registry(), any()) :: [{pid(), any()}]
  def lookup(reg, key) do
    LGS.call(reg, {:lookup, key})
  end

  @doc """
  Number of registered entries across all keys.
  """
  @spec count(registry()) :: non_neg_integer()
  def count(reg) do
    LGS.call(reg, :count)
  end

  @doc """
  Run an ETS match spec against the registry's entries. Each entry
  is shaped as `{key, pid, value}`, matching OTP `Registry.select/2`.
  Useful for libraries that enumerate all subscribers via
  `Registry.select`.
  """
  @spec select(registry(), :ets.match_spec()) :: [any()]
  def select(reg, match_spec) do
    LGS.call(reg, {:select, match_spec})
  end

  @doc """
  Keys registered by `pid` in `reg`. Same shape as `Registry.keys/2`.
  """
  @spec keys(registry(), pid()) :: [any()]
  def keys(reg, pid) when is_pid(pid) do
    LGS.call(reg, {:keys, pid})
  end

  @doc """
  Dispatch a callback over every {pid, value} registered under `key`.
  The callback may be:

    * a 1-arity function -- receives `[{pid, value}, ...]`
    * an `{module, function, args}` MFA tuple -- invoked as
      `apply(module, function, [entries | args])`, matching the
      shape OTP's `Registry.dispatch/4` accepts.
  """
  @spec dispatch(registry(), any(), ([{pid(), any()}] -> any()) | {module(), atom(), [any()]}) ::
          :ok
  def dispatch(reg, key, fun) when is_function(fun, 1) do
    entries = lookup(reg, key)
    fun.(entries)
    :ok
  end

  def dispatch(reg, key, {mod, fun, args}) when is_atom(mod) and is_atom(fun) and is_list(args) do
    entries = lookup(reg, key)
    apply(mod, fun, [entries | args])
    :ok
  end

  @doc false
  def dispatch(reg, key, fun, _opts), do: dispatch(reg, key, fun)

  @doc """
  Read a `key`-keyed metadata value previously stored via `put_meta/3`.
  Returns `{:ok, value}` if found, `:error` otherwise. Same shape as
  OTP `Registry.meta/2`.
  """
  @spec meta(registry(), any()) :: {:ok, any()} | :error
  def meta(reg, key) do
    LGS.call(reg, {:get_meta, key})
  end

  @doc """
  Store metadata under `key`. Returns `:ok`. Same shape as
  OTP `Registry.put_meta/3`.
  """
  @spec put_meta(registry(), any(), any()) :: :ok
  def put_meta(reg, key, value) do
    LGS.call(reg, {:put_meta, key, value})
  end

  # ============================================================
  # :via callback contract (Module.register_name/2 and friends)
  # ============================================================

  @doc false
  def register_name({reg, key}, pid) do
    case LGS.call(reg, {:register, pid, key, nil}) do
      {:ok, _} -> :yes
      {:error, _} -> :no
    end
  end

  @doc false
  def unregister_name({reg, key}) do
    LGS.call(reg, {:unregister_pid_via, key})
    :ok
  end

  @doc false
  def whereis_name({reg, key}) do
    case LGS.call(reg, {:lookup, key}) do
      [{pid, _}] -> pid
      _ -> :undefined
    end
  end

  @doc false
  def send({reg, key}, msg) do
    case whereis_name({reg, key}) do
      :undefined -> exit({:badarg, {{reg, key}, msg}})
      pid -> Lockstep.send(pid, msg)
    end
  end
end

defmodule Lockstep.Registry.Server do
  @moduledoc false

  # State:
  #   keys: :unique | :duplicate
  #   table: %{key => entry | [entry]}      (entry = {pid, value})
  #   monitors: %{ref => {pid, key}}        (entries by monitor ref)
  #   pid_index: %{pid => MapSet.new(keys)} (for fast Registry.keys/2)

  def init(%{keys: keys} = init) do
    initial_meta = Map.get(init, :initial_meta, [])
    meta_map = Enum.into(initial_meta, %{})
    {:ok, %{keys: keys, table: %{}, monitors: %{}, pid_index: %{}, meta: meta_map}}
  end

  def handle_call({:register, pid, key, value}, _from, state) do
    case state.keys do
      :unique ->
        case Map.get(state.table, key) do
          nil ->
            ref = Lockstep.monitor(pid)

            state =
              state
              |> put_in([:table, key], {pid, value})
              |> put_in([:monitors, ref], {pid, key})
              |> add_pid_key(pid, key)

            {:reply, {:ok, self()}, state}

          {existing_pid, _} ->
            {:reply, {:error, {:already_registered, existing_pid}}, state}
        end

      :duplicate ->
        ref = Lockstep.monitor(pid)
        existing = Map.get(state.table, key, [])

        state =
          state
          |> put_in([:table, key], [{pid, value} | existing])
          |> put_in([:monitors, ref], {pid, key})
          |> add_pid_key(pid, key)

        {:reply, {:ok, self()}, state}
    end
  end

  def handle_call({:unregister, pid, key}, _from, state) do
    state = drop_registration(state, pid, key)
    {:reply, :ok, state}
  end

  # :via path -- caller is OTP framework, but it's calling on behalf of
  # an action that's about to terminate. The "pid to drop" is the one
  # currently registered under this key (any of them).
  def handle_call({:unregister_pid_via, key}, _from, state) do
    case Map.get(state.table, key) do
      nil ->
        {:reply, :ok, state}

      {pid, _} ->
        state = drop_registration(state, pid, key)
        {:reply, :ok, state}

      list when is_list(list) ->
        state =
          Enum.reduce(list, state, fn {pid, _}, acc ->
            drop_registration(acc, pid, key)
          end)

        {:reply, :ok, state}
    end
  end

  def handle_call({:lookup, key}, _from, state) do
    result =
      case Map.get(state.table, key) do
        nil -> []
        {_, _} = entry -> [entry]
        list when is_list(list) -> list
      end

    {:reply, result, state}
  end

  def handle_call(:count, _from, state) do
    n =
      Enum.reduce(state.table, 0, fn
        {_k, list}, acc when is_list(list) -> acc + length(list)
        {_k, _entry}, acc -> acc + 1
      end)

    {:reply, n, state}
  end

  def handle_call({:keys, pid}, _from, state) do
    keys =
      case Map.get(state.pid_index, pid) do
        nil -> []
        set -> MapSet.to_list(set)
      end

    {:reply, keys, state}
  end

  def handle_call({:select, match_spec}, _from, state) do
    # Flatten all registrations into {key, pid, value} tuples
    # (the OTP Registry layout that match_spec expects).
    entries =
      Enum.flat_map(state.table, fn
        {key, list} when is_list(list) ->
          for {pid, value} <- list, do: {key, pid, value}

        {key, {pid, value}} ->
          [{key, pid, value}]
      end)

    # Compile and run the match spec against the in-memory list.
    compiled = :ets.match_spec_compile(match_spec)
    result = :ets.match_spec_run(entries, compiled)

    {:reply, result, state}
  end

  def handle_call({:get_meta, key}, _from, state) do
    {:reply, Map.fetch(state.meta, key), state}
  end

  def handle_call({:put_meta, key, value}, _from, state) do
    {:reply, :ok, %{state | meta: Map.put(state.meta, key, value)}}
  end

  def handle_info({:DOWN, ref, :process, _pid, _reason}, state) do
    case Map.pop(state.monitors, ref) do
      {nil, _} ->
        {:noreply, state}

      {{pid, key}, monitors} ->
        state = %{state | monitors: monitors}
        state = drop_pid_under_key(state, pid, key)
        {:noreply, state}
    end
  end

  def handle_info(_, state), do: {:noreply, state}

  # ----- helpers -----

  defp drop_registration(state, pid, key) do
    case Map.get(state.table, key) do
      nil ->
        state

      {^pid, _} ->
        state =
          state
          |> Map.update!(:table, &Map.delete(&1, key))
          |> drop_monitor_for(pid, key)
          |> remove_pid_key(pid, key)

        state

      {_other, _} ->
        # Different pid owns key -- nothing to drop.
        state

      list when is_list(list) ->
        new_list = Enum.reject(list, fn {p, _} -> p == pid end)

        state =
          if new_list == [] do
            Map.update!(state, :table, &Map.delete(&1, key))
          else
            put_in(state, [:table, key], new_list)
          end

        state
        |> drop_monitor_for(pid, key)
        |> remove_pid_key(pid, key)
    end
  end

  defp drop_pid_under_key(state, pid, key) do
    drop_registration(state, pid, key)
  end

  defp drop_monitor_for(state, pid, key) do
    monitors =
      Enum.reject(state.monitors, fn {_ref, {p, k}} -> p == pid and k == key end)
      |> Map.new()

    %{state | monitors: monitors}
  end

  defp add_pid_key(state, pid, key) do
    update_in(state.pid_index, fn idx ->
      Map.update(idx, pid, MapSet.new([key]), &MapSet.put(&1, key))
    end)
  end

  defp remove_pid_key(state, pid, key) do
    update_in(state.pid_index, fn idx ->
      case Map.get(idx, pid) do
        nil ->
          idx

        set ->
          set = MapSet.delete(set, key)
          if MapSet.size(set) == 0, do: Map.delete(idx, pid), else: Map.put(idx, pid, set)
      end
    end)
  end
end