diff --git a/src/Simplex/Messaging/Server.hs b/src/Simplex/Messaging/Server.hs index fd80c86c5..29de4ff6d 100644 --- a/src/Simplex/Messaging/Server.hs +++ b/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 diff --git a/src/Simplex/Messaging/Server/Env/STM.hs b/src/Simplex/Messaging/Server/Env/STM.hs index a485ab73e..debd36c12 100644 --- a/src/Simplex/Messaging/Server/Env/STM.hs +++ b/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 diff --git a/src/Simplex/Messaging/Transport/Server.hs b/src/Simplex/Messaging/Transport/Server.hs index fe4bd5879..8913ba0c3 100644 --- a/src/Simplex/Messaging/Transport/Server.hs +++ b/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