smp server: save stats periodically

src/Simplex/Messaging/Server.hs

@@ -115,6 +115,7 @@ import Simplex.Messaging.Server.QueueStore.QueueInfo
 import Simplex.Messaging.Server.QueueStore.STM (withLoadedQueues)
 import Simplex.Messaging.Server.QueueStore.Types
 import Simplex.Messaging.Server.Stats
+import Simplex.Messaging.Server.StatsStore
 import Simplex.Messaging.Server.StoreLog (foldLogLines)
 import Simplex.Messaging.SystemTime
 import Simplex.Messaging.TMap (TMap)
@@ -143,6 +144,7 @@ import GHC.Conc.Sync (threadLabel)
 
 #if defined(dbServerPostgres)
 import Simplex.Messaging.Server.MsgStore.Postgres (getDbMessageStats)
+import qualified Simplex.Messaging.Server.QueueStore.Postgres as PQ
 #endif
 
 -- | Runs an SMP server using passed configuration.
@@ -195,6 +197,7 @@ smpServer started cfg@ServerConfig {transports, transportConfig = tCfg, startOpt
         : map runServer transports
             <> expireMessagesThread_ cfg
             <> serverStatsThread_ cfg
+            <> saveStatsThread_ cfg
             <> prometheusMetricsThread_ cfg
             <> controlPortThread_ cfg
     )
