From e60b25b8fe397ceb076251166e38ecadd7bdc4e0 Mon Sep 17 00:00:00 2001
From: Luca Succi <luca.succi@stritzinger.com>
Date: Wed, 22 May 2024 17:02:13 +0200
Subject: [PATCH] Fix WS disconnection handling

- Demonitor the gun process before shutting it down through gun.
- Handle disconnections detected by monitor ('DOWN' monitor message)
---
 src/grisp_connect_ws.erl           |  9 +++++++--
 test/grisp_connect_api_SUITE.erl   | 16 ++++++++++++++++
 test/grisp_connect_manager.erl     | 11 +++++++----
 test/grisp_connect_test_client.erl | 20 ++++++++++++++++++--
 4 files changed, 48 insertions(+), 8 deletions(-)

diff --git a/src/grisp_connect_ws.erl b/src/grisp_connect_ws.erl
index 792f240..7e8be38 100644
--- a/src/grisp_connect_ws.erl
+++ b/src/grisp_connect_ws.erl
@@ -91,12 +91,17 @@ handle_info({gun_down, Pid, ws, closed, [Stream]}, #state{gun_pid = Pid, ws_stre
     ?LOG_WARNING(#{event => ws_closed}),
     grisp_connect_client:disconnected(),
     {noreply, shutdown_gun(S)};
+handle_info({'DOWN', _, process, Pid, Reason}, #state{gun_pid = Pid} = S) ->
+    ?LOG_WARNING(#{event => gun_crash, reason => Reason}),
+    grisp_connect_client:disconnected(),
+    {noreply, S?disconnected_state};
 handle_info(M, S) ->
-    ?LOG_WARNING(#{event => unhandled_info, info => M}),
+    ?LOG_WARNING(#{event => unhandled_info, info => M, state => S}),
     {noreply, S}.
 
 % internal functions -----------------------------------------------------------
 
-shutdown_gun(#state{gun_pid = Pid} = State) ->
+shutdown_gun(#state{gun_pid = Pid, gun_ref = GunRef} = State) ->
+    demonitor(GunRef),
     gun:shutdown(Pid),
     State?disconnected_state.
diff --git a/test/grisp_connect_api_SUITE.erl b/test/grisp_connect_api_SUITE.erl
index 66550f8..1f17526 100644
--- a/test/grisp_connect_api_SUITE.erl
+++ b/test/grisp_connect_api_SUITE.erl
@@ -7,6 +7,8 @@
 -compile([export_all, nowarn_export_all]).
 
 -import(grisp_connect_test_client, [wait_connection/0]).
+-import(grisp_connect_test_client, [wait_connection/1]).
+-import(grisp_connect_test_client, [wait_disconnection/0]).
 -import(grisp_connect_test_client, [serial_number/0]).
 -import(grisp_connect_test_client, [cert_dir/0]).
 
@@ -69,6 +71,20 @@ link_device_test(_) ->
     ?assertMatch({ok, <<"ok">>}, grisp_connect:link_device()),
     ?assertMatch({ok, <<"pong">>}, grisp_connect:ping()).
 
+reconnect_on_gun_crash_test(_) ->
+    {state, GunPid, _, _, _} = sys:get_state(grisp_connect_ws),
+    proc_lib:stop(GunPid),
+    ?assertMatch(ok, wait_disconnection()),
+    ?assertMatch(ok, wait_connection()).
+
+reconnect_on_disconnection_test(Config) ->
+    KraftRef = ?config(kraft_instance, Config),
+    kraft:stop(KraftRef),
+    ?assertMatch(ok, wait_disconnection()),
+    KraftRef2 = grisp_connect_manager:kraft_start(cert_dir()),
+    ?assertMatch(ok, wait_connection(100)),
+    [{kraft_instance, KraftRef2} | proplists:delete(kraft_instance, Config)].
+
 %--- Internal ------------------------------------------------------------------
 
 flush() ->
diff --git a/test/grisp_connect_manager.erl b/test/grisp_connect_manager.erl
index 151fd08..8886618 100644
--- a/test/grisp_connect_manager.erl
+++ b/test/grisp_connect_manager.erl
@@ -18,6 +18,12 @@ start(CertDir, Config) ->
     application:start(mnesia),
 
     {ok, Started2} = application:ensure_all_started(kraft),
+    KraftRef = kraft_start(CertDir),
+
+    {ok, Started3} = application:ensure_all_started(grisp_manager),
+    [{apps, Started1 ++ Started2 ++ Started3}, {kraft_instance , KraftRef} | Config].
+
+kraft_start(CertDir) ->
     SslOpts = [
         {verify, verify_peer},
         {keyfile, filename:join(CertDir, "server.key")},
@@ -33,10 +39,7 @@ start(CertDir, Config) ->
         {"/grisp-connect/ws",
             {ws, grisp_manager_device_api}, #{}, #{type => json_rpc}}
     ],
-    kraft:start(KraftOpts, KraftRoutes),
-
-    {ok, Started3} = application:ensure_all_started(grisp_manager),
-    [{apps, Started1 ++ Started2 ++ Started3} | Config].
+    kraft:start(KraftOpts, KraftRoutes).
 
 cleanup_apps(Apps) ->
     mnesia:delete_table(eresu_user),
diff --git a/test/grisp_connect_test_client.erl b/test/grisp_connect_test_client.erl
index 5e3c7ca..fb007b9 100644
--- a/test/grisp_connect_test_client.erl
+++ b/test/grisp_connect_test_client.erl
@@ -6,10 +6,12 @@
 -export([cert_dir/0]).
 -export([serial_number/0]).
 -export([wait_connection/0]).
+-export([wait_connection/1]).
+-export([wait_disconnection/0]).
 
 %--- API -----------------------------------------------------------------------
 
-cert_dir() -> filename:join(code:lib_dir(grisp_connect, test), "certs"). 
+cert_dir() -> filename:join(code:lib_dir(grisp_connect, test), "certs").
 
 serial_number() -> <<"0000">>.
 
@@ -17,7 +19,7 @@ wait_connection() ->
     wait_connection(20).
 
 wait_connection(0) ->
-    ct:pal("grisp_connect state:~n~p~n", [sys:get_state(grisp_connect_client)]),
+    ct:pal("grisp_connect_ws state:~n~p~n", [sys:get_state(grisp_connect_ws)]),
     {error, timeout};
 wait_connection(N) ->
     case grisp_connect:is_connected() of
@@ -26,3 +28,17 @@ wait_connection(N) ->
            ct:sleep(100),
            wait_connection(N - 1)
     end.
+
+wait_disconnection() ->
+    wait_disconnection(20).
+
+wait_disconnection(0) ->
+    ct:pal("grisp_connect_ws state:~n~p~n", [sys:get_state(grisp_connect_ws)]),
+    {error, timeout};
+wait_disconnection(N) ->
+    case grisp_connect:is_connected() of
+        true ->
+            ct:sleep(100),
+            wait_disconnection(N - 1);
+        false -> ok
+    end.