From 953d32066786d868cffca88ad43f9fa2cbae26c2 Mon Sep 17 00:00:00 2001
From: GilbertWong <gilbertwong96@icloud.com>
Date: Thu, 18 Jul 2019 17:40:25 +0800
Subject: [PATCH] Implement better websocket exit mechanism

---
 src/emqx_session.erl    |  7 +++++++
 src/emqx_ws_channel.erl | 21 ++++++++++++---------
 2 files changed, 19 insertions(+), 9 deletions(-)

diff --git a/src/emqx_session.erl b/src/emqx_session.erl
index 392ce77f5..c5ea950ce 100644
--- a/src/emqx_session.erl
+++ b/src/emqx_session.erl
@@ -331,6 +331,10 @@ update_expiry_interval(SPid, Interval) ->
 close(SPid) ->
     gen_server:call(SPid, close).
 
+-spec(close(spid(), atom()) -> ok).
+close(SPid, Reason) ->
+    gen_server:call(SPid, {close, Reason}).
+
 %%------------------------------------------------------------------------------
 %% gen_server callbacks
 %%------------------------------------------------------------------------------
@@ -457,6 +461,9 @@ handle_call({pubrel, PacketId, _ReasonCode}, _From, State = #state{awaiting_rel
 handle_call(close, _From, State) ->
     {stop, normal, ok, State};
 
+handle_call({close, Reason}, _From, State) ->
+    {stop, Reason, ok, State};
+
 handle_call(Req, _From, State) ->
     ?LOG(error, "Unexpected call: ~p", [Req]),
     {reply, ignored, State}.
diff --git a/src/emqx_ws_channel.erl b/src/emqx_ws_channel.erl
index 9571baab8..6d60b968f 100644
--- a/src/emqx_ws_channel.erl
+++ b/src/emqx_ws_channel.erl
@@ -303,20 +303,23 @@ websocket_info(Info, State) ->
     ?LOG(error, "Unexpected info: ~p", [Info]),
     {ok, State}.
 
-terminate(SockError, _Req, #state{keepalive   = Keepalive,
-                                  proto_state = ProtoState,
-                                  shutdown    = Shutdown}) ->
-    ?LOG(debug, "Terminated for ~p, sockerror: ~p",
-         [Shutdown, SockError]),
+terminate(WsReason, _Req, #state{keepalive   = Keepalive,
+                                 proto_state = ProtoState,
+                                 shutdown    = Shutdown}) ->
+    ?LOG(debug, "Terminated for ~p, websocket reason: ~p",
+         [Shutdown, WsReason]),
     emqx_keepalive:cancel(Keepalive),
     case {ProtoState, Shutdown} of
         {undefined, _} -> ok;
         {_, {shutdown, Reason}} ->
+            SessionPid = emqx_protocol:session(ProtoState),
             emqx_protocol:terminate(Reason, ProtoState),
-            exit(Reason);
-        {_, Error} ->
-            emqx_protocol:terminate(Error, ProtoState),
-            exit({error, SockError})
+            SessionPid ! {'EXIT', self(), Reason};
+        {_, _Error} ->
+            ?LOG(info, "Terminate for unexpected error: ~p", [WsReason]),
+            SessionPid = emqx_protocol:session(ProtoState),
+            emqx_protocol:terminate(unknown, ProtoState),
+            SessionPid ! {'EXIT', self(), unknown}
     end.
 
 %%--------------------------------------------------------------------