Azolo · May 16, 2017 · May 15, 2017 · May 16, 2017 · May 16, 2017 · May 16, 2017
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -1,6 +1,9 @@
 ## unreleased
 - `WebSockex.start_link` will no longer cause the calling process to exit on
   connection failure and will return a proper error tuple instead.
+- Change `WebSockex.Conn.RequestError` to `WebSockex.RequestError`.
+- Add `handle_connect_failure` to be invoked after initiating a connection
+  fails. Fixes #5
 
 ## 0.1.2
 - Rework how disconnects are handled which should improve the

diff --git a/lib/websockex/client.ex b/lib/websockex/client.ex
@@ -53,6 +53,16 @@ defmodule WebSockex.Client do
                         | {:remote, :closed}
                         | {:error, term}
 
+  @typedoc """
+  A map that contains information about the failure to connect.
+
+  This map contains the error, attempt number, and the `t:WebSockex.Conn.t/0`
+  that was used to attempt the connection.
+  """
+  @type connect_failure_map :: %{error: %WebSockex.RequestError{} | %WebSockex.ConnError{},
+                                 attempt_number: integer,
+                                 conn: WebSockex.Conn.t}
+
   @doc """
   Invoked after connection is established.
   """
@@ -119,6 +129,24 @@ defmodule WebSockex.Client do
     | {:close, new_state}
     | {:close, close_frame, new_state} when new_state: term
 
+  @doc """
+  Invoked when there is a failure trying to open the websocket.
+
+  The failure map is the `t:connect_failure_map.t/0` type, and contains the
+  error attempt number and `t:WebSockex.Conn.t/0` used to connect. You can
+  modify the `Conn` struct to change various things about the way you're
+  attmpting to connect.
+
+  - `{:ok, state}` will continue the process termination or error.
+  - `{:reconnect, state}` will attempt to reconnect instead of terminating.
+  - `{:reconnect, conn, state}` will attempt to reconnect with the connection
+    data in `conn`. `conn` is expected to be a `t:WebSockex.Conn.t/0`.
+  """
+  @callback handle_connect_failure(connect_failure_map, state :: term) ::
+    {:ok, new_state}
+    | {:reconnect, new_state}
+    | {:reconnect, WebSockex.Conn.t, new_state} when new_state: term
+
   @doc """
   Invoked when the process is terminating.
   """
@@ -132,8 +160,8 @@ defmodule WebSockex.Client do
     {:ok, new_state :: term}
     | {:error, reason :: term}
 
-  @optional_callbacks [handle_disconnect: 2, handle_ping: 2, handle_pong: 2, terminate: 2,
-                       code_change: 3]
+  @optional_callbacks [handle_disconnect: 2, handle_ping: 2, handle_pong: 2, handle_connect_failure: 2,
+                       terminate: 2, code_change: 3]
 
   defmacro __using__(_) do
     quote location: :keep do
@@ -178,14 +206,18 @@ defmodule WebSockex.Client do
       def handle_pong(:pong, state), do: {:ok, state}
       def handle_pong({:pong, _}, state), do: {:ok, state}
 
+      @doc false
+      def handle_connect_failure(_failure_map, state), do: {:ok, state}
+
       @doc false
       def terminate(_close_reason, _state), do: :ok
 
       @doc false
       def code_change(_old_vsn, state, _extra), do: {:ok, state}
 
       defoverridable [init: 2, handle_frame: 2, handle_cast: 2, handle_info: 2, handle_ping: 2,
-                      handle_pong: 2, handle_disconnect: 2, terminate: 2, code_change: 3]
+                      handle_pong: 2, handle_disconnect: 2, handle_connect_failure: 2,
+                      terminate: 2, code_change: 3]
     end
   end
 
@@ -245,23 +277,9 @@ defmodule WebSockex.Client do
     # OTP stuffs
     debug = :sys.debug_options([])
 
-    with conn <- WebSockex.Conn.new(uri, opts),
-         {:ok, conn} <- open_connection(conn),
-         {:ok, new_module_state} <- module_init(module, module_state, conn) do
-           :proc_lib.init_ack(parent, {:ok, self()})
-            websocket_loop(parent, debug, %{conn: conn,
-                                            module: module,
-                                            module_state: new_module_state,
-                                            buffer: <<>>,
-                                            fragment: nil})
-    else
-      {:error, reason} ->
-        error = Exception.normalize(:error, reason)
-        :proc_lib.init_ack(parent, {:error, error})
-    end
-  rescue
-    exception ->
-      exit({exception, System.stacktrace})
+    conn = WebSockex.Conn.new(uri, opts)
+
+    init_connect(parent, debug, conn, module, module_state, opts)
   end
 
   ## OTP Stuffs
