defmodule GenRegistry do
@moduledoc """
GenRegistry provides a `Registry` like interface for managing processes.
"""
@behaviour GenRegistry.Behaviour
use GenServer
alias :ets, as: ETS
alias GenRegistry.Types
defstruct [:worker_module, :worker_type, :workers]
@typedoc """
GenRegistry State.
- worker_module: Module to spawn
- worker_type: `:supervisor` if the worker_module is a supervisor, `:worker` otherwise
- workers: ETS table id holding the worker tracking records.
"""
@type t :: %__MODULE__{
worker_module: module,
worker_type: :supervisor | :worker,
workers: ETS.tab()
}
@gen_module Application.get_env(:gen_registry, :gen_module, GenServer)
## Client
@doc """
Callback called by `Supervisor.init/2`
It is required that you provide a `:worker_module` argument or the call will fail.
"""
@spec child_spec(opts :: Keyword.t()) :: Supervisor.child_spec()
def child_spec(opts) do
worker_module = Keyword.fetch!(opts, :worker_module)
opts =
opts
|> Keyword.delete(:worker_module)
|> Keyword.put_new(:name, worker_module)
%{
id: opts[:name],
start: {__MODULE__, :start_link, [worker_module, opts]},
type: :supervisor
}
end
@doc """
Start a registry instance.
GenRegistry should be run under a supervision tree, it is not recommended to call this directly.
"""
@spec start_link(module, Keyword.t()) :: {:ok, pid} | {:error, any}
def start_link(module, opts \\ []) do
GenServer.start_link(__MODULE__, {module, opts[:name]}, opts)
end
@doc """
Lookup a running a process.
This is a fast path to the ETS table.
"""
@spec lookup(table :: ETS.tab(), id :: Types.id()) :: {:ok, pid} | {:error, :not_found}
def lookup(table, id) do
case ETS.lookup(table, id) do
[{^id, pid}] -> {:ok, pid}
[] -> {:error, :not_found}
end
end
@doc """
Attempts to lookup a running process by id.
If the id is not associated with a running process then it is spawned, the optional third
argument will be passed to `start_link` of the `worker_module` to spawn a new process.
"""
@spec lookup_or_start(
registry :: GenServer.server(),
id :: Types.id(),
args :: [any],
timeout :: integer
) ::
{:ok, pid} | {:error, any}
def lookup_or_start(registry, id, args \\ [], timeout \\ 5_000)
def lookup_or_start(registry, id, args, timeout) when is_atom(registry) do
case lookup(registry, id) do
{:ok, pid} ->
{:ok, pid}
{:error, :not_found} ->
@gen_module.call(registry, {:lookup_or_start, id, args}, timeout)
end
end
def lookup_or_start(registry, id, args, timeout) do
@gen_module.call(registry, {:lookup_or_start, id, args}, timeout)
end
@doc """
Starts a process by id
If the id is already associated with a running process `{:error, {:already_started, pid}}` is
returned.
If the id is not associated with a running process then it is spawned, the optional third
argument will be passed to `start_link` of the `worker_module` to spawn a new process.
"""
@spec start(
registry :: GenServer.server(),
id :: Types.id(),
args :: [any()],
timeout :: integer()
) :: {:ok, pid()} | {:error, {:already_started, pid()}} | {:error, any()}
def start(registry, id, args \\ [], timeout \\ 5_000)
def start(registry, id, args, timeout) when is_atom(registry) do
case lookup(registry, id) do
{:ok, pid} ->
{:error, {:already_started, pid}}
{:error, :not_found} ->
@gen_module.call(registry, {:start, id, args}, timeout)
end
end
def start(registry, id, args, timeout) do
@gen_module.call(registry, {:start, id, args}, timeout)
end
@doc """
Safely stops a process managed by the GenRegistry
In addition to stopping the process, the id is also removed from the GenRegistry
If the id provided is not registered this will return `{:error, :not_found}`
"""
@spec stop(registry :: GenServer.server(), id :: Types.id()) :: :ok | {:error, :not_found}
def stop(registry, id) do
@gen_module.call(registry, {:stop, id})
end
@doc """
Return the number of running processes in this registry.
"""
@spec count(table :: ETS.tab()) :: non_neg_integer()
def count(table) do
ETS.info(table, :size)
end
@doc """
Return a sample entry from the registry.
If the registry is empty returns `nil`.
"""
@spec sample(table :: ETS.tab()) :: {Types.id(), pid()} | nil
def sample(table) do
case ETS.first(table) do
:"$end_of_table" ->
nil
key ->
case ETS.lookup(table, key) do
[entry] ->
entry
_ ->
sample(table)
end
end
end
@doc """
Loop over all the processes and return result.
The function will be called with two arguments, a two-tuple of `{id, pid}` and then accumulator,
the function should return the accumulator.
There is no ordering guarantee when reducing.
"""
@spec reduce(table :: ETS.tab(), acc :: any, ({Types.id(), pid()}, any() -> any())) :: any
def reduce(table, acc, func) do
ETS.foldr(func, acc, table)
end
@doc """
Returns all the entries of the GenRegistry as a list.
There is no ordering guarantee for the list.
"""
@spec to_list(table :: ETS.tab()) :: [{Types.id(), pid()}]
def to_list(table) do
ETS.tab2list(table)
end
## Server Callbacks
def init({worker_module, name}) do
Process.flag(:trap_exit, true)
worker_type =
case worker_module.module_info[:attributes][:behaviour] do
[:supervisor] -> :supervisor
_ -> :worker
end
state = %__MODULE__{
workers:
ETS.new(name, [
:public,
:set,
:named_table,
{:read_concurrency, true}
]),
worker_module: worker_module,
worker_type: worker_type
}
{:ok, state}
end
def terminate(_reason, _state) do
for pid <- Process.get_keys(), is_pid(pid) do
Process.unlink(pid)
Process.exit(pid, :kill)
end
:ok
end
def handle_call({:lookup_or_start, id, args}, _from, state) do
{:reply, do_lookup_or_start(state, id, args), state}
end
def handle_call({:start, id, args}, _from, state) do
{:reply, do_start(state, id, args), state}
end
def handle_call({:stop, id}, _from, state) do
{:reply, do_stop(state, id), state}
end
# Call from supervisor module.
def handle_call(
:which_children,
_from,
%__MODULE__{worker_type: worker_type, worker_module: worker_module} = state
) do
children =
for pid <- Process.get_keys(), is_pid(pid) do
{:undefined, pid, worker_type, [worker_module]}
end
{:reply, children, state}
end
def handle_call(
:count_children,
_from,
%__MODULE__{worker_type: worker_type, worker_module: worker_module} = state
) do
counts = [
specs: 1,
active: count(worker_module),
supervisors: 0,
workers: 0
]
counts =
case worker_type do
:worker -> Keyword.put(counts, :workers, counts[:active])
:supervisor -> Keyword.put(counts, :supervisors, counts[:active])
end
{:reply, counts, state}
end
def handle_call(_message, _from, state) do
{:reply, {:error, __MODULE__}, state}
end
def handle_info({:EXIT, pid, _reason}, %__MODULE__{workers: workers} = state) do
ETS.delete(workers, Process.delete(pid))
{:noreply, state}
end
def handle_info(_message, state) do
{:noreply, state}
end
## Private
@spec do_lookup_or_start(state :: t, id :: Types.id(), args :: [any]) ::
{:ok, pid} | {:error, any}
defp do_lookup_or_start(%__MODULE__{worker_module: worker_module, workers: workers}, id, args) do
case lookup(workers, id) do
{:ok, pid} ->
{:ok, pid}
{:error, :not_found} ->
case apply(worker_module, :start_link, args) do
{:ok, pid} ->
ETS.insert_new(workers, {id, pid})
Process.put(pid, id)
{:ok, pid}
{:error, reason} ->
{:error, reason}
end
end
end
@spec do_start(state :: t(), id :: Types.id(), args :: [any()]) ::
{:ok, pid()} | {:error, {:already_started, pid()}} | {:error, any()}
defp do_start(%__MODULE__{worker_module: worker_module, workers: workers}, id, args) do
case lookup(workers, id) do
{:ok, pid} ->
{:error, {:already_started, pid}}
{:error, :not_found} ->
case apply(worker_module, :start_link, args) do
{:ok, pid} ->
ETS.insert_new(workers, {id, pid})
Process.put(pid, id)
{:ok, pid}
{:error, reason} ->
{:error, reason}
end
end
end
@spec do_stop(state :: t, id :: Types.id()) :: :ok | {:error, :not_found}
defp do_stop(%__MODULE__{workers: workers}, id) do
with {:ok, pid} <- lookup(workers, id) do
Process.unlink(pid)
Process.exit(pid, :shutdown)
ETS.delete(workers, Process.delete(pid))
:ok
end
end
end
