From 475db5ae306e1f608018200a271914cf8e517e5b Mon Sep 17 00:00:00 2001
From: Yura Shelyag <yuriy.shelyag@gmail.com>
Date: Fri, 1 Nov 2024 13:33:55 +0100
Subject: [PATCH] Fixed die and SIGINT exit

---
 hbs2-peer/app/PeerMain.hs                     | 27 ++++++++++++++-----
 hbs2-peer/app/RPC2/Die.hs                     |  3 +--
 hbs2-peer/lib/HBS2/Peer/RPC/Internal/Types.hs |  3 ++-
 3 files changed, 23 insertions(+), 10 deletions(-)

diff --git a/hbs2-peer/app/PeerMain.hs b/hbs2-peer/app/PeerMain.hs
index 3c2ff57c..5201fabb 100644
--- a/hbs2-peer/app/PeerMain.hs
+++ b/hbs2-peer/app/PeerMain.hs
@@ -126,6 +126,8 @@ import Streaming.Prelude qualified as S
 import Graphics.Vty qualified as Vty
 import Graphics.Vty.Platform.Unix qualified as Vty
 
+import Control.Concurrent.Async (ExceptionInLinkedThread(..))
+
 data GoAgainException = GoAgainException
                         deriving (Eq,Ord,Show,Typeable)
 
@@ -670,6 +672,20 @@ respawn opts =
     print (self, args)
     executeFile self False args Nothing
 
+respawnOnError :: PeerOpts -> IO () -> IO ()
+respawnOnError opts act =
+  Exception.try act >>= match
+  where match (Right r) = return r
+        match (Left e)
+          | Just (ExceptionInLinkedThread _ e') <- Exception.fromException e = do
+             match (Left e')
+          | Just ec <- Exception.fromException @ExitCode e =
+            notice $ viaShow ec
+          | Just UserInterrupt <- Exception.fromException e =
+            notice "Interrupted by user"
+          | otherwise =
+            myException e >> performGC >> respawn opts
+
 runPeer :: forall e s . ( e ~ L4Proto
                         , FromStringMaybe (PeerAddr e)
                         , s ~ Encryption e
@@ -678,10 +694,7 @@ runPeer :: forall e s . ( e ~ L4Proto
                         , HasStorage (PeerM e IO)
                         )=> PeerOpts -> IO ()
 
-runPeer opts = Exception.handle (\e -> myException e
-                                  >> performGC
-                                  >> respawn opts
-                                ) $ runResourceT do
+runPeer opts = respawnOnError opts $ runResourceT do
 
   myself <- liftIO myThreadId
 
@@ -1223,8 +1236,8 @@ runPeer opts = Exception.handle (\e -> myException e
   rpcProto <- async $ flip runReaderT rpcctx do
     env <- newNotifyEnvServer @(RefChanEvents L4Proto) refChanNotifySource
     envrl <- newNotifyEnvServer @(RefLogEvents L4Proto) refLogNotifySource
-    w1 <- asyncLinked $ runNotifyWorkerServer env
-    w2 <- asyncLinked $ runNotifyWorkerServer envrl
+    w1 <- async $ runNotifyWorkerServer env
+    w2 <- async $ runNotifyWorkerServer envrl
     wws <- replicateM 1 $ async $ runProto @UNIX
       [ makeResponse (makeServer @PeerAPI)
       , makeResponse (makeServer @RefLogAPI)
@@ -1234,7 +1247,7 @@ runPeer opts = Exception.handle (\e -> myException e
       , makeResponse (makeNotifyServer @(RefChanEvents L4Proto) env)
       , makeResponse (makeNotifyServer @(RefLogEvents L4Proto) envrl)
       ]
-    mapM_ wait (w1 : w2 : wws )
+    void $ waitAnyCancel (w1 : w2 : wws )
 
   void $ waitAnyCancel $ w <> [ loop
                               , m1
diff --git a/hbs2-peer/app/RPC2/Die.hs b/hbs2-peer/app/RPC2/Die.hs
index 15d3a587..34f3263e 100644
--- a/hbs2-peer/app/RPC2/Die.hs
+++ b/hbs2-peer/app/RPC2/Die.hs
@@ -17,8 +17,7 @@ instance (MonadIO m) => HandleMethod m RpcDie where
   handleMethod _ = do
     debug $ "rpc.die: exiting"
     void $ liftIO $ do
-      w <- async $ pause @'Seconds 0.5 >> Exit.exitSuccess
-      link w
+      pause @'Seconds 0.5 >> Exit.exitSuccess
 
 
 
diff --git a/hbs2-peer/lib/HBS2/Peer/RPC/Internal/Types.hs b/hbs2-peer/lib/HBS2/Peer/RPC/Internal/Types.hs
index dfb23ede..38d7c357 100644
--- a/hbs2-peer/lib/HBS2/Peer/RPC/Internal/Types.hs
+++ b/hbs2-peer/lib/HBS2/Peer/RPC/Internal/Types.hs
@@ -25,6 +25,7 @@ import Control.Monad
 import Control.Monad.Reader
 import Data.ByteString ( ByteString )
 import UnliftIO
+import HBS2.Prelude (asyncLinked)
 
 data RPC2Context =
   RPC2Context
@@ -50,6 +51,6 @@ instance Monad m => HasOwnPeer UNIX (ReaderT RPC2Context m) where
 
 instance (MonadUnliftIO m, HasProtocol UNIX (ServiceProto (api :: [Type]) UNIX))
   => HasDeferred (ServiceProto api UNIX) UNIX  m where
-  deferred m = void $ async m
+  deferred m = void $ asyncLinked m