@@ -289,6 +307,49 @@ defmodule WebSockex.Client do
 
   # Internals! Yay
 
+  defp init_connect(parent, debug, conn, module, module_state, opts, attempt \\ 1) do
+    case open_connection(conn) do
+      {:ok, conn} ->
+        module_init(parent, debug, conn, module, module_state, opts)
+      {:error, %{__struct__: struct} = reason}
+      when struct in [WebSockex.ConnError, WebSockex.RequestError] ->
+        case handle_connect_failure(conn, module, module_state, reason, attempt) do
+          {:ok, _} ->
+            :proc_lib.init_ack(parent, {:error, reason})
+          {:reconnect, conn, new_state} ->
+            init_connect(parent, debug, conn, module, new_state, opts, attempt+1)
+        end
+      {:error, reason} ->
+        error = Exception.normalize(:error, reason)
+        :proc_lib.init_ack(parent, {:error, error})
+    end
+  end
+
+  defp reconnect(parent, debug, state, attempt \\ 1) do
+    case open_connection(state.conn) do
+      {:ok, conn} ->
+        websocket_loop(parent, debug, %{state | conn: conn,
+                                                buffer: <<>>})
+      {:error, %{__struct__: struct} = reason}
+      when struct in [WebSockex.ConnError, WebSockex.RequestError] ->
+        case handle_connect_failure(state.conn,
+                                    state.module,
+                                    state.module_state,
+                                    reason,
+                                    attempt) do
+          {:ok, _} ->
+            raise reason
+          {:reconnect, new_conn, new_module_state} ->
+            new_state = %{state | conn: new_conn,
+                                  module_state: new_module_state}
+            reconnect(parent, debug, new_state, attempt+1)
+        end
+      {:error, reason} ->
+        error = Exception.normalize(:error, reason)
+        raise error
+    end
+  end
+
   defp open_connection(conn) do
     with {:ok, conn} <- WebSockex.Conn.open_socket(conn),
          key <- :crypto.strong_rand_bytes(16) |> Base.encode64,
@@ -418,14 +479,7 @@ defmodule WebSockex.Client do
       {:ok, new_state} ->
         terminate(reason, parent, debug, %{state | module_state: new_state})
       {:reconnect, new_state} ->