@@ -238,10 +241,11 @@ smpServer started cfg@ServerConfig {transports, transportConfig = tCfg, startOpt
 
     saveServer :: Bool -> M s ()
     saveServer drainMsgs = do
+      -- the database stats store is closed together with the message store
+      saveServerStats
       ms <- asks msgStore_
       liftIO $ saveServerMessages drainMsgs ms >> closeMsgStore (fromMsgStore ms)
       saveServerNtfs
-      saveServerStats
 
     closeServer :: M s ()
     closeServer = do
@@ -512,6 +516,18 @@ smpServer started cfg@ServerConfig {transports, transportConfig = tCfg, startOpt
       [logServerStats logStatsStartTime interval serverStatsLogFile]
     serverStatsThread_ _ = []
 
+    saveStatsThread_ :: ServerConfig s -> [M s ()]
+    saveStatsThread_ ServerConfig {serverStatsBackupFile = Just _, statsSaveInterval} = [saveStatsThread statsSaveInterval]
+    saveStatsThread_ _ = []
+
+    saveStatsThread :: Int64 -> M s ()
+    saveStatsThread interval = do
+      labelMyThread "saveStatsThread"
+      lock <- asks $ savingLock . server
+      forever $ do
+        liftIO $ threadDelay' $ 1000000 * interval
+        withLock' lock "saveStats" saveServerStats
+
     logServerStats :: Int64 -> Int64 -> FilePath -> M s ()
     logServerStats startAt logInterval statsFilePath = do
       labelMyThread "logServerStats"
@@ -523,9 +539,12 @@ smpServer started cfg@ServerConfig {transports, transportConfig = tCfg, startOpt
       st <- asks msgStore
       EntityCounts {queueCount, notifierCount, rcvServiceCount, ntfServiceCount, rcvServiceQueuesCount, ntfServiceQueuesCount} <-
         liftIO $ getEntityCounts @(StoreQueue s) $ queueStore st
+      lock <- asks $ savingLock . server
+      u <- askUnliftIO
       let interval = 1000000 * logInterval
       forever $ do
-        withFile statsFilePath AppendMode $ \h -> liftIO $ do
+        withFile statsFilePath AppendMode $ \h -> withLock' lock "logStats" $ liftIO $ do
+          unliftIO u saveServerStats
           hSetBuffering h LineBuffering
           ts <- getCurrentTime
           fromTime' <- atomicSwapIORef fromTime ts
@@ -657,6 +676,7 @@ smpServer started cfg@ServerConfig {transports, transportConfig = tCfg, startOpt
                        <> showServiceStats ntfServices'
                        <> showNameResolverStats rslvStats'
               )
+          unliftIO u saveServerStats
         liftIO $ threadDelay' interval
       where
         showProxyStats ProxyStatsData {_pRequests, _pSuccesses, _pErrorsConnect, _pErrorsCompat, _pErrorsOther} =
@@ -2512,38 +2532,39 @@ restoreServerNtfs =
             ntfErr :: Show e => String -> e -> String
             ntfErr op e = op <> " error (" <> show e <> "): " <> B.unpack (B.take 100 s)
 
-saveServerStats :: M s ()
-saveServerStats =
-  asks (serverStatsBackupFile . config)
-    >>= mapM_ (\f -> asks serverStats >>= liftIO . getServerStatsData >>= liftIO . saveStats f)
+serverStatsStore :: M s (Maybe StatsStore)
+serverStatsStore = do
+  ms <- asks msgStore_
+  fmap (statsStore ms) <$> asks (serverStatsBackupFile . config)
   where
-    saveStats f stats = do
-      logNote $ "saving server stats to file " <> T.pack f
-      B.writeFile f $ strEncode stats
-      logNote "server stats saved"
+    statsStore :: MsgStore s -> FilePath -> StatsStore
+    statsStore = \case
+      StoreMemory _ -> StatsFile
+#if defined(dbServerPostgres)
+      StoreDatabase ms -> StatsDB (PQ.dbStore $ queueStore ms)
+#endif
+
+saveServerStats :: M s ()
+saveServerStats = do
+  ss <- asks serverStats
+  serverStatsStore >>= mapM_ (liftIO . (`saveStats` ss))
 
 restoreServerStats :: forall s. MsgStoreClass s => Maybe MessageStats -> MessageStats -> M s ()
-restoreServerStats msgStats_ ntfStats = asks (serverStatsBackupFile . config) >>= mapM_ restoreStats
+restoreServerStats msgStats_ ntfStats = do
+  s <- asks serverStats
+  serverStatsStore >>= mapM_ (liftIO . tryAny . loadStats >=> restore s)
+  st <- asks msgStore
+  EntityCounts {queueCount} <- liftIO $ getEntityCounts @(StoreQueue s) $ queueStore st
+  liftIO $ do
+    writeIORef (qCount s) queueCount
+    writeIORef (msgCount s) $ maybe 0 storedMsgsCount msgStats_
+    writeIORef (ntfCount s) $ storedMsgsCount ntfStats
+    atomicModifyIORef'_ (msgExpired s) (+ maybe 0 expiredMsgsCount msgStats_)
+    atomicModifyIORef'_ (msgNtfExpired s) (+ expiredMsgsCount ntfStats)
   where
-    restoreStats f = whenM (doesFileExist f) $ do
-      logNote $ "restoring server stats from file " <> T.pack f
-      liftIO (strDecode <$> B.readFile f) >>= \case
-        Right d@ServerStatsData {_qCount = statsQCount, _msgCount = statsMsgCount, _ntfCount = statsNtfCount} -> do
-          s <- asks serverStats
-          st <- asks msgStore
-          EntityCounts {queueCount = _qCount} <- liftIO $ getEntityCounts @(StoreQueue s) $ queueStore st
-          let _msgCount = maybe statsMsgCount storedMsgsCount msgStats_
-              _ntfCount = storedMsgsCount ntfStats
-              _msgExpired' = _msgExpired d + maybe 0 expiredMsgsCount msgStats_
-              _msgNtfExpired' = _msgNtfExpired d + expiredMsgsCount ntfStats
-          liftIO $ setServerStats s d {_qCount, _msgCount, _ntfCount, _msgExpired = _msgExpired', _msgNtfExpired = _msgNtfExpired'}
-          renameFile f $ f <> ".bak"
-          logNote "server stats restored"
-          compareCounts "Queue" statsQCount _qCount
-          compareCounts "Message" statsMsgCount _msgCount
-          compareCounts "Notification" statsNtfCount _ntfCount
-        Left e -> do
-          logNote $ "error restoring server stats: " <> T.pack e
-          liftIO exitFailure
-    compareCounts name statsCnt storeCnt =
-      when (statsCnt /= storeCnt) $ logWarn $ name <> " count differs: stats: " <> tshow statsCnt <> ", store: " <> tshow storeCnt
+    restore s = \case
+      Right (Just d) -> liftIO (setServerStats s d) >> logNote "server stats restored"
+      Right Nothing -> pure ()
+      Left e -> do
+        logError $ "error restoring server stats: " <> tshow e
+        liftIO exitFailure

tests/ServerTests.hs

@@ -33,7 +33,7 @@ import Data.Hashable (hash)
 import qualified Data.IntSet as IS
 import Data.List.NonEmpty (NonEmpty)
 import Data.List (isPrefixOf)
-import Data.Maybe (catMaybes)
+import Data.Maybe (catMaybes, fromMaybe)
 import Data.String (IsString (..))
 import Data.Time.Clock (NominalDiffTime, diffUTCTime, getCurrentTime)
 import Text.Read (readMaybe)
@@ -66,6 +66,8 @@ import Simplex.Messaging.Server.Env.STM (readWriteQueueStore)
 import Simplex.Messaging.Server.MsgStore.Postgres (PostgresMsgStoreCfg (..), exportDbMessages)
 import Simplex.Messaging.Server.MsgStore.Types (MsgStoreClass (..), newMsgStore)
 import Simplex.Messaging.Server.StoreLog (closeStoreLog)
+import CoreTests.ServerStatsTests (withStatsDB)
+import Simplex.Messaging.Server.StatsStore (getDBStats)
 #endif
 
 serverTests :: SpecWith (ASrvTransport, AStoreType)
@@ -94,6 +96,9 @@ serverTests = do
   describe "Store log" testWithStoreLog
   describe "Restore messages" testRestoreMessages
   describe "Restore messages (old / v2)" testRestoreExpireMessages
+  describe "Save server stats" $ do
+    testSaveStatsPeriodically
+    testStartupStoreCounts
   describe "Save prometheus metrics" testPrometheusMetrics
   describe "Delay of AUTH error" testAuthErrorDelay
   describe "Message notifications" $ do
@@ -1098,8 +1103,7 @@ testRestoreMessages =
       pure ()
     rId <- readTVarIO recipientId
     logSize testStoreLogFile `shouldReturn` 2
-    logSize testServerStatsBackupFile `shouldReturn` 101
-    Right stats1 <- strDecode <$> B.readFile testServerStatsBackupFile
+    stats1 <- readSavedStats msType
     checkStats stats1 [rId] 5 1
     withSmpServerConfigOn at cfg' testPort . runTest t $ \h -> do
       Just rKey <- readTVarIO recipientKey
@@ -1114,8 +1118,7 @@ testRestoreMessages =
     logSize testStoreLogFile `shouldReturn` (if compacting then 1 else 2)
     -- the last message is not removed because it was not ACK'd
     -- logSize testStoreMsgsFile `shouldReturn` 3
-    logSize testServerStatsBackupFile `shouldReturn` 101
-    Right stats2 <- strDecode <$> B.readFile testServerStatsBackupFile
+    stats2 <- readSavedStats msType
     checkStats stats2 [rId] 5 3
 
     withSmpServerConfigOn at cfg' testPort . runTest t $ \h -> do
@@ -1132,11 +1135,10 @@ testRestoreMessages =
       pure ()
     logSize testStoreLogFile `shouldReturn` (if compacting then 1 else 2)
     removeFile testStoreLogFile
-    logSize testServerStatsBackupFile `shouldReturn` 101
-    Right stats3 <- strDecode <$> B.readFile testServerStatsBackupFile
+    stats3 <- readSavedStats msType
     checkStats stats3 [rId] 5 5
     removeFileIfExists testStoreMsgsFile
-    removeFile testServerStatsBackupFile
+    removeFileIfExists testServerStatsBackupFile
   where
     runTest :: (HasTestEnv, Transport c) => TProxy c 'TServer -> (THandleSMP c 'TClient -> IO ()) -> ThreadId -> Expectation
     runTest _ test' server = do
@@ -1162,6 +1164,68 @@ checkStats s qs sent received = do
   IS.toList _week `shouldBe` map (hash . unEntityId) qs
   IS.toList _month `shouldBe` map (hash . unEntityId) qs
 
+readSavedStats :: HasTestEnv => AStoreType -> IO ServerStatsData
+readSavedStats = \case
+  ASType _ SMSMemory -> do
+    logSize testServerStatsBackupFile `shouldReturn` 101
+    either error id . strDecode <$> B.readFile testServerStatsBackupFile
+#if defined(dbServerPostgres)
+  ASType _ SMSPostgres -> withStatsDB $ fmap (fromMaybe $ error "no saved stats") . getDBStats
+#else
+  ASType _ SMSPostgres -> error "compiled without server_postgres flag"
+#endif
+
+testSaveStatsPeriodically :: SpecWith (ASrvTransport, AStoreType)
+testSaveStatsPeriodically =
+  it "should save stats while running and restore them on start" $ \(at@(ATransport t), msType) -> do
+    removeFileIfExists testServerStatsBackupFile
+    g <- C.newRandom
+    (sPub, sKey) <- atomically $ C.generateAuthKeyPair C.SEd25519 g
+    let cfg' = updateCfg (cfgMS msType) $ \cfg_ -> cfg_ {serverStatsBackupFile = Just testServerStatsBackupFile, statsSaveInterval = 1}
+    (rId, stats1, snapshot) <- withSmpServerConfigOn at cfg' testPort $ \_ -> do
+      rId <- runClient t $ \h -> runClient t $ \rh -> do
+        (sId, rId, _, _) <- createAndSecureQueue rh sPub
+        Resp "1" _ OK <- signSendRecv h sKey ("1", sId, _SEND "hello")
+        pure rId
+      threadDelay 1500000
+      stats1 <- readSavedStats msType
+      snapshot <- case msType of
+        ASType _ SMSMemory -> Just <$> B.readFile testServerStatsBackupFile
+        _ -> pure Nothing
+      pure (rId, stats1, snapshot)
+    checkStats stats1 [rId] 1 0
+    -- a crash leaves the last periodic save, not the save on shutdown
+    mapM_ (B.writeFile testServerStatsBackupFile) snapshot
+    stats2 <- withSmpServerConfigOn at cfg' testPort $ \_ -> threadDelay 1500000 >> readSavedStats msType
+    checkStats stats2 [rId] 1 0
+    removeFileIfExists testServerStatsBackupFile
+  where
+    runClient :: (HasTestEnv, Transport c) => TProxy c 'TServer -> (THandleSMP c 'TClient -> IO a) -> IO a
+    runClient _ = testSMPClient
+
+testStartupStoreCounts :: SpecWith (ASrvTransport, AStoreType)
+testStartupStoreCounts =
+  it "should count stored queues and messages on start without saved stats" $ \(at@(ATransport t), msType) -> do
+    g <- C.newRandom
+    (sPub, sKey) <- atomically $ C.generateAuthKeyPair C.SEd25519 g
+    let cfg' = cfgMS msType
+    withSmpServerConfigOn at cfg' testPort $ \_ -> runClient t $ \h -> runClient t $ \rh -> do
+      (sId, _, _, _) <- createAndSecureQueue rh sPub
+      Resp "1" _ OK <- signSendRecv h sKey ("1", sId, _SEND "hello 1")
+      Resp "2" _ OK <- signSendRecv h sKey ("2", sId, _SEND "hello 2")
+      pure ()
+    let cfg2 = updateCfg cfg' $ \cfg_ -> cfg_ {prometheusInterval = Just 1}
+    metrics <- withSmpServerConfigOn at cfg2 testPort $ \_ -> threadDelay 1500000 >> B.readFile testPrometheusMetricsFile
+    let expected :: [ByteString]
+        expected = case msType of
+          ASType SQSMemory _ -> ["simplex_smp_messages_total 2", "simplex_smp_queues_total1 1"]
+          -- Postgres queue count is a reltuples estimate
+          _ -> ["simplex_smp_messages_total 2"]
+    forM_ expected $ \l -> map (B.unwords . take 2 . B.words) (B.lines metrics) `shouldContain` [l]
+  where
+    runClient :: (HasTestEnv, Transport c) => TProxy c 'TServer -> (THandleSMP c 'TClient -> IO a) -> IO a
+    runClient _ = testSMPClient
+
 testRestoreExpireMessages :: SpecWith (ASrvTransport, AStoreType)
 testRestoreExpireMessages =
   it "should store messages on exit and restore on start (old / v2)" $ \(at@(ATransport t), msType) -> do
@@ -1210,7 +1274,7 @@ testRestoreExpireMessages =
     msgs'' <- B.readFile testStoreMsgsFile
     length (B.lines msgs'') `shouldBe` 2
     B.lines msgs'' `shouldBe` drop 2 (B.lines msgs)
-    Right ServerStatsData {_msgExpired} <- strDecode <$> B.readFile testServerStatsBackupFile
+    ServerStatsData {_msgExpired} <- readSavedStats msType
     _msgExpired `shouldBe` 2
   where
     exportStoreMessages :: HasTestEnv => AStoreType -> IO ()