From 0c54d899dab1dfca7771a768068a497c1d01be45 Mon Sep 17 00:00:00 2001
From: terry-xiaoyu <506895667@qq.com>
Date: Sat, 20 Jul 2019 13:19:13 +0800
Subject: [PATCH] Trap and handle exit in channel and ws_channel

---
 src/emqx_channel.erl    | 9 +++++++++
 src/emqx_ws_channel.erl | 8 ++++++++
 2 files changed, 17 insertions(+)

diff --git a/src/emqx_channel.erl b/src/emqx_channel.erl
index 65d34d740..0aeedcdfa 100644
--- a/src/emqx_channel.erl
+++ b/src/emqx_channel.erl
@@ -148,6 +148,7 @@ call(CPid, Req) ->
 %%--------------------------------------------------------------------
 
 init({Transport, RawSocket, Options}) ->
+    process_flag(trap_exit, true),
     {ok, Socket} = Transport:wait(RawSocket),
     {ok, Peername} = Transport:ensure_ok_or_exit(peername, [Socket]),
     {ok, Sockname} = Transport:ensure_ok_or_exit(sockname, [Socket]),
@@ -365,6 +366,14 @@ handle(info, {shutdown, conflict, {ClientId, NewPid}}, State) ->
 handle(info, {shutdown, Reason}, State) ->
     shutdown(Reason, State);
 
+handle(info, Info = {'EXIT', SessionPid, Reason}, State = #state{proto_state = ProtoState}) ->
+    case emqx_protocol:session(ProtoState) of
+        undefined ->
+            ?LOG(error, "Unexpected EXIT: ~p", [Info]),
+            {keep_state, State};
+        SessionPid -> shutdown(Reason, State)
+    end;
+
 handle(info, Info, State) ->
     ?LOG(error, "Unexpected info: ~p", [Info]),
     {keep_state, State}.
diff --git a/src/emqx_ws_channel.erl b/src/emqx_ws_channel.erl
index 74a5e2a71..c308b687e 100644
--- a/src/emqx_ws_channel.erl
+++ b/src/emqx_ws_channel.erl
@@ -299,6 +299,14 @@ websocket_info({shutdown, Reason}, State) ->
 websocket_info({stop, Reason}, State) ->
     {stop, State#state{shutdown = Reason}};
 
+websocket_info(Info = {'EXIT', SessionPid, Reason}, State = #state{proto_state = ProtoState}) ->
+    case emqx_protocol:session(ProtoState) of
+        undefined ->
+            ?LOG(error, "Unexpected EXIT: ~p", [Info]),
+            {ok, State};
+        SessionPid -> shutdown(Reason, State)
+    end;
+
 websocket_info(Info, State) ->
     ?LOG(error, "Unexpected info: ~p", [Info]),
     {ok, State}.