Incresed resilience to crashes, and error reporting with on_error function

src/wsecli.erl

@@ -109,11 +109,12 @@ start(URI, Options) ->
   Port     :: inet:port_number(),
   Resource :: string(),
   Options  :: list({atom(), any()})
-) -> {ok, pid()}.
+) -> {ok, pid()} | {error, {shutdown, any()}} | {error, any()}.
 start(Host, Port, Path, Options) ->
   UseSSL        = proplists:get_value(ssl, Options, false),
-  GenOptions    = [{timeout, 5000}],
-  ClientOptions = {Host, Port, Path, UseSSL},
+  Timeout       = proplists:get_value(timeout, Options, 5000),
+  GenOptions    = [{timeout, Timeout}],
+  ClientOptions = {Host, Port, Path, UseSSL, Timeout},
   case proplists:get_value(register, Options, false) of
       false ->
         gen_fsm:start_link(?MODULE, ClientOptions, GenOptions);
@@ -262,16 +263,20 @@ on_close(Client, Callback) ->
     Resource :: string(),
     SSL      :: boolean()
   }
-  ) -> {ok, connecting, #data{}}.
-init({Host, Port, Resource, SSL}) ->
+  ) -> {ok, connecting, #data{}} | {stop, {shutdown, any()}}.
+init({Host, Port, Resource, SSL, Timeout}) ->
   SocketType = case SSL of true -> ssl; false -> plain end,
-  {ok, Socket}    = wsecli_socket:open(Host, Port, SocketType, self()),
-  {ok, Handshake} = wsock_handshake:open(Resource, Host, Port),
-  Request         = wsock_http:encode(Handshake#handshake.message),
+  case wsecli_socket:open(Host, Port, SocketType, Timeout, self()) of
+    {ok, Socket} ->
+      {ok, Handshake} = wsock_handshake:open(Resource, Host, Port),
+      Request         = wsock_http:encode(Handshake#handshake.message),
 
-  wsecli_socket:send(Request, Socket),
+      wsecli_socket:send(Request, Socket),
 
-  {ok, connecting, #data{ socket = Socket, handshake = Handshake }}.
+      {ok, connecting, #data{ socket = Socket, handshake = Handshake }};
+    {error, Error} ->
+      {stop, {shutdown, Error}}
+  end.
 
 
 %% @hidden
@@ -366,8 +371,9 @@ handle_info({socket, {data, Data}}, connecting, StateData) ->
     {ok, _Handshake} ->
       spawn(StateData#data.cb#callbacks.on_open),
       {next_state, open, StateData};
-    {error, _Error} ->
-      {stop, failed_handshake, StateData}
+    {error, Error} ->
+      (StateData#data.cb#callbacks.on_error)(Error),
+      {stop, {shutdown, failed_handshake}, StateData}
   end;
 handle_info({socket, {data, Data}}, open, StateData) ->
   {Messages, State} = case StateData#data.fragmented_message of
@@ -389,7 +395,10 @@ handle_info({socket, {data, Data}}, closing, StateData) ->
       {next_state, closing, StateData}
   end;
 handle_info({socket, close}, _StateName, StateData) ->
-  {stop, normal, StateData}.
+  {stop, normal, StateData};
+handle_info({socket, {error, Error}}, StateName, StateData) ->
+  (StateData#data.cb#callbacks.on_error)(Error),
+  {next_state, StateName, StateData}.
 
 %% @hidden
 -spec terminate(

src/wsecli_socket.erl

@@ -2,7 +2,7 @@
 -module(wsecli_socket).
 -include("wsecli.hrl").
 
--export([open/4]).
+-export([open/5]).
 -export([send/2]).
 -export([close/1]).
 
@@ -18,7 +18,8 @@
 %%========================================
 %% Constants
 %%========================================
--define(DEFAULT_SOCKET_OPTIONS, [binary, {reuseaddr, true}, {packet, raw}]).
+-define(DEFAULT_SOCKET_OPTIONS(Timeout), [{reuseaddr, true}, {packet, raw},
+                                         binary, {send_timeout, Timeout}]).
 
 %%========================================
 %% Client API
@@ -27,14 +28,17 @@
   Host    :: string(),
   Port    :: inet:port_number(),
   Type    :: socket_type(),
+  Timeout :: timeout(),
   Client  :: pid()
   ) ->
   {ok, socket()} |
   {error, term()}.
-open(Host, Port, plain, Client) ->
-  wsecli_socket_plain:start_link(Host, Port, Client, ?DEFAULT_SOCKET_OPTIONS);
-open(Host, Port, ssl, Client) ->
-  wsecli_socket_ssl:start_link(Host, Port, Client, ?DEFAULT_SOCKET_OPTIONS).
+open(Host, Port, plain, Timeout, Client) ->
+  wsecli_socket_plain:start_link(Host, Port, Client,
+                                 ?DEFAULT_SOCKET_OPTIONS(Timeout), Timeout);
+open(Host, Port, ssl, Timeout, Client) ->
+  wsecli_socket_ssl:start_link(Host, Port, Client,
+                              ?DEFAULT_SOCKET_OPTIONS(Timeout), Timeout).
 
 -spec send(
   Data   :: iolist(),

src/wsecli_socket_plain.erl

@@ -2,8 +2,8 @@
 -module(wsecli_socket_plain).
 -include("wsecli.hrl").
 
--export([start_link/4]).
--export([init/4]).
+-export([start_link/5]).
+-export([init/5]).
 -export([loop/1]).
 
 
@@ -22,30 +22,41 @@
   Host    :: string(),
   Port    :: inet:port_number(),
   Client  :: pid(),
-  Options :: list(gen_tcp:connect_option())
+  Options :: list(gen_tcp:connect_option()),
+  Timeout :: timeout()
 ) -> {ok, socket()} | {error, term()}.
-start_link(Host, Port, Client, Options) ->
-  proc_lib:start_link(?MODULE, init, [Host, Port, Client, Options]).
+start_link(Host, Port, Client, Options, Timeout) ->
+  proc_lib:start_link(?MODULE, init, [Host, Port, Client, Options, Timeout]).
 
 -spec init(
   Host    :: string(),
   Port    :: inet:port_number(),
   Client  :: pid(),
-  Options :: list(gen_tcp:connect_option())
+  Options :: list(gen_tcp:connect_option()),
+  Timeout :: timeout()
   ) -> ok.
-init(Host, Port, Client, Options) ->
-  {ok, Socket} = gen_tcp:connect(Host, Port, Options),
-  State        = #state{ client = Client, socket = Socket },
-  proc_lib:init_ack({ok, self()}),
-  loop(State).
+init(Host, Port, Client, Options, Timeout) ->
+  case gen_tcp:connect(Host, Port, Options, Timeout) of
+    {ok, Socket} ->
+      State        = #state{ client = Client, socket = Socket },
+      proc_lib:init_ack({ok, self()}),
+      loop(State);
+    {error, _} = Error ->
+      proc_lib:init_ack(Error)
+  end.
 
 %%========================================
 %% Internal
 %%========================================
 loop(State) ->
   receive
     {socket, send, Data}   ->
-      gen_tcp:send(State#state.socket, Data),
+      case gen_tcp:send(State#state.socket, Data) of
+        ok ->
+          ok;
+        {error, _} = Error ->
+          wsecli_socket:notify_client(Error, State#state.client)
+      end,
       loop(State);
     {socket, close}        ->
       gen_tcp:close(State#state.socket);

src/wsecli_socket_ssl.erl

@@ -2,8 +2,8 @@
 -module(wsecli_socket_ssl).
 -include("wsecli.hrl").
 
--export([start_link/4]).
--export([init/4]).
+-export([start_link/5]).
+-export([init/5]).
 -export([loop/1]).
 
 
@@ -22,30 +22,41 @@
   Host    :: string(),
   Port    :: inet:port_number(),
   Client  :: pid(),
-  Options :: list(ssl:connect_option())
+  Options :: list(ssl:connect_option()),
+  Timeout :: timeout()
 ) -> {ok, socket()}.
-start_link(Host, Port, Client, Options) ->
-  proc_lib:start_link(?MODULE, init, [Host, Port, Client, Options]).
+start_link(Host, Port, Client, Options, Timeout) ->
+  proc_lib:start_link(?MODULE, init, [Host, Port, Client, Options, Timeout]).
 
 -spec init(
   Host    :: string(),
   Port    :: inet:port_number(),
   Client  :: pid(),
-  Options :: list(ssl:connect_option())
+  Options :: list(ssl:connect_option()),
+  Timeout :: timeout()
   ) -> ok.
-init(Host, Port, Client, Options) ->
-  {ok, Socket} = ssl:connect(Host, Port, Options),
-  State        = #state{ client = Client, socket = Socket },
-  proc_lib:init_ack({ok, self()}),
-  loop(State).
+init(Host, Port, Client, Options, Timeout) ->
+  case ssl:connect(Host, Port, Options, Timeout) of
+    {ok, Socket} ->
+      State        = #state{ client = Client, socket = Socket },
+      proc_lib:init_ack({ok, self()}),
+      loop(State);
+    {error, _} = Error ->
+      proc_lib:init_ack(Error)
+  end.
 
 %%========================================
 %% Internal
 %%========================================
 loop(State) ->
   receive
     {socket, send, Data}   ->
-      ssl:send(State#state.socket, Data),
+      case ssl:send(State#state.socket, Data) of
+        ok ->
+          ok;
+        {error, _} = Error ->
+          wsecli_socket:notify_client(Error, State#state.client)
+      end,
       loop(State);
     {socket, close}        ->
       ssl:close(State#state.socket);