smp server: stop server faster (#1371)

* smp server: stop server faster

* .702

* asks

* fix leak

* kill client when server stopped

* remove version

* fix version

src/Simplex/Messaging/Server.hs

@@ -160,7 +160,8 @@ smpServer started cfg@ServerConfig {transports, transportConfig = tCfg} attachHT
     runServer (tcpPort, ATransport t, addHTTP) = do
       smpCreds <- asks tlsServerCreds
       httpCreds_ <- asks httpServerCreds
-      ss <- asks sockets
+      ss <- liftIO newSocketState
+      asks sockets >>= atomically . (`modifyTVar'` ((tcpPort, ss) :))
       serverSignKey <- either fail pure $ fromTLSCredentials smpCreds
       env <- ask
       liftIO $ case (httpCreds_, attachHTTP_) of
@@ -190,6 +191,7 @@ smpServer started cfg@ServerConfig {transports, transportConfig = tCfg} attachHT
 
     stopServer :: Server -> M ()
     stopServer s = do
+      asks serverActive >>= atomically . (`writeTVar` False)
       logInfo "Saving server state..."
       withLock' (savingLock s) "final" $ saveServer False >> closeServer
       logInfo "Server stopped"
@@ -657,25 +659,28 @@ smpServer started cfg@ServerConfig {transports, transportConfig = tCfg} attachHT
 #else
                 hPutStrLn h "Not available on GHC 8.10"
 #endif
-              CPSockets -> withUserRole $ do
-                (accepted', closed', active') <- unliftIO u $ asks sockets
-                (accepted, closed, active) <- (,,) <$> readTVarIO accepted' <*> readTVarIO closed' <*> readTVarIO active'
-                hPutStrLn h "Sockets: "
-                hPutStrLn h $ "accepted: " <> show accepted
-                hPutStrLn h $ "closed: " <> show closed
-                hPutStrLn h $ "active: " <> show (IM.size active)
-                hPutStrLn h $ "leaked: " <> show (accepted - closed - IM.size active)
+              CPSockets -> withUserRole $ unliftIO u (asks sockets) >>= readTVarIO >>= mapM_ putSockets
+                where
+                  putSockets (tcpPort, (accepted', closed', active')) = do
+                    (accepted, closed, active) <- (,,) <$> readTVarIO accepted' <*> readTVarIO closed' <*> readTVarIO active'
+                    hPutStrLn h $ "Sockets for port " <> tcpPort <> ":"
+                    hPutStrLn h $ "accepted: " <> show accepted
+                    hPutStrLn h $ "closed: " <> show closed
+                    hPutStrLn h $ "active: " <> show (IM.size active)
+                    hPutStrLn h $ "leaked: " <> show (accepted - closed - IM.size active)
               CPSocketThreads -> withAdminRole $ do
 #if MIN_VERSION_base(4,18,0)
-                (_, _, active') <- unliftIO u $ asks sockets
-                active <- readTVarIO active'
-                forM_ (IM.toList active) $ \(sid, tid') ->
-                  deRefWeak tid' >>= \case
-                    Nothing -> hPutStrLn h $ intercalate "," [show sid, "", "gone", ""]
-                    Just tid -> do
-                      label <- threadLabel tid
-                      status <- threadStatus tid
-                      hPutStrLn h $ intercalate "," [show sid, show tid, show status, fromMaybe "" label]
+                unliftIO u (asks sockets) >>= readTVarIO >>= mapM_ putSocketThreads
+                where
+                  putSocketThreads (tcpPort, (_, _, active')) = do
+                    active <- readTVarIO active'
+                    forM_ (IM.toList active) $ \(sid, tid') ->
+                      deRefWeak tid' >>= \case
+                        Nothing -> hPutStrLn h $ intercalate "," [tcpPort, show sid, "", "gone", ""]
+                        Just tid -> do
+                          label <- threadLabel tid
+                          status <- threadStatus tid
+                          hPutStrLn h $ intercalate "," [tcpPort, show sid, show tid, show status, fromMaybe "" label]
 #else
                 hPutStrLn h "Not available on GHC 8.10"
 #endif
@@ -830,12 +835,13 @@ clientDisconnected c@Client {clientId, subscriptions, ntfSubscriptions, connecte
   subs <- atomically $ swapTVar subscriptions M.empty
   ntfSubs <- atomically $ swapTVar ntfSubscriptions M.empty
   liftIO $ mapM_ cancelSub subs
-  Server {subscribers, notifiers, subClients, ntfSubClients} <- asks server
-  liftIO $ updateSubscribers subs subscribers
-  liftIO $ updateSubscribers ntfSubs notifiers
-  asks clients >>= atomically . (`modifyTVar'` IM.delete clientId)
-  atomically $ modifyTVar' subClients $ IM.delete clientId
-  atomically $ modifyTVar' ntfSubClients $ IM.delete clientId
+  whenM (asks serverActive >>= readTVarIO) $ do
+    Server {subscribers, notifiers, subClients, ntfSubClients} <- asks server
+    liftIO $ updateSubscribers subs subscribers
+    liftIO $ updateSubscribers ntfSubs notifiers
+    asks clients >>= atomically . (`modifyTVar'` IM.delete clientId)
+    atomically $ modifyTVar' subClients $ IM.delete clientId
+    atomically $ modifyTVar' ntfSubClients $ IM.delete clientId
   tIds <- atomically $ swapTVar endThreads IM.empty
   liftIO $ mapM_ (mapM_ killThread <=< deRefWeak) tIds
   where
@@ -862,8 +868,10 @@ cancelSub s = case subThread s of
 receive :: Transport c => THandleSMP c 'TServer -> Client -> M ()
 receive h@THandle {params = THandleParams {thAuth}} Client {rcvQ, sndQ, rcvActiveAt, sessionId} = do
   labelMyThread . B.unpack $ "client $" <> encode sessionId <> " receive"
+  sa <- asks serverActive
   forever $ do
     ts <- L.toList <$> liftIO (tGet h)
+    unlessM (readTVarIO sa) $ throwIO $ userError "server stopped"
     atomically . (writeTVar rcvActiveAt $!) =<< liftIO getSystemTime
     stats <- asks serverStats
     (errs, cmds) <- partitionEithers <$> mapM (cmdAction stats) ts

src/Simplex/Messaging/Server/Env/STM.hs

@@ -138,6 +138,7 @@ defaultProxyClientConcurrency = 32
 
 data Env = Env
   { config :: ServerConfig,
+    serverActive :: TVar Bool,
     serverInfo :: ServerInformation,
     server :: Server,
     serverIdentity :: KeyHash,
@@ -149,7 +150,7 @@ data Env = Env
     tlsServerCreds :: T.Credential,
     httpServerCreds :: Maybe T.Credential,
     serverStats :: ServerStats,
-    sockets :: SocketState,
+    sockets :: TVar [(ServiceName, SocketState)],
     clientSeq :: TVar ClientId,
     clients :: TVar (IntMap (Maybe Client)),
     proxyAgent :: ProxyAgent -- senders served on this proxy
@@ -243,6 +244,7 @@ newProhibitedSub = do
 
 newEnv :: ServerConfig -> IO Env
 newEnv config@ServerConfig {smpCredentials, httpCredentials, storeLogFile, smpAgentCfg, information, messageExpiration} = do
+  serverActive <- newTVarIO True
   server <- newServer
   queueStore <- newQueueStore
   msgStore <- newMsgStore
@@ -258,11 +260,11 @@ newEnv config@ServerConfig {smpCredentials, httpCredentials, storeLogFile, smpAg
   Fingerprint fp <- loadFingerprint smpCredentials
   let serverIdentity = KeyHash fp
   serverStats <- newServerStats =<< getCurrentTime
-  sockets <- newSocketState
+  sockets <- newTVarIO []
   clientSeq <- newTVarIO 0
   clients <- newTVarIO mempty
   proxyAgent <- newSMPProxyAgent smpAgentCfg random
-  pure Env {config, serverInfo, server, serverIdentity, queueStore, msgStore, ntfStore, random, storeLog, tlsServerCreds, httpServerCreds, serverStats, sockets, clientSeq, clients, proxyAgent}
+  pure Env {serverActive, config, serverInfo, server, serverIdentity, queueStore, msgStore, ntfStore, random, storeLog, tlsServerCreds, httpServerCreds, serverStats, sockets, clientSeq, clients, proxyAgent}
   where
     getCredentials protocol creds = do
       files <- missingCreds

src/Simplex/Messaging/Transport/Server.hs

@@ -132,7 +132,7 @@ runTCPServerSocket :: SocketState -> TMVar Bool -> IO Socket -> (Socket -> IO ()
 runTCPServerSocket (accepted, gracefullyClosed, clients) started getSocket server =
   E.bracket getSocket (closeServer started clients) $ \sock ->
     forever . E.bracketOnError (safeAccept sock) (close . fst) $ \(conn, _peer) -> do
-      cId <- atomically $ stateTVar accepted $ \cId -> let cId' = cId + 1 in cId `seq` (cId', cId')
+      cId <- atomically $ stateTVar accepted $ \cId -> let cId' = cId + 1 in cId' `seq` (cId', cId')
       let closeConn _ = do
             atomically $ modifyTVar' clients $ IM.delete cId
             gracefulClose conn 5000 `catchAll_` pure () -- catchAll_ is needed here in case the connection was closed earlier
@@ -168,8 +168,8 @@ newSocketState = (,,) <$> newTVarIO 0 <*> newTVarIO 0 <*> newTVarIO mempty
 
 closeServer :: TMVar Bool -> TVar (IntMap (Weak ThreadId)) -> Socket -> IO ()
 closeServer started clients sock = do
-  readTVarIO clients >>= mapM_ (deRefWeak >=> mapM_ killThread)
   close sock
+  readTVarIO clients >>= mapM_ (deRefWeak >=> mapM_ killThread)
   void . atomically $ tryPutTMVar started False
 
 startTCPServer :: TMVar Bool -> Maybe HostName -> ServiceName -> IO Socket