Clean up gateway lifecycle processes

lib/bot/discord_bot.ex

@@ -14,6 +14,10 @@ defmodule Lanyard.DiscordBot do
   end
 
   def init(state) do
+    # The gateway is both linked and monitored; trap exits so the :DOWN restart
+    # path below can run instead of losing the bot process first.
+    Process.flag(:trap_exit, true)
+
     {:ok, %__MODULE__{token: state.token, gateway_client_pid: nil}, {:continue, :setup_bot}}
   end
 
@@ -42,9 +46,21 @@ defmodule Lanyard.DiscordBot do
     {:noreply, %{state | resume_data: nil}}
   end
 
-  def handle_info({:DOWN, _ref, :process, _pid, reason}, state) do
-    Logger.warning("Discord bot crashed with reason: #{reason}. Restarting.")
+  def handle_info({:EXIT, pid, _reason}, %{gateway_client_pid: pid} = state) do
+    {:noreply, state}
+  end
+
+  def handle_info({:EXIT, _pid, _reason}, state) do
+    {:noreply, state}
+  end
+
+  def handle_info({:DOWN, _ref, :process, pid, reason}, %{gateway_client_pid: pid} = state) do
+    Logger.warning("Discord bot crashed with reason: #{inspect(reason)}. Restarting.")
 
     {:noreply, state, {:continue, :setup_bot}}
   end
+
+  def handle_info({:DOWN, _ref, :process, _pid, _reason}, state) do
+    {:noreply, state}
+  end
 end

lib/gateway/client.ex

@@ -94,7 +94,7 @@ defmodule Lanyard.Gateway.Client do
       )
     end
 
-    {:close, reason, state}
+    close_gateway(reason, state)
   end
 
   def websocket_handle({:text, payload}, _socket, state) do
@@ -122,6 +122,8 @@ defmodule Lanyard.Gateway.Client do
     # Start sending heartbeat with interval defined by the hello packet
     Logger.debug("Discord: Hello")
 
+    stop_process(state[:heartbeat_pid])
+
     {:ok, heartbeat_pid} =
       Heartbeat.start_link(
         state[:agent_seq_num],
@@ -169,13 +171,13 @@ defmodule Lanyard.Gateway.Client do
        }}
     )
 
-    {:close, "Reconnecting for resume", state}
+    close_gateway("Reconnecting for resume", state)
   end
 
   defp _handle_data(%{op: :invalid_session} = _data, state) do
     Logger.warning("Discord: Invalid session, starting new session")
     send(:discord_bot, :clear_resume)
-    {:close, "Invalid session, starting new session", state}
+    close_gateway("Invalid session, starting new session", state)
   end
 
   def websocket_info(:start, _connection, state) do
@@ -220,12 +222,16 @@ defmodule Lanyard.Gateway.Client do
        }}
     )
 
-    {:close, "Heartbeat stale", state}
+    close_gateway("Heartbeat stale", state)
   end
 
   @spec websocket_terminate(any(), any(), nil | keyword() | map()) :: :ok
   def websocket_terminate(reason, _conn_state, state) do
-    Logger.info("Discord: Websocket closed in state #{inspect(state)} with reason #{inspect(reason)}")
+    Logger.info(
+      "Discord: Websocket closed in state #{inspect(state)} with reason #{inspect(reason)}"
+    )
+
+    cleanup_transient_processes(state)
 
     :ok
   end
@@ -371,6 +377,35 @@ defmodule Lanyard.Gateway.Client do
     end
   end
 
+  defp close_gateway(reason, state) do
+    {:close, reason, cleanup_transient_processes(state)}
+  end
+
+  defp cleanup_transient_processes(state) when is_map(state) do
+    # Reconnect/close paths can exit normally, which won't take linked helpers
+    # down with the gateway process.
+    stop_process(state[:heartbeat_pid])
+    stop_process(state[:agent_seq_num])
+
+    state
+    |> Map.put(:heartbeat_pid, nil)
+    |> Map.put(:agent_seq_num, nil)
+  end
+
+  defp cleanup_transient_processes(state), do: state
+
+  defp stop_process(pid) when is_pid(pid) do
+    if Process.alive?(pid) do
+      try do
+        GenServer.stop(pid, :normal, 1_000)
+      catch
+        :exit, _ -> :ok
+      end
+    end
+  end
+
+  defp stop_process(_pid), do: :ok
+
   defp create_member_presences(payload) do
     Task.start(fn ->
       Enum.each(payload.data["members"], fn member ->

lib/gateway/heartbeat.ex

@@ -21,6 +21,9 @@ defmodule Lanyard.Gateway.Heartbeat do
       agent_seq_num: agent_seq_num,
       interval: interval,
       socket_pid: socket_pid,
+      # Normal gateway exits do not stop linked processes, so tie heartbeat
+      # shutdown to the owning socket process explicitly.
+      socket_ref: Process.monitor(socket_pid),
       timer: nil,
       ack?: true
     }
@@ -42,6 +45,13 @@ defmodule Lanyard.Gateway.Heartbeat do
     {:noreply, %{state | timer: nil}}
   end
 
+  def handle_info(
+        {:DOWN, ref, :process, pid, _reason},
+        %{socket_pid: pid, socket_ref: ref} = state
+      ) do
+    {:stop, :normal, state}
+  end
+
   def handle_call(:ack, _from, state) do
     {:reply, :ok, %{state | ack?: true}}
   end