smp server: correctly track if ntf service is subscribed and total subscribed queues count (fixes race condition between NSUB and NSUBS from notification server) (#1583)

Author: epoberezkin

src/Simplex/Messaging/Server.hs

@@ -1321,7 +1321,7 @@ client
   -- TODO [certs rcv] rcv subscriptions
   Server {subscribers, ntfSubscribers}
   ms
-  clnt@Client {clientId, subscriptions, ntfSubscriptions, serviceSubsCount = _todo', ntfServiceSubsCount, rcvQ, sndQ, clientTHParams = thParams'@THandleParams {sessionId}, procThreads} = do
+  clnt@Client {clientId, subscriptions, ntfSubscriptions, ntfServiceSubscribed, serviceSubsCount = _todo', ntfServiceSubsCount, rcvQ, sndQ, clientTHParams = thParams'@THandleParams {sessionId}, procThreads} = do
     labelMyThread . B.unpack $ "client $" <> encode sessionId <> " commands"
     let THandleParams {thVersion} = thParams'
         service = peerClientService =<< thAuth thParams'
@@ -1721,18 +1721,19 @@ client
 
         subscribeServiceNotifications :: THPeerClientService -> M s BrokerMsg
         subscribeServiceNotifications THClientService {serviceId} = do
-          srvSubs <- readTVarIO ntfServiceSubsCount
-          if srvSubs == 0
-            then
+          subscribed <- readTVarIO ntfServiceSubscribed
+          if subscribed
+            then SOKS <$> readTVarIO ntfServiceSubsCount
+            else
               liftIO (getNtfServiceQueueCount @(StoreQueue s) (queueStore ms) serviceId) >>= \case
                 Left e -> pure $ ERR e
-                Right count -> do
+                Right !count' -> do
                   atomically $ do
-                    modifyTVar' ntfServiceSubsCount (+ count) -- service count
-                    modifyTVar' (totalServiceSubs ntfSubscribers) (+ count) -- server count for all services
+                    writeTVar ntfServiceSubscribed True
+                    count <- swapTVar ntfServiceSubsCount count'
+                    modifyTVar' (totalServiceSubs ntfSubscribers) (+ (count' - count)) -- server count for all services
                   atomically $ writeTQueue (subQ ntfSubscribers) (CSService serviceId, clientId)
-                  pure $ SOKS count
-            else pure $ SOKS srvSubs
+                  pure $ SOKS count'
 
         acknowledgeMsg :: MsgId -> StoreQueue s -> QueueRec -> M s (Transmission BrokerMsg)
         acknowledgeMsg msgId q qr = time "ACK" $ do

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

@@ -389,6 +389,8 @@ data Client s = Client
   { clientId :: ClientId,
     subscriptions :: TMap RecipientId Sub,
     ntfSubscriptions :: TMap NotifierId (),
+    serviceSubscribed :: TVar Bool, -- set independently of serviceSubsCount, to track whether service subscription command was received
+    ntfServiceSubscribed :: TVar Bool,
     serviceSubsCount :: TVar Int64, -- only one service can be subscribed, based on its certificate, this is subscription count
     ntfServiceSubsCount :: TVar Int64, -- only one service can be subscribed, based on its certificate, this is subscription count
     rcvQ :: TBQueue (NonEmpty (VerifiedTransmission s)),
@@ -458,6 +460,8 @@ newClient :: ClientId -> Natural -> THandleParams SMPVersion 'TServer -> SystemT
 newClient clientId qSize clientTHParams createdAt = do
   subscriptions <- TM.emptyIO
   ntfSubscriptions <- TM.emptyIO
+  serviceSubscribed <- newTVarIO False
+  ntfServiceSubscribed <- newTVarIO False
   serviceSubsCount <- newTVarIO 0
   ntfServiceSubsCount <- newTVarIO 0
   rcvQ <- newTBQueueIO qSize
@@ -474,6 +478,8 @@ newClient clientId qSize clientTHParams createdAt = do
       { clientId,
         subscriptions,
         ntfSubscriptions,
+        serviceSubscribed,
+        ntfServiceSubscribed,
         serviceSubsCount,
         ntfServiceSubsCount,
         rcvQ,