Trap and handle exit in channel and ws_channel

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}.

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}.