fix: retry failed connections and stop gracefully on error

Author: SpaceEEC

lib/gateway/connection.ex

@@ -387,7 +387,8 @@ defmodule Crux.Gateway.Connection do
     end
   end
 
-  def terminate(:fatal_close_code, _state, data) do
+  def terminate(reason, _state, data)
+      when reason in [:fatal_close_code, :connection_failure] do
     Gateway.stop_shard(data.name, {data.shard_id, data.shard_count})
   end
 
@@ -472,6 +473,12 @@ defmodule Crux.Gateway.Connection do
     {:next_state, @disconnected, data, :postpone}
   end
 
+  defp handle_common(_state, :info, {:EXIT, conn, :connection_failure}, %{conn: conn}) do
+    Logger.error(fn -> "Connection failed to establish connection, shutting down." end)
+
+    {:stop, :connection_failure}
+  end
+
   defp handle_packet(packet, data)
 
   defp handle_packet(%{op: @heartbeat}, %{conn: conn, seq: seq} = data) do

lib/gateway/connection/gun.ex

@@ -84,6 +84,8 @@ defmodule Crux.Gateway.Connection.Gun do
   @connecting :connecting
   @connected :connected
 
+  @attempt_limit 5
+
   defstruct [
     :parent,
     :host,
@@ -93,6 +95,7 @@ defmodule Crux.Gateway.Connection.Gun do
     :zlib,
     :buffer,
     :conn,
+    :attempts,
     :expect_disconnect
   ]
 
@@ -109,6 +112,8 @@ defmodule Crux.Gateway.Connection.Gun do
            buffer: binary(),
            # WS connection wrapper process
            conn: pid() | nil,
+           # Limit the amount of attempts to establish a connection
+           attempts: non_neg_integer(),
            # Whether we are expecting a gun_down / disconnect
            # and do not want to notify the spawning process again
            expect_disconnect: boolean()
@@ -136,6 +141,7 @@ defmodule Crux.Gateway.Connection.Gun do
           zlib: nil,
           buffer: <<>>,
           conn: nil,
+          attempts: 0,
           expect_disconnect: false
         }
 
@@ -184,30 +190,54 @@ defmodule Crux.Gateway.Connection.Gun do
     z = Zlib.open()
     Zlib.inflateInit(z)
 
-    Logger.debug(fn -> "Starting a process to connect to #{data.host}:#{data.port}" end)
+    attempts = data.attempts + 1
+    data = %{data | attempts: attempts}
+
+    Logger.debug(fn ->
+      "Starting a process to connect to #{data.host}:#{data.port} (Attempt: #{attempts} / #{@attempt_limit})"
+    end)
 
     # > Gun does not currently support Websocket over HTTP/2.
     {:ok, conn} = Gun.open(data.host, data.port, %{protocols: [:http]})
 
     Logger.debug(fn -> "Process started, waiting for its connection to be up." end)
 
-    {:ok, :http} = Gun.await_up(conn)
+    conn
+    |> Gun.await_up()
+    |> case do
+      {:ok, :http} ->
+        Logger.debug(fn ->
+          "Connection is up, now upgrading it to use the WebSocket protocol, using " <>
+            data.path <> data.query
+        end)
 
-    Logger.debug(fn ->
-      "Connection is up, now upgrading it to use the WebSocket protocol, using " <>
-        data.path <> data.query
-    end)
+        stream_ref = Gun.ws_upgrade(conn, data.path <> data.query)
+        :ok = await_upgrade(conn, stream_ref)
+
+        Logger.debug(fn ->
+          "Connection upgraded to use the WebSocket protocol, we are good to go!"
+        end)
+
+        send_connected(data)
 
-    stream_ref = Gun.ws_upgrade(conn, data.path <> data.query)
-    :ok = await_upgrade(conn, stream_ref)
+        data = %{data | conn: conn, zlib: z, attempts: 0}
 
-    Logger.debug(fn -> "Connection upgraded to use the WebSocket protocol, we are good to go!" end)
+        {:keep_state, data, {:timeout, 0, :connected}}
 
-    send_connected(data)
+      {:error, :timeout} when attempts >= @attempt_limit ->
+        Logger.error(fn ->
+          "Connection timed out, no attempts remaining, won't retry. (#{attempts} / #{@attempt_limit})"
+        end)
 
-    data = %{data | conn: conn, zlib: z}
+        {:stop, :connection_failure, data}
 
-    {:keep_state, data, {:timeout, 0, :connected}}
+      {:error, :timeout} ->
+        Logger.warn(fn ->
+          "Connection timed out, will retry. (#{attempts} / #{@attempt_limit})"
+        end)
+
+        {:repeat_state, data}
+    end
   end
 
   def connecting(:timeout, :connected, data) do