Migrate branch bug18732

3 files changed, 41 insertions, 11 deletions

diff --git a/src/rabbit_misc.erl b/src/rabbit_misc.erl
index 927d7712..b71aba42 100644
--- a/src/rabbit_misc.erl
+++ b/src/rabbit_misc.erl
@@ -35,7 +35,7 @@
 -export([r/3, r/2, rs/1]).
 -export([permission_list/1]).
 -export([enable_cover/0, report_cover/0]).
--export([with_exit_handler/2]).
+-export([throw_on_error/2, with_exit_handler/2]).
 -export([with_user/2, with_vhost/2, with_realm/2, with_user_and_vhost/3]).
 -export([execute_mnesia_transaction/1]).
 -export([ensure_ok/2]).
@@ -79,6 +79,8 @@
 -spec(permission_list/1 :: (ticket()) -> [permission()]).
 -spec(enable_cover/0 :: () -> 'ok' | {'error', any()}).
 -spec(report_cover/0 :: () -> 'ok').
+-spec(throw_on_error/2 ::
+      (atom(), thunk({error, any()} | {ok, A} | A)) -> A). 
 -spec(with_exit_handler/2 :: (thunk(A), thunk(A)) -> A). 
 -spec(with_user/2 :: (username(), thunk(A)) -> A). 
 -spec(with_vhost/2 :: (vhost(), thunk(A)) -> A).
@@ -231,6 +233,13 @@ report_coverage_percentage(File, Cov, NotCov, Mod) ->
                end,
                Mod]).
 
+throw_on_error(E, Thunk) ->
+    case Thunk() of
+        {error, Reason} -> throw({E, Reason});
+        {ok, Res}       -> Res;
+        Res             -> Res
+    end.
+
 with_exit_handler(Handler, Thunk) ->
     try
         Thunk()
diff --git a/src/rabbit_reader.erl b/src/rabbit_reader.erl
index 1d11cbaa..38349a1c 100644
--- a/src/rabbit_reader.erl
+++ b/src/rabbit_reader.erl
@@ -157,14 +157,27 @@ teardown_profiling(Value) ->
             fprof:analyse([{dest, []}, {cols, 100}])
     end.
 
+inet_op(F) -> rabbit_misc:throw_on_error(inet_error, F).
+
+peername(Sock) ->   
+    try
+        {Address, Port} = inet_op(fun () -> inet:peername(Sock) end),
+        AddressS = inet_parse:ntoa(Address),
+        {AddressS, Port}
+    catch
+        Ex -> rabbit_log:error("error on TCP connection ~p:~p~n",
+                               [self(), Ex]),
+              rabbit_log:info("closing TCP connection ~p", [self()]),
+              exit(normal)
+    end.
+
 start_connection(Parent, Deb, ClientSock) ->
-    ProfilingValue = setup_profiling(),
     process_flag(trap_exit, true),
-    {ok, {PeerAddress, PeerPort}} = inet:peername(ClientSock),
-    PeerAddressS = inet_parse:ntoa(PeerAddress),
-    rabbit_log:info("starting TCP connection ~p from ~s:~p~n",
-                    [self(), PeerAddressS, PeerPort]),
+    {PeerAddressS, PeerPort} = peername(ClientSock),
+    ProfilingValue = setup_profiling(),
     try 
+        rabbit_log:info("starting TCP connection ~p from ~s:~p~n",
+                        [self(), PeerAddressS, PeerPort]),
         erlang:send_after(?HANDSHAKE_TIMEOUT * 1000, self(),
                           handshake_timeout),
         mainloop(Parent, Deb, switch_callback(
@@ -255,7 +268,8 @@ mainloop(Parent, Deb, State = #v1{sock= Sock, recv_ref = Ref}) ->
     end.
 
 switch_callback(OldState, NewCallback, Length) ->
-    {ok, Ref} = prim_inet:async_recv(OldState#v1.sock, Length, -1),
+    Ref = inet_op(fun () -> prim_inet:async_recv(
+                              OldState#v1.sock, Length, -1) end),
     OldState#v1{callback = NewCallback,
                 recv_ref = Ref}.
 
@@ -470,7 +484,10 @@ handle_input(handshake, <<"AMQP",1,1,ProtocolMajor,ProtocolMinor>>,
     end;
 
 handle_input(handshake, Other, #v1{sock = Sock}) ->
-    ok = gen_tcp:send(Sock, <<"AMQP",1,1,?PROTOCOL_VERSION_MAJOR,?PROTOCOL_VERSION_MINOR>>),
+    ok = inet_op(fun () -> gen_tcp:send(
+                             Sock, <<"AMQP",1,1,
+                                    ?PROTOCOL_VERSION_MAJOR,
+                                    ?PROTOCOL_VERSION_MINOR>>) end),
     throw({bad_header, Other});
 
 handle_input(Callback, Data, _State) ->
diff --git a/src/rabbit_writer.erl b/src/rabbit_writer.erl
index eda871ec..c3c7db53 100644
--- a/src/rabbit_writer.erl
+++ b/src/rabbit_writer.erl
@@ -127,12 +127,16 @@ assemble_frames(Channel, MethodRecord, Content, FrameMax) ->
                       Channel, Content, FrameMax),
     [MethodFrame | ContentFrames].
 
+tcp_send(Sock, Data) ->
+    rabbit_misc:throw_on_error(inet_error,
+                               fun () -> gen_tcp:send(Sock, Data) end).
+
 internal_send_command(Sock, Channel, MethodRecord) ->
-    ok = gen_tcp:send(Sock, assemble_frames(Channel, MethodRecord)).
+    ok = tcp_send(Sock, assemble_frames(Channel, MethodRecord)).
 
 internal_send_command(Sock, Channel, MethodRecord, Content, FrameMax) ->
-    ok = gen_tcp:send(Sock, assemble_frames(Channel, MethodRecord,
-                                            Content, FrameMax)).
+    ok = tcp_send(Sock, assemble_frames(Channel, MethodRecord,
+                                        Content, FrameMax)).
 
 %% gen_tcp:send/2 does a selective receive of {inet_reply, Sock,
 %% Status} to obtain the result. That is bad when it is called from