Trap channels

trap.ex

defmodule PushEx.Instrumentation.Tracker do
  use GenServer

  def start_link(_) do
    GenServer.start_link(__MODULE__, [], name: __MODULE__)
  end

  def init(_) do
    # trap_exit to trap non-kill error messages coming from the channel PIDs
    Process.flag(:trap_exit, true)

    {:ok, %{channel_pids: %{}}}
  end

  # Call this when a channel connects
  def track_channel(socket = %Phoenix.Socket{}) do
    GenServer.call(__MODULE__, {:track, socket})
  end

  def state() do
    GenServer.call(__MODULE__, :state)
  end

  ## Callbacks

  def handle_call(:connected_channel_count, _from, state = %{channel_pids: pids}) do
    {:reply, map_size(pids), state}
  end

  def handle_call(
        {:track, socket = %Phoenix.Socket{channel_pid: pid, topic: topic}},
        _from,
        state = %{channel_pids: pids}
      ) do
    link_processes_to_capture_bidirectional_exits(pid)

    new_channel_pids =
      Map.put(pids, pid, %{
        channel: topic
      })

    {:reply, :ok, %{state | channel_pids: new_channel_pids}}
  end

  def handle_call(:state, _from, state) do
    {:reply, state, state}
  end

  # The channel pid just went offline. Could handle some type of callback here
  def handle_info({:EXIT, pid, _reason}, state = %{channel_pids: pids}) do
    new_channel_pids = Map.delete(pids, pid)

    {:noreply, %{state | channel_pids: new_channel_pids}}
  end

  # Link allows exit trapping
  defp link_processes_to_capture_bidirectional_exits(pid), do: Process.link(pid)
end