dotfiles/app/lib/prymn/agents/connection.ex

defmodule Prymn.Agents.Connection do
  @moduledoc false

  require Logger

  use GenServer, restart: :transient

  @timeout :timer.minutes(2)

  @spec start_link(String.t()) :: GenServer.on_start()
  def start_link(addr) when is_binary(addr) do
    GenServer.start_link(__MODULE__, addr, name: via(addr))
  end

  @spec keep_alive(pid | String.t()) :: :ok
  def keep_alive(server) when is_pid(server), do: GenServer.cast(server, :keep_alive)
  def keep_alive(server) when is_binary(server), do: GenServer.cast(via(server), :keep_alive)

  @spec stop(pid | String.t()) :: :ok
  def stop(server) when is_pid(server), do: GenServer.stop(server, :shutdown)
  def stop(server) when is_binary(server), do: GenServer.stop(via(server), :shutdown)

  @spec get_channel(String.t()) :: GRPC.Channel.t()
  def get_channel(server) do
    GenServer.call(via(server), :get_channel)
  end

  @impl true
  def init(host) do
    Process.flag(:trap_exit, true)
    {:ok, {host, nil}, {:continue, :connect}}
  end

  @impl true
  def handle_continue(:connect, {host_address, _} = state) do
    case GRPC.Stub.connect(host_address, 50012, []) do
      {:ok, channel} ->
        keep_alive(self())
        {:noreply, {host_address, channel}, {:continue, :health}}

      {:error, error} ->
        {:stop, error, state}
    end
  end

  @impl true
  def handle_continue(:health, {_, channel} = state) do
    pid = self()

    Task.start_link(fn ->
      {:ok, stream} = PrymnProto.Prymn.Agent.Stub.health(channel, %Google.Protobuf.Empty{})

      # Read from the stream forever and send data back to parent
      stream
      |> Stream.each(fn {_, data} -> send(pid, data) end)
      |> Enum.take_while(fn _ -> true end)
    end)

    {:noreply, state, @timeout}
  end

  @impl true
  def handle_call(:get_channel, _from, {_, channel} = state) do
    {:reply, channel, state, @timeout}
  end

  @impl true
  def handle_cast(:keep_alive, state) do
    {:noreply, state, @timeout}
  end

  @impl true
  def handle_info(%PrymnProto.Prymn.HealthResponse{} = response, {host, _} = state) do
    response
    |> Prymn.Agents.Health.new_from_proto()
    |> Prymn.Agents.Health.update(host)

    {:noreply, state, @timeout}
  end

  def handle_info(%GRPC.RPCError{} = response, {host, _} = state) do
    Logger.debug("received a GRPC error: #{inspect(response)}")

    response
    |> Prymn.Agents.Health.new_from_proto()
    |> Prymn.Agents.Health.update(host)

    {:noreply, state, @timeout}
  end

  def handle_info({:gun_up, _pid, _protocol}, state) do
    # TODO: If it's possible for the GRPC connection to be down when we receive
    # this message, we should `{:continue, state.channel.host}`
    {:noreply, state, {:continue, :health}}
  end

  def handle_info({:gun_down, _pid, _proto, _reason, _}, {host, _} = state) do
    Logger.debug("disconnected from #{host}")
    {:noreply, state, @timeout}
  end

  def handle_info(:timeout, state) do
    {:stop, {:shutdown, :timeout}, state}
  end

  def handle_info({:EXIT, _from, _reason}, state) do
    # TODO: How to handle this... (this happens when a linked process exists
    # e.g the one with the health stream)
    {:noreply, state, @timeout}
  end

  def handle_info(msg, state) do
    Logger.warning("received unexpected message: #{inspect(msg)}")
    {:noreply, state, @timeout}
  end

  @impl true
  def terminate(reason, {host, channel}) do
    Logger.debug("terminating Agent connection (host: #{host}, reason: #{inspect(reason)})")
    if channel, do: GRPC.Stub.disconnect(channel)
  end

  defp via(name) do
    {:via, Registry, {Prymn.Agents.Registry, name}}
  end
end