-        case open_connection(state.conn) do
-          {:ok, conn} ->
-            websocket_loop(parent, debug, %{state | conn: conn,
-                                                    buffer: <<>>,
-                                                    module_state: new_state})
-          {:error, term} ->
-            raise term
-        end
+        reconnect(parent, debug, %{state | module_state: new_state})
       badreply ->
         raise %WebSockex.BadResponseError{module: state.module,
           function: :handle_disconnect, args: [reason, state.module_state],
@@ -436,10 +490,15 @@ defmodule WebSockex.Client do
       terminate({exception, System.stacktrace}, parent, debug, state)
   end
 
-  defp module_init(module, module_state, conn) do
+  defp module_init(parent, debug, conn, module, module_state, _opts) do
     case apply(module, :init, [module_state, conn]) do
-      {:ok, new_state} ->
-        {:ok, new_state}
+      {:ok, new_module_state} ->
+         :proc_lib.init_ack(parent, {:ok, self()})
+          websocket_loop(parent, debug, %{conn: conn,
+                                          module: module,
+                                          module_state: new_module_state,
+                                          buffer: <<>>,
+                                          fragment: nil})
       badreply ->
         raise %WebSockex.BadResponseError{module: module, function: :init,
           args: [module_state, conn], response: badreply}
@@ -458,6 +517,24 @@ defmodule WebSockex.Client do
     end
   end
 
+  defp handle_connect_failure(conn, module, module_state, reason, attempt) do
+    failure_map = %{conn: conn,
+                    error: reason,
+                    attempt_number: attempt}
+
+    case apply(module, :handle_connect_failure, [failure_map, module_state]) do
+      {:ok, new_state} ->
+        {:ok, new_state}
+      {:reconnect, new_state} ->
+        {:reconnect, conn, new_state}
+      {:reconnect, new_conn, new_state} ->
+        {:reconnect, new_conn, new_state}
+      badreply ->
+        raise %WebSockex.BadResponseError{module: module, function: :handle_connect_failure,
+          args: [failure_map, module_state], response: badreply}
+    end
+  end
+
   defp handle_send(binary_frame, parent, debug, %{conn: conn} = state) do
     case WebSockex.Conn.socket_send(conn, binary_frame) do
       :ok ->

diff --git a/lib/websockex/conn.ex b/lib/websockex/conn.ex
@@ -193,7 +193,7 @@ defmodule WebSockex.Conn do
       {:ok, {:http_response, _version, 101, _message}, rest} ->
          decode_headers(rest)
       {:ok, {:http_response, _, code, message}, _} ->
-        {:error, %WebSockex.Conn.RequestError{code: code, message: message}}
+        {:error, %WebSockex.RequestError{code: code, message: message}}
       {:error, error} ->
         {:error, error}
     end

diff --git a/lib/websockex/errors.ex b/lib/websockex/errors.ex
@@ -6,7 +6,7 @@ defmodule WebSockex.ConnError do
   def message(%__MODULE__{original: error}), do: "Connection Error: #{inspect error}"
 end
 
-defmodule WebSockex.Conn.RequestError do
+defmodule WebSockex.RequestError do
   defexception [:code, :message]
 
   def message(%__MODULE__{code: code, message: message}) do

diff --git a/test/support/test_server.ex b/test/support/test_server.ex
@@ -15,12 +15,17 @@ defmodule WebSockex.TestServer do
   def start(pid) when is_pid(pid) do
     ref = make_ref()
     port = get_port()
+    {:ok, agent_pid} = Agent.start_link(fn -> :ok end)
     url = "ws://localhost:#{port}/ws"
-    case Plug.Adapters.Cowboy.http(__MODULE__, [], [dispatch: dispatch(pid), port: port, ref: ref]) do
+
+    opts = [dispatch: dispatch({pid, agent_pid}),
+            port: port,
+            ref: ref]
+
+    case Plug.Adapters.Cowboy.http(__MODULE__, [], opts) do
       {:ok, _} ->
         {:ok, {ref, url}}
       {:error, :eaddrinuse} ->
-        IO.puts "Address #{port} in use!"
         start(pid)
     end
   end
@@ -29,7 +34,9 @@ defmodule WebSockex.TestServer do
     ref = make_ref()
     port = get_port()
     url = "wss://localhost:#{port}/ws"
-    opts = [dispatch: dispatch(pid),
+    {:ok, agent_pid} = Agent.start_link(fn -> :ok end)
+
+    opts = [dispatch: dispatch({pid, agent_pid}),
             certfile: @certfile,
             keyfile: @keyfile,
             port: port,
@@ -51,7 +58,6 @@ defmodule WebSockex.TestServer do
   def receive_socket_pid do
     receive do
       pid when is_pid(pid) -> pid
-      _ -> receive_socket_pid()
     after
       500 -> raise "No Server Socket pid"
     end
@@ -62,8 +68,8 @@ defmodule WebSockex.TestServer do
     [cert]
   end
 
-  defp dispatch(pid) do
-    [{:_, [{"/ws", WebSockex.TestSocket, [pid]}]}]
+  defp dispatch(tuple) do
+    [{:_, [{"/ws", WebSockex.TestSocket, [tuple]}]}]
   end
 
   defp get_port do
@@ -80,13 +86,20 @@ end
 defmodule WebSockex.TestSocket do
   @behaviour :cowboy_websocket_handler
 
-  def init(_, _, _) do
-    {:upgrade, :protocol, :cowboy_websocket}
+  def init(_, req, [{_, agent_pid}]) do
+    case Agent.get(agent_pid, fn x -> x end) do
+      :ok -> {:upgrade, :protocol, :cowboy_websocket}
+      int when is_integer(int) ->
+        :cowboy_req.reply(int, req)
+        {:shutdown, req, :tests_are_fun}
+    end
   end
 
-  def websocket_init(_, req, [pid]) do
+  def terminate(_,_,_), do: :ok
+
+  def websocket_init(_, req, [{pid, agent_pid}]) do
     send(pid, self())
-    {:ok, req, %{pid: pid}}
+    {:ok, req, %{pid: pid, agent_pid: agent_pid}}
   end
 
   def websocket_terminate({:remote, :closed}, _, state) do
@@ -128,5 +141,12 @@ defmodule WebSockex.TestSocket do
   def websocket_info({:send, frame}, req, state) do
     {:reply, frame, req, state}
   end
+  def websocket_info({:set_code, code}, req, state) do
+    Agent.update(state.agent_pid, fn _ -> code end)
+    {:ok, req, state}
+  end
+  def websocket_info(:shutdown, req, state) do
+    {:shutdown, req, state}
+  end
   def websocket_info(_, req, state), do: {:ok, req, state}
 end