diff --git a/src/Simplex/FileTransfer/Agent.hs b/src/Simplex/FileTransfer/Agent.hs index d83983363..0f7b19b65 100644 --- a/src/Simplex/FileTransfer/Agent.hs +++ b/src/Simplex/FileTransfer/Agent.hs @@ -210,7 +210,7 @@ runXFTPRcvWorker c srv Worker {doWork} = do withRetryIntervalLimit xftpConsecutiveRetries ri' $ \delay' loop -> do liftIO $ waitWhileSuspended c liftIO $ waitForUserNetwork c - atomically $ incXFTPServerStat c userId srv downloadAttempts + incXFTPServerStat c userId srv downloadAttempts downloadFileChunk fc replica approvedRelays `catchAllErrors` \e -> retryOnError "XFTP rcv worker" (retryLoop loop e delay') (retryDone e) e where @@ -222,7 +222,7 @@ runXFTPRcvWorker c srv Worker {doWork} = do liftIO $ assertAgentForeground c loop retryDone e = do - atomically . incXFTPServerStat c userId srv $ case e of + incXFTPServerStat c userId srv $ case e of XFTP _ XFTP.AUTH -> downloadAuthErrs _ -> downloadErrs rcvWorkerInternalError c rcvFileId rcvFileEntityId redirectEntityId_ (Just fileTmpPath) e @@ -247,8 +247,8 @@ runXFTPRcvWorker c srv Worker {doWork} = do Just RcvFileRedirect {redirectFileInfo = RedirectFileInfo {size = FileSize finalSize}, redirectEntityId} -> (redirectEntityId, finalSize) liftIO . when complete $ updateRcvFileStatus db rcvFileId RFSReceived pure (entityId, complete, RFPROG rcvd total) - atomically $ incXFTPServerStat c userId srv downloads - atomically $ incXFTPServerSizeStat c userId srv downloadsSize (fromIntegral $ toKB chSize) + incXFTPServerStat c userId srv downloads + incXFTPServerSizeStat c userId srv downloadsSize (fromIntegral $ toKB chSize) notify c entityId progress when complete . lift . void $ getXFTPRcvWorker True c Nothing @@ -536,7 +536,7 @@ runXFTPSndWorker c srv Worker {doWork} = do withRetryIntervalLimit xftpConsecutiveRetries ri' $ \delay' loop -> do liftIO $ waitWhileSuspended c liftIO $ waitForUserNetwork c - atomically $ incXFTPServerStat c userId srv uploadAttempts + incXFTPServerStat c userId srv uploadAttempts uploadFileChunk cfg fc replica `catchAllErrors` \e -> retryOnError "XFTP snd worker" (retryLoop loop e delay') (retryDone e) e where @@ -548,7 +548,7 @@ runXFTPSndWorker c srv Worker {doWork} = do liftIO $ assertAgentForeground c loop retryDone e = do - atomically $ incXFTPServerStat c userId srv uploadErrs + incXFTPServerStat c userId srv uploadErrs sndWorkerInternalError c sndFileId sndFileEntityId (Just filePrefixPath) e uploadFileChunk :: AgentConfig -> SndFileChunk -> SndFileChunkReplica -> AM () uploadFileChunk AgentConfig {xftpMaxRecipientsPerRequest = maxRecipients} sndFileChunk@SndFileChunk {sndFileId, userId, chunkSpec = chunkSpec@XFTPChunkSpec {filePath, chunkSize = chSize}, digest = chunkDigest} replica = do @@ -566,8 +566,8 @@ runXFTPSndWorker c srv Worker {doWork} = do let uploaded = uploadedSize chunks total = totalSize chunks complete = all chunkUploaded chunks - atomically $ incXFTPServerStat c userId srv uploads - atomically $ incXFTPServerSizeStat c userId srv uploadsSize (fromIntegral $ toKB chSize) + incXFTPServerStat c userId srv uploads + incXFTPServerSizeStat c userId srv uploadsSize (fromIntegral $ toKB chSize) notify c sndFileEntityId $ SFPROG uploaded total when complete $ do (sndDescr, rcvDescrs) <- sndFileToDescrs sf @@ -715,7 +715,7 @@ runXFTPDelWorker c srv Worker {doWork} = do withRetryIntervalLimit xftpConsecutiveRetries ri' $ \delay' loop -> do liftIO $ waitWhileSuspended c liftIO $ waitForUserNetwork c - atomically $ incXFTPServerStat c userId srv deleteAttempts + incXFTPServerStat c userId srv deleteAttempts deleteChunkReplica `catchAllErrors` \e -> retryOnError "XFTP del worker" (retryLoop loop e delay') (retryDone e) e where @@ -727,12 +727,12 @@ runXFTPDelWorker c srv Worker {doWork} = do liftIO $ assertAgentForeground c loop retryDone e = do - atomically $ incXFTPServerStat c userId srv deleteErrs + incXFTPServerStat c userId srv deleteErrs delWorkerInternalError c deletedSndChunkReplicaId e deleteChunkReplica = do agentXFTPDeleteChunk c userId replica withStore' c $ \db -> deleteDeletedSndChunkReplica db deletedSndChunkReplicaId - atomically $ incXFTPServerStat c userId srv deletions + incXFTPServerStat c userId srv deletions delWorkerInternalError :: AgentClient -> Int64 -> AgentErrorType -> AM () delWorkerInternalError c deletedSndChunkReplicaId e = do diff --git a/src/Simplex/FileTransfer/Client/Agent.hs b/src/Simplex/FileTransfer/Client/Agent.hs index d5faea92f..55cae5d3c 100644 --- a/src/Simplex/FileTransfer/Client/Agent.hs +++ b/src/Simplex/FileTransfer/Client/Agent.hs @@ -75,7 +75,7 @@ type ME a = ExceptT XFTPClientAgentError IO a getXFTPServerClient :: XFTPClientAgent -> XFTPServer -> ME XFTPClient getXFTPServerClient XFTPClientAgent {xftpClients, startedAt, config} srv = do - atomically getClientVar >>= either newXFTPClient waitForXFTPClient + liftIO getClientVar >>= either newXFTPClient waitForXFTPClient where connectClient :: ME XFTPClient connectClient = @@ -88,14 +88,8 @@ getXFTPServerClient XFTPClientAgent {xftpClients, startedAt, config} srv = do atomically $ TM.delete srv xftpClients logInfo $ "disconnected from " <> showServer srv - getClientVar :: STM (Either XFTPClientVar XFTPClientVar) - getClientVar = maybe (Left <$> newClientVar) (pure . Right) =<< TM.lookup srv xftpClients - where - newClientVar :: STM XFTPClientVar - newClientVar = do - var <- newEmptyTMVar - TM.insert srv var xftpClients - pure var + getClientVar :: IO (Either XFTPClientVar XFTPClientVar) + getClientVar = TM.lookupOrInsertIO srv newEmptyTMVar xftpClients waitForXFTPClient :: XFTPClientVar -> ME XFTPClient waitForXFTPClient clientVar = do diff --git a/src/Simplex/Messaging/Agent.hs b/src/Simplex/Messaging/Agent.hs index 5862f3b7c..a474c86a4 100644 --- a/src/Simplex/Messaging/Agent.hs +++ b/src/Simplex/Messaging/Agent.hs @@ -1169,7 +1169,7 @@ createRcvQueue c nm userId connId srvWithAuth@(ProtoServerWithAuth srv _) enable (rq, qUri, tSess, sessId, serviceId_) <- newRcvQueue_ c nm userId connId srvWithAuth vr qd (isJust ntfServer_) subMode nonce_ e2eKeys `catchAllErrors` \e -> liftIO (print e) >> throwE e - atomically $ incSMPServerStat c userId srv connCreated + incSMPServerStat c userId srv connCreated rq' <- withStore c $ \db -> updateNewConnRcv db connId rq subMode lift . when (subMode == SMSubscribe) $ addNewQueueSubscription c rq' tSess sessId serviceId_ mapM_ (newQueueNtfSubscription c rq') ntfServer_ @@ -1709,7 +1709,7 @@ createReplyQueue :: AgentClient -> NetworkRequestMode -> ConnData -> SndQueue -> createReplyQueue c nm ConnData {userId, connId, enableNtfs} SndQueue {smpClientVersion} subMode srv = do ntfServer_ <- if enableNtfs then newQueueNtfServer else pure Nothing (rq, qUri, tSess, sessId, serviceId_) <- newRcvQueue c nm userId connId srv (versionToRange smpClientVersion) SCMInvitation (isJust ntfServer_) subMode - atomically $ incSMPServerStat c userId (qServer rq) connCreated + incSMPServerStat c userId (qServer rq) connCreated let qInfo = toVersionT qUri smpClientVersion rq' <- withStore c $ \db -> upgradeSndConnToDuplex db connId rq subMode lift . when (subMode == SMSubscribe) $ addNewQueueSubscription c rq' tSess sessId serviceId_ @@ -2359,7 +2359,7 @@ runCommandProcessing c@AgentClient {subQ} connId server_ Worker {doWork} = do secureQueue c NRMBackground rq' senderKey -- we may add more statistics special to queue rotation later on, -- not accounting secure during rotation for now: - -- atomically $ incSMPServerStat c userId server connSecured + -- incSMPServerStat c userId server connSecured rq1' <- withStore' c $ \db -> setRcvQueueStatus db rq' Secured >> setRcvSwitchStatus db rq1 (Just RSSendingQUSE) void . enqueueMessages c cData sqs SMP.noMsgFlags $ QUSE [((server, sndId), True)] withStore' c (`deleteCommand` cmdId) @@ -2426,7 +2426,7 @@ runCommandProcessing c@AgentClient {subQ} connId server_ Worker {doWork} = do secure :: RcvQueue -> SMP.SndPublicAuthKey -> AM () secure rq@RcvQueue {server} senderKey = do secureQueue c NRMBackground rq senderKey - atomically $ incSMPServerStat c userId server connSecured + incSMPServerStat c userId server connSecured withStore' c $ \db -> setRcvQueueStatus db rq Secured where withServer a = case server_ of @@ -2647,7 +2647,7 @@ runSmpQueueMsgDelivery c@AgentClient {subQ} sq@SndQueue {userId, connId, server, let err = if msgType == AM_A_MSG_ then MERR mId e else ERR e case e of SMP _ SMP.QUOTA -> do - atomically $ incSMPServerStat c userId server sentQuotaErrs + incSMPServerStat c userId server sentQuotaErrs case msgType of AM_CONN_INFO -> connError msgId NOT_AVAILABLE AM_CONN_INFO_REPLY -> connError msgId NOT_AVAILABLE @@ -2659,7 +2659,7 @@ runSmpQueueMsgDelivery c@AgentClient {subQ} sq@SndQueue {userId, connId, server, notify $ MWARN (unId msgId) e retrySndMsg RISlow SMP _ SMP.AUTH -> do - atomically $ incSMPServerStat c userId server sentAuthErrs + incSMPServerStat c userId server sentAuthErrs case msgType of AM_CONN_INFO -> connError msgId NOT_AVAILABLE AM_CONN_INFO_REPLY -> connError msgId NOT_AVAILABLE @@ -2698,7 +2698,7 @@ runSmpQueueMsgDelivery c@AgentClient {subQ} sq@SndQueue {userId, connId, server, when (serverHostError e) $ notify $ MWARN (unId msgId) e retrySndMsg RIFast | otherwise -> do - atomically $ incSMPServerStat c userId server sentOtherErrs + incSMPServerStat c userId server sentOtherErrs notifyDel msgId err where retrySndMsg riMode = do @@ -2725,7 +2725,7 @@ runSmpQueueMsgDelivery c@AgentClient {subQ} sq@SndQueue {userId, connId, server, -- because it can be sent before HELLO is received -- With `status == Active` condition, CON is sent here only by the accepting party, that previously received HELLO when (status == Active) $ do - atomically $ incSMPServerStat c userId (qServer rq) connCompleted + incSMPServerStat c userId (qServer rq) connCompleted notify $ CON pqEncryption -- this branch should never be reached as receive queue is created before the confirmation, _ -> logError "HELLO sent without receive queue" @@ -2800,7 +2800,7 @@ runSmpQueueMsgDelivery c@AgentClient {subQ} sq@SndQueue {userId, connId, server, pure msgIds_ forM_ (L.nonEmpty msgIds_) $ \msgIds -> do notify $ MERRS (L.map unId msgIds) err - atomically $ incSMPServerStat' c userId server sentExpiredErrs (length msgIds_ + 1) + incSMPServerStat' c userId server sentExpiredErrs (length msgIds_ + 1) delMsg :: InternalId -> AM () delMsg = delMsgKeep False delMsgKeep :: Bool -> InternalId -> AM () @@ -2959,16 +2959,16 @@ synchronizeRatchet' c connId pqSupport' force = withConnLock c connId "synchroni ackQueueMessage :: AgentClient -> RcvQueue -> SMP.MsgId -> AM (Maybe ATransmission) ackQueueMessage c rq@RcvQueue {userId, connId, server, rcvSwchStatus} srvMsgId = do - atomically $ incSMPServerStat c userId server ackAttempts + incSMPServerStat c userId server ackAttempts tryAllErrors (sendAck c rq srvMsgId) >>= \case Right drained -> sendMsgNtf ackMsgs $ drained && rcvSwchStatus == Just RSReceivedQEND Left (SMP _ SMP.NO_MSG) -> sendMsgNtf ackNoMsgErrs False Left e -> do - unless (temporaryOrHostError e) $ atomically $ incSMPServerStat c userId server ackOtherErrs + unless (temporaryOrHostError e) $ incSMPServerStat c userId server ackOtherErrs throwE e where sendMsgNtf stat ended = do - atomically $ incSMPServerStat c userId server stat + incSMPServerStat c userId server stat ifM (liftIO $ hasGetLock c rq) (do atomically $ releaseGetLock c rq brokerTs_ <- eitherToMaybe <$> tryAllErrors (withStore c $ \db -> getRcvMsgBrokerTs db connId srvMsgId) @@ -3084,7 +3084,7 @@ deleteConnQueues c nm waitDelivery ntf rqs = do | otherwise -> do deleteConnRcvQueue db rq -- attempts and successes are counted in deleteQueues function - atomically $ incSMPServerStat c userId server connDeleted + incSMPServerStat c userId server connDeleted pure ((rq, Right ()), Just (Just e)) notify = when ntf . atomically . writeTBQueue (subQ c) connResults :: [(RcvQueue, Either AgentErrorType ())] -> Map ConnId (Either AgentErrorType ()) @@ -3589,7 +3589,7 @@ data ACKd = ACKd | ACKPending -- It cannot be finally, as sometimes it needs to be ACK+DEL, -- and sometimes ACK has to be sent from the consumer. processSMPTransmissions :: AgentClient -> ServerTransmissionBatch SMPVersion ErrorType BrokerMsg -> AM' () -processSMPTransmissions c@AgentClient {subQ} (tSess@(userId, srv, _), THandleParams {thAuth, sessionId = sessId}, ts) = do +processSMPTransmissions c@AgentClient {subQ, smpServersStats} (tSess@(userId, srv, _), THandleParams {thAuth, sessionId = sessId}, ts) = do upConnIds <- newTVarIO [] serviceRQs <- newTVarIO ([] :: [RcvQueue]) forM_ ts $ \(entId, t) -> case t of @@ -3626,7 +3626,7 @@ processSMPTransmissions c@AgentClient {subQ} (tSess@(userId, srv, _), THandlePar connIds <- readTVarIO upConnIds unless (null connIds) $ do notify' "" $ UP srv connIds - atomically $ incSMPServerStat' c userId srv connSubscribed $ length connIds + incSMPServerStat' c userId srv connSubscribed $ length connIds readTVarIO serviceRQs >>= processRcvServiceAssocs c srv where withRcvConn :: SMP.RecipientId -> (forall c. RcvQueue -> Connection c -> AM ()) -> AM' () @@ -3638,23 +3638,27 @@ processSMPTransmissions c@AgentClient {subQ} (tSess@(userId, srv, _), THandlePar Left e -> notify' connId (ERR e) Right () -> pure () processSubOk :: RcvQueue -> TVar [ConnId] -> TVar [RcvQueue] -> Maybe SMP.ServiceId -> IO () - processSubOk rq@RcvQueue {connId} upConnIds serviceRQs serviceId_ = - atomically . whenM (isPendingSub rq) $ do + processSubOk rq@RcvQueue {connId} upConnIds serviceRQs serviceId_ = do + srvStats <- getSrvStats + atomically . whenM (isPendingSub srvStats rq) $ do SS.addActiveSub tSess sessId serviceId_ rq $ currentSubs c modifyTVar' upConnIds (connId :) when (isJust serviceId_ && serviceId_ == clientServiceId_) $ modifyTVar' serviceRQs (rq :) clientServiceId_ = (\THClientService {serviceId} -> serviceId) <$> (clientService =<< thAuth) processSubErr :: RcvQueue -> SMPClientError -> AM' () processSubErr rq@RcvQueue {connId} e = do - atomically . whenM (isPendingSub rq) $ - failSubscription c tSess rq e >> incSMPServerStat c userId srv connSubErrs + srvStats <- liftIO getSrvStats + atomically . whenM (isPendingSub srvStats rq) $ + failSubscription c tSess rq e >> modifyTVar' (connSubErrs srvStats) (+ 1) processClientNotice rq e notifyErr connId e - isPendingSub :: RcvQueue -> STM Bool - isPendingSub rq = do + isPendingSub :: AgentSMPServerStats -> RcvQueue -> STM Bool + isPendingSub srvStats rq = do pending <- (&&) <$> SS.hasPendingSub tSess (queueId rq) (currentSubs c) <*> activeClientSession c tSess sessId - unless pending $ incSMPServerStat c userId srv connSubIgnored + unless pending $ modifyTVar' (connSubIgnored srvStats) (+ 1) pure pending + getSrvStats :: IO AgentSMPServerStats + getSrvStats = either id id <$> TM.lookupOrInsertIO (userId, srv) newAgentSMPServerStats smpServersStats processClientNotice rq e = forM_ (smpErrorClientNotice e) $ \notice_ -> E.bracket_ @@ -3679,7 +3683,7 @@ processSMPTransmissions c@AgentClient {subQ} (tSess@(userId, srv, _), THandlePar pendingMsgs = withConnLock c connId "processSMP" $ case smpMsg of SMP.MSG msg@SMP.RcvMessage {msgId = srvMsgId} -> do - atomically $ incSMPServerStat c userId srv recvMsgs + incSMPServerStat c userId srv recvMsgs void . handleNotifyAck $ do msg' <- decryptSMPMessage rq msg handleNotifyAck $ case msg' of @@ -3765,7 +3769,7 @@ processSMPTransmissions c@AgentClient {subQ} (tSess@(userId, srv, _), THandlePar | otherwise = pure conn' Right Nothing -> prohibited "msg: bad agent msg" >> ack Left e@(AGENT A_DUPLICATE {}) -> do - atomically $ incSMPServerStat c userId srv recvDuplicates + incSMPServerStat c userId srv recvDuplicates withStore' c (\db -> getLastMsg db connId srvMsgId) >>= \case Just RcvMsg {internalId, msgMeta, msgBody = agentMsgBody, userAck} | userAck -> ackDel internalId @@ -3788,7 +3792,7 @@ processSMPTransmissions c@AgentClient {subQ} (tSess@(userId, srv, _), THandlePar _ -> ack _ -> checkDuplicateHash e encryptedMsgHash >> ack Left (AGENT (A_CRYPTO e)) -> do - atomically $ incSMPServerStat c userId srv recvCryptoErrs + incSMPServerStat c userId srv recvCryptoErrs exists <- withStore' c $ \db -> checkRcvMsgHashExists db connId encryptedMsgHash unless exists notifySync ack @@ -3803,7 +3807,7 @@ processSMPTransmissions c@AgentClient {subQ} (tSess@(userId, srv, _), THandlePar notify $ RSYNC rss' (Just e) cStats withStore' c $ \db -> setConnRatchetSync db connId rss' Left e -> do - atomically $ incSMPServerStat c userId srv recvErrs + incSMPServerStat c userId srv recvErrs checkDuplicateHash e encryptedMsgHash >> ack where checkDuplicateHash :: AgentErrorType -> ByteString -> AM () @@ -4047,7 +4051,7 @@ processSMPTransmissions c@AgentClient {subQ} (tSess@(userId, srv, _), THandlePar -- this branch is executed by the accepting party in duplexHandshake mode (v2) -- (was executed by initiating party in v1 that is no longer supported) | sndStatus == Active -> do - atomically $ incSMPServerStat c userId srv connCompleted + incSMPServerStat c userId srv connCompleted notify $ CON pqEncryption | otherwise -> enqueueDuplexHello sq _ -> pure () diff --git a/src/Simplex/Messaging/Agent/Client.hs b/src/Simplex/Messaging/Agent/Client.hs index d892fcad0..5276b93cf 100644 --- a/src/Simplex/Messaging/Agent/Client.hs +++ b/src/Simplex/Messaging/Agent/Client.hs @@ -174,7 +174,6 @@ module Simplex.Messaging.Agent.Client incSMPServerStat, incSMPServerStat', incXFTPServerStat, - incXFTPServerStat', incXFTPServerSizeStat, incNtfServerStat, incNtfServerStat', @@ -408,15 +407,12 @@ getAgentWorker = getAgentWorker' id pure getAgentWorker' :: forall a k e m. (Ord k, Show k, AnyError e, MonadUnliftIO m) => (a -> Worker) -> (Worker -> STM a) -> String -> Bool -> AgentClient -> k -> TMap k a -> (a -> ExceptT e m ()) -> m a getAgentWorker' toW fromW name hasWork c@AgentClient {agentEnv} key ws work = do - atomically (getWorker >>= maybe createWorker whenExists) >>= \w -> runWorker w $> w + liftIO (TM.lookupOrInsertIO key createWorker ws) >>= either pure whenExists >>= \w -> runWorker w $> w where getWorker = TM.lookup key ws - createWorker = do - w <- fromW =<< newWorker c - TM.insert key w ws - pure w + createWorker = fromW =<< newWorker c whenExists w - | hasWork = hasWorkToDo (toW w) $> w + | hasWork = atomically $ hasWorkToDo (toW w) $> w | otherwise = pure w runWorker w = runWorkerAsync (toW w) runWork where @@ -696,12 +692,12 @@ getSMPProxyClient c@AgentClient {active, smpClients, smpProxiedRelays, workerSeq unlessM (readTVarIO active) $ throwE INACTIVE proxySrv <- maybe (getNextServer c userId proxySrvs [destSrv]) pure proxySrv_ ts <- liftIO getCurrentTime - (tSess, auth) <- atomically $ proxiedTransportSession proxySrv + (tSess, auth) <- liftIO $ proxiedTransportSession proxySrv withGetSessVar workerSeq tSess smpClients ts (newProxyClient tSess auth ts) (waitForProxyClient tSess auth) where - proxiedTransportSession :: SMPServerWithAuth -> STM (SMPTransportSession, Maybe SMP.BasicAuth) + proxiedTransportSession :: SMPServerWithAuth -> IO (SMPTransportSession, Maybe SMP.BasicAuth) proxiedTransportSession proxySrv = do - ProtoServerWithAuth srv auth <- TM.lookup destSess smpProxiedRelays >>= maybe (TM.insert destSess proxySrv smpProxiedRelays $> proxySrv) pure + ProtoServerWithAuth srv auth <- either id id <$> TM.lookupOrInsertIO destSess (pure proxySrv) smpProxiedRelays pure ((userId, srv, qId), auth) newProxyClient :: SMPTransportSession -> Maybe SMP.BasicAuth -> UTCTime -> SMPClientVar -> AM (SMPConnectedClient, Either AgentErrorType ProxiedRelay) newProxyClient tSess auth ts v = do @@ -1109,10 +1105,8 @@ withLocksMap :: (Ord k, MonadUnliftIO m) => TMap k Lock -> Set k -> Text -> m a withLocksMap = withGetLocks . getMapLock {-# INLINE withLocksMap #-} -getMapLock :: Ord k => TMap k Lock -> k -> STM Lock -getMapLock locks key = TM.lookup key locks >>= maybe newLock pure - where - newLock = createLock >>= \l -> TM.insert key l locks $> l +getMapLock :: Ord k => TMap k Lock -> k -> IO Lock +getMapLock locks key = either id id <$> TM.lookupOrInsertIO key createLock locks withClient_ :: forall a v err msg. ProtocolServerClient v err msg => AgentClient -> NetworkRequestMode -> TransportSession msg -> (Client msg -> AM a) -> AM a withClient_ c nm tSess@(_, srv, _) action = do @@ -1166,11 +1160,11 @@ sendOrProxySMPMessage c nm userId destSrv connId cmdStr spKey_ senderId msgFlags fst <$> sendOrProxySMPCommand c nm userId destSrv connId cmdStr senderId sendViaProxy sendDirectly where sendViaProxy smp proxySess = do - atomically $ incSMPServerStat c userId destSrv sentViaProxyAttempts - atomically $ incSMPServerStat c userId (protocolClientServer' smp) sentProxiedAttempts + incSMPServerStat c userId destSrv sentViaProxyAttempts + incSMPServerStat c userId (protocolClientServer' smp) sentProxiedAttempts proxySMPMessage smp nm proxySess spKey_ senderId msgFlags msg sendDirectly smp = do - atomically $ incSMPServerStat c userId destSrv sentDirectAttempts + incSMPServerStat c userId destSrv sentDirectAttempts sendSMPMessage smp nm spKey_ senderId msgFlags msg sendOrProxySMPCommand :: @@ -1246,8 +1240,8 @@ sendOrProxySMPCommand c nm userId destSrv@ProtocolServer {host = destHosts} conn sameProxiedRelay proxySess' = prSessionId proxySess == prSessionId proxySess' case r of Right r'@(srv_, _) -> do - atomically $ incSMPServerStat c userId destSrv sentViaProxy - forM_ srv_ $ \proxySrv -> atomically $ incSMPServerStat c userId proxySrv sentProxied + incSMPServerStat c userId destSrv sentViaProxy + forM_ srv_ $ \proxySrv -> incSMPServerStat c userId proxySrv sentProxied pure r' Left e | serverHostError e -> ifM directAllowed ((Nothing,) <$> sendDirectly destSess) (throwE e) @@ -1255,7 +1249,7 @@ sendOrProxySMPCommand c nm userId destSrv@ProtocolServer {host = destHosts} conn sendDirectly tSess = withLogClient_ c nm tSess (unEntityId entId) ("SEND " <> cmdStr) $ \(SMPConnectedClient smp _) -> do tryAllErrors (liftClient SMP (clientServer smp) $ sendCmdDirectly smp) >>= \case - Right r -> r <$ atomically (incSMPServerStat c userId destSrv sentDirect) + Right r -> r <$ incSMPServerStat c userId destSrv sentDirect Left e -> throwE e ipAddressProtected :: NetworkConfig -> ProtocolServer p -> Bool @@ -1561,18 +1555,15 @@ newRcvQueue_ c nm userId connId (ProtoServerWithAuth srv auth) vRange cqrd enabl newErr :: String -> AM (Maybe ShortLinkCreds) newErr = throwE . BROKER (B.unpack $ strEncode srv) . UNEXPECTED . ("Create queue: " <>) -processSubResults :: AgentClient -> SMPTransportSession -> SessionId -> Maybe ServiceId -> NonEmpty (RcvQueueSub, Either SMPClientError (Maybe ServiceId)) -> STM ([RcvQueueSub], [(RcvQueueSub, Maybe ClientNotice)]) -processSubResults c tSess@(userId, srv, _) sessId serviceId_ rs = do +processSubResults :: AgentClient -> SMPTransportSession -> SessionId -> Maybe ServiceId -> NonEmpty (RcvQueueSub, Either SMPClientError (Maybe ServiceId)) -> STM ([RcvQueueSub], [(RcvQueueSub, Maybe ClientNotice)], [(AgentSMPServerStats -> TVar Int, Int)]) +processSubResults c tSess sessId serviceId_ rs = do pendingSubs <- SS.getPendingQueueSubs tSess $ currentSubs c let (failed, subscribed@(qs, sQs), notices, ignored) = foldr (partitionResults pendingSubs) (M.empty, ([], []), [], 0) rs - unless (M.null failed) $ do - incSMPServerStat' c userId srv connSubErrs $ M.size failed + unless (M.null failed) $ failSubscriptions c tSess failed - unless (null qs && null sQs) $ do - incSMPServerStat' c userId srv connSubscribed $ length qs + length sQs + unless (null qs && null sQs) $ SS.batchAddActiveSubs tSess sessId serviceId_ subscribed $ currentSubs c - unless (ignored == 0) $ incSMPServerStat' c userId srv connSubIgnored ignored - pure (sQs, notices) + pure (sQs, notices, [(connSubErrs, M.size failed), (connSubscribed, length qs + length sQs), (connSubIgnored, ignored)]) where partitionResults :: Map SMP.RecipientId RcvQueueSub -> @@ -1711,20 +1702,21 @@ subscribeSessQueues_ c withEvents qs = sendClientBatch_ "SUB" False subscribe_ c subscribe_ :: SMPClient -> NonEmpty RcvQueueSub -> IO (BatchResponses RcvQueueSub SMPClientError (Maybe ServiceId), Bool) subscribe_ smp qs' = do let (userId, srv, _) = tSess - atomically $ incSMPServerStat' c userId srv connSubAttempts $ length qs' + incSMPServerStat' c userId srv connSubAttempts $ length qs' rs <- sendBatch (\smp' _ -> subscribeSMPQueues smp') smp NRMBackground qs' cs_ <- if withEvents then Just . S.fromList . map qConnId . M.elems <$> atomically (SS.getActiveSubs tSess $ currentSubs c) else pure Nothing active <- E.uninterruptibleMask_ $ do - (active, (serviceQs, notices)) <- atomically $ do - r@(_, (_, notices)) <- ifM + (active, (serviceQs, notices, stats)) <- atomically $ do + r@(_, (_, notices, _)) <- ifM (activeClientSession c tSess sessId) ((True,) <$> processSubResults c tSess sessId (smpClientServiceId smp) rs) - ((False, ([], [])) <$ incSMPServerStat' c userId srv connSubIgnored (length rs)) + (pure (False, ([], [], [(connSubIgnored, length rs)]))) unless (null notices) $ takeTMVar $ clientNoticesLock c pure r + forM_ stats $ \(sel, n) -> unless (n == 0) $ incSMPServerStat' c userId srv sel n unless (null serviceQs) $ void $ processRcvServiceAssocs c srv serviceQs `runReaderT` agentEnv c unless (null notices) $ void $ @@ -2116,10 +2108,10 @@ deleteQueues c nm = sendTSessionBatches "DEL" mkSMPTSession deleteQueues_ c nm where deleteQueues_ smp rqs = do let (userId, srv, _) = transportSession' smp - atomically $ incSMPServerStat' c userId srv connDelAttempts $ length rqs + incSMPServerStat' c userId srv connDelAttempts $ length rqs rs <- sendBatch deleteSMPQueues smp nm rqs let successes = foldl' (\n (_, r) -> if isRight r then n + 1 else n) 0 rs - atomically $ incSMPServerStat' c userId srv connDeleted successes + incSMPServerStat' c userId srv connDeleted successes pure rs -- This is only used in background @@ -2573,44 +2565,36 @@ withNextSrv c userId srvsSel triedHosts usedSrvs action = do atomically $ writeTVar triedHosts tried' action srvAuth -incSMPServerStat :: AgentClient -> UserId -> SMPServer -> (AgentSMPServerStats -> TVar Int) -> STM () +incSMPServerStat :: MonadIO m => AgentClient -> UserId -> SMPServer -> (AgentSMPServerStats -> TVar Int) -> m () incSMPServerStat c userId srv sel = incSMPServerStat' c userId srv sel 1 -incSMPServerStat' :: AgentClient -> UserId -> SMPServer -> (AgentSMPServerStats -> TVar Int) -> Int -> STM () +incSMPServerStat' :: MonadIO m => AgentClient -> UserId -> SMPServer -> (AgentSMPServerStats -> TVar Int) -> Int -> m () incSMPServerStat' = incServerStat (\AgentClient {smpServersStats = s} -> s) newAgentSMPServerStats -incXFTPServerStat :: AgentClient -> UserId -> XFTPServer -> (AgentXFTPServerStats -> TVar Int) -> STM () +incXFTPServerStat :: MonadIO m => AgentClient -> UserId -> XFTPServer -> (AgentXFTPServerStats -> TVar Int) -> m () incXFTPServerStat c userId srv sel = incXFTPServerStat_ c userId srv sel 1 {-# INLINE incXFTPServerStat #-} -incXFTPServerStat' :: AgentClient -> UserId -> XFTPServer -> (AgentXFTPServerStats -> TVar Int) -> Int -> STM () -incXFTPServerStat' = incXFTPServerStat_ -{-# INLINE incXFTPServerStat' #-} - -incXFTPServerSizeStat :: AgentClient -> UserId -> XFTPServer -> (AgentXFTPServerStats -> TVar Int64) -> Int64 -> STM () +incXFTPServerSizeStat :: MonadIO m => AgentClient -> UserId -> XFTPServer -> (AgentXFTPServerStats -> TVar Int64) -> Int64 -> m () incXFTPServerSizeStat = incXFTPServerStat_ {-# INLINE incXFTPServerSizeStat #-} -incXFTPServerStat_ :: Num n => AgentClient -> UserId -> XFTPServer -> (AgentXFTPServerStats -> TVar n) -> n -> STM () +incXFTPServerStat_ :: (MonadIO m, Num n) => AgentClient -> UserId -> XFTPServer -> (AgentXFTPServerStats -> TVar n) -> n -> m () incXFTPServerStat_ = incServerStat (\AgentClient {xftpServersStats = s} -> s) newAgentXFTPServerStats {-# INLINE incXFTPServerStat_ #-} -incNtfServerStat :: AgentClient -> UserId -> NtfServer -> (AgentNtfServerStats -> TVar Int) -> STM () +incNtfServerStat :: MonadIO m => AgentClient -> UserId -> NtfServer -> (AgentNtfServerStats -> TVar Int) -> m () incNtfServerStat c userId srv sel = incNtfServerStat' c userId srv sel 1 {-# INLINE incNtfServerStat #-} -incNtfServerStat' :: AgentClient -> UserId -> NtfServer -> (AgentNtfServerStats -> TVar Int) -> Int -> STM () +incNtfServerStat' :: MonadIO m => AgentClient -> UserId -> NtfServer -> (AgentNtfServerStats -> TVar Int) -> Int -> m () incNtfServerStat' = incServerStat (\AgentClient {ntfServersStats = s} -> s) newAgentNtfServerStats {-# INLINE incNtfServerStat' #-} -incServerStat :: Num n => (AgentClient -> TMap (UserId, ProtocolServer p) s) -> STM s -> AgentClient -> UserId -> ProtocolServer p -> (s -> TVar n) -> n -> STM () -incServerStat statsSel mkNewStats c userId srv sel n = do - TM.lookup (userId, srv) (statsSel c) >>= \case - Just v -> modifyTVar' (sel v) (+ n) - Nothing -> do - newStats <- mkNewStats - modifyTVar' (sel newStats) (+ n) - TM.insert (userId, srv) newStats (statsSel c) +incServerStat :: (MonadIO m, Num n) => (AgentClient -> TMap (UserId, ProtocolServer p) s) -> STM s -> AgentClient -> UserId -> ProtocolServer p -> (s -> TVar n) -> n -> m () +incServerStat statsSel mkNewStats c userId srv sel n = liftIO $ do + v <- either id id <$> TM.lookupOrInsertIO (userId, srv) mkNewStats (statsSel c) + atomically $ modifyTVar' (sel v) (+ n) data AgentServersSummary = AgentServersSummary { smpServersStats :: Map (UserId, SMPServer) AgentSMPServerStatsData, diff --git a/src/Simplex/Messaging/Agent/Lock.hs b/src/Simplex/Messaging/Agent/Lock.hs index 081c53b68..0068911f0 100644 --- a/src/Simplex/Messaging/Agent/Lock.hs +++ b/src/Simplex/Messaging/Agent/Lock.hs @@ -40,20 +40,18 @@ withLock' lock name = (atomically $ putTMVar lock name) (void . atomically $ takeTMVar lock) -withGetLock :: MonadUnliftIO m => (k -> STM Lock) -> k -> Text -> m a -> m a +withGetLock :: MonadUnliftIO m => (k -> IO Lock) -> k -> Text -> m a -> m a withGetLock getLock key name a = E.bracket - (atomically $ getPutLock getLock key name) + (liftIO $ getPutLock getLock key name) (atomically . takeTMVar) (const a) -withGetLocks :: MonadUnliftIO m => (k -> STM Lock) -> Set k -> Text -> m a -> m a +withGetLocks :: MonadUnliftIO m => (k -> IO Lock) -> Set k -> Text -> m a -> m a withGetLocks getLock keys name = E.bracket holdLocks releaseLocks . const where - holdLocks = forConcurrently (S.toList keys) $ \key -> atomically $ getPutLock getLock key name + holdLocks = forConcurrently (S.toList keys) $ \key -> liftIO $ getPutLock getLock key name releaseLocks = mapM_ (atomically . takeTMVar) --- getLock and putTMVar can be in one transaction on the assumption that getLock doesn't write in case the lock already exists, --- and in case it is created and added to some shared resource (we use TMap) it also helps avoid contention for the newly created lock. -getPutLock :: (k -> STM Lock) -> k -> Text -> STM Lock -getPutLock getLock key name = getLock key >>= \l -> putTMVar l name $> l +getPutLock :: (k -> IO Lock) -> k -> Text -> IO Lock +getPutLock getLock key name = getLock key >>= \l -> atomically (putTMVar l name) $> l diff --git a/src/Simplex/Messaging/Agent/NtfSubSupervisor.hs b/src/Simplex/Messaging/Agent/NtfSubSupervisor.hs index 7153abcbf..b86fbd5eb 100644 --- a/src/Simplex/Messaging/Agent/NtfSubSupervisor.hs +++ b/src/Simplex/Messaging/Agent/NtfSubSupervisor.hs @@ -323,7 +323,7 @@ runNtfWorker c srv Worker {doWork} = incStatByUserId :: NtfServer -> (AgentNtfServerStats -> TVar Int) -> [NtfSubscription] -> AM' () incStatByUserId ntfServer sel ss = forM_ (M.assocs userIdsCounts) $ \(userId, count) -> - atomically $ incNtfServerStat' c userId ntfServer sel count + incNtfServerStat' c userId ntfServer sel count where userIdsCounts = foldl' (\acc NtfSubscription {userId} -> M.insertWith (+) userId 1 acc) M.empty ss -- NSADelete and NSARotate are deprecated, but their processing is kept for legacy db records; @@ -361,10 +361,10 @@ runNtfWorker c srv Worker {doWork} = Just nSubId -> lift getNtfToken >>= \case Just tkn@NtfToken {ntfServer} -> do - atomically $ incNtfServerStat c userId ntfServer ntfDelAttempts + incNtfServerStat c userId ntfServer ntfDelAttempts tryAllErrors (agentNtfDeleteSubscription c nSubId tkn) >>= \case Right _ -> do - atomically $ incNtfServerStat c userId ntfServer ntfDeleted + incNtfServerStat c userId ntfServer ntfDeleted continue' Left e | temporaryOrHostError e -> pure $ Just sub -- don't continue, retry diff --git a/src/Simplex/Messaging/Client/Agent.hs b/src/Simplex/Messaging/Client/Agent.hs index 0d59f00b8..8020e5251 100644 --- a/src/Simplex/Messaging/Client/Agent.hs +++ b/src/Simplex/Messaging/Client/Agent.hs @@ -420,7 +420,7 @@ subscribeQueuesNtfs = subscribeQueues_ subscribeQueues_ :: ServiceParty p => SMPClientAgent p -> SMPServer -> NonEmpty (QueueId, C.APrivateAuthKey) -> IO () subscribeQueues_ ca srv subs = do - atomically $ addPendingSubs ca srv $ L.toList subs + addPendingSubs ca srv $ L.toList subs runExceptT (getSMPServerClient' ca srv) >>= \case Right smp -> smpSubscribeQueues ca smp srv subs Left _ -> pure () -- no call to reconnectClient - failing getSMPServerClient' does that @@ -534,8 +534,10 @@ addActiveSubs :: SMPClientAgent p -> SMPServer -> [(QueueId, (SessionId, C.APriv addActiveSubs = addSubsList_ . activeQueueSubs {-# INLINE addActiveSubs #-} -addPendingSubs :: SMPClientAgent p -> SMPServer -> [(QueueId, C.APrivateAuthKey)] -> STM () -addPendingSubs = addSubsList_ . pendingQueueSubs +addPendingSubs :: SMPClientAgent p -> SMPServer -> [(QueueId, C.APrivateAuthKey)] -> IO () +addPendingSubs ca srv ss = TM.lookupOrInsertIO srv (newTVar ss') (pendingQueueSubs ca) >>= mapM_ (atomically . TM.union ss') + where + ss' = M.fromList ss {-# INLINE addPendingSubs #-} addSubsList_ :: TMap SMPServer (TMap QueueId s) -> SMPServer -> [(QueueId, s)] -> STM () diff --git a/src/Simplex/Messaging/Notifications/Server.hs b/src/Simplex/Messaging/Notifications/Server.hs index 51687abdd..735774394 100644 --- a/src/Simplex/Messaging/Notifications/Server.hs +++ b/src/Simplex/Messaging/Notifications/Server.hs @@ -495,7 +495,7 @@ subscribeNtfs NtfSubscriber {smpSubscribers, subscriberSeq, smpAgent = ca} st sm >>= mapM_ (\sub -> atomically $ writeTQueue (subscriberSubQ sub) ntfSub) where getSubscriberVar :: IO (Either SMPSubscriberVar SMPSubscriberVar) - getSubscriberVar = atomically . getSessVar subscriberSeq smpServer smpSubscribers =<< getCurrentTime + getSubscriberVar = getSessVarIO subscriberSeq smpServer smpSubscribers =<< getCurrentTime createSMPSubscriber :: SMPSubscriberVar -> IO (Maybe SMPSubscriber) createSMPSubscriber v = diff --git a/src/Simplex/Messaging/Notifications/Server/Env.hs b/src/Simplex/Messaging/Notifications/Server/Env.hs index d7f772a87..1ad6e09b6 100644 --- a/src/Simplex/Messaging/Notifications/Server/Env.hs +++ b/src/Simplex/Messaging/Notifications/Server/Env.hs @@ -207,7 +207,7 @@ getPushClient :: NtfPushServer -> PushProvider -> IO (PushProviderClient, PushCl getPushClient s@NtfPushServer {apnsConfig = APNSPushClientConfig {reconnectInterval}} pp = withRetryIntervalCount reconnectInterval $ \n _delay loop -> do ts <- getCurrentTime - E.try (atomically (getSessVar (pushClientSeq s) pp (pushClients s) ts) >>= either (newPushClient s pp) waitForPushClient) >>= \case + E.try (getSessVarIO (pushClientSeq s) pp (pushClients s) ts >>= either (newPushClient s pp) waitForPushClient) >>= \case Right result -> pure result Left e | n < 2 -> do diff --git a/src/Simplex/Messaging/Notifications/Server/Stats.hs b/src/Simplex/Messaging/Notifications/Server/Stats.hs index 18a746a8d..b23fc6301 100644 --- a/src/Simplex/Messaging/Notifications/Server/Stats.hs +++ b/src/Simplex/Messaging/Notifications/Server/Stats.hs @@ -338,8 +338,5 @@ getStatsByServer s = readTVarIO s >>= fmap (StatsByServerData . M.toList) . mapM setStatsByServer :: TMap Text (TVar Int) -> StatsByServerData -> IO () setStatsByServer s (StatsByServerData d) = mapM newTVarIO (M.fromList d) >>= atomically . writeTVar s --- double lookup avoids STM transaction with a shared map in most cases incServerStat :: Text -> TMap Text (TVar Int) -> IO () -incServerStat h s = TM.lookupIO h s >>= atomically . maybe newServerStat (`modifyTVar'` (+ 1)) - where - newServerStat = TM.lookup h s >>= maybe (TM.insertM h (newTVar 1) s) (`modifyTVar'` (+ 1)) +incServerStat h s = TM.lookupOrInsertIO h (newTVar 1) s >>= mapM_ (atomically . (`modifyTVar'` (+ 1))) diff --git a/src/Simplex/Messaging/Server.hs b/src/Simplex/Messaging/Server.hs index f38f0f829..30133885d 100644 --- a/src/Simplex/Messaging/Server.hs +++ b/src/Simplex/Messaging/Server.hs @@ -1768,20 +1768,14 @@ client fmap (either ((,Nothing) . err) id) $ liftIO $ runExceptT $ do msg' <- forM msg_ $ \msg -> liftIO $ do ts <- getSystemSeconds - sub <- maybe (atomically getSub) pure sub_ + sub <- maybe getSub pure sub_ atomically $ setDelivered sub msg ts unless hasSub $ incStat $ qSub stats pure (NoCorrId, entId, MSG (encryptMsg qr msg)) pure ((corrId, entId, SOK clntServiceId), msg') - getSub :: STM Sub - getSub = - TM.lookup entId (subscriptions clnt) >>= \case - Just sub -> pure sub - Nothing -> do - sub <- newSubscription NoSub - TM.insert entId sub $ subscriptions clnt - pure sub + getSub :: IO Sub + getSub = either id id <$> TM.lookupOrInsertIO entId (newSubscription NoSub) (subscriptions clnt) subscribeNewQueue :: RecipientId -> QueueRec -> M s () subscribeNewQueue rId QueueRec {rcvServiceId} = do @@ -1795,10 +1789,13 @@ client -- clients that use GET are not added to server subscribers getMessage :: StoreQueue s -> QueueRec -> M s (Transmission BrokerMsg) getMessage q qr = do - atomically (TM.lookup entId $ subscriptions clnt) >>= \case - Nothing -> - atomically newSub >>= (`getMessage_` Nothing) - Just s@Sub {subThread} -> + -- Here we don't account for this client as subscribed in the server + -- and don't notify other subscribed clients. + -- This is tracked as "subscription" in the client to prevent these + -- clients from being able to subscribe. + liftIO (TM.lookupOrInsertIO entId newProhibitedSub $ subscriptions clnt) >>= \case + Left s -> getMessage_ s Nothing + Right s@Sub {subThread} -> case subThread of ProhibitSub -> atomically (swapTVar (delivered s) Nothing) @@ -1809,15 +1806,6 @@ client incStat $ msgGetProhibited stats pure $ err $ CMD PROHIBITED where - newSub :: STM Sub - newSub = do - s <- newProhibitedSub - TM.insert entId s $ subscriptions clnt - -- Here we don't account for this client as subscribed in the server - -- and don't notify other subscribed clients. - -- This is tracked as "subscription" in the client to prevent these - -- clients from being able to subscribe. - pure s getMessage_ :: Sub -> Maybe (MsgId, SystemSeconds) -> M s (Transmission BrokerMsg) getMessage_ s delivered_ = do stats <- asks serverStats @@ -1897,23 +1885,17 @@ client Just _ -> do -- unassociation already done in prepareBatchSubs incSrvStat srvAssocRemoved - -- getSubscription may be Just for receiving service, where clientSubs also hold active deliveries for service subscriptions. - -- For notification service it can only be Just if storage and session states diverge. - r <- atomically $ getSubscription >>= newSub + -- getSubscription may find a subscription for receiving service, where clientSubs also hold active deliveries for service subscriptions. + -- For notification service it can only find one if storage and session states diverge. + r <- getSubscription atomically writeSub pure r Nothing -> do - r@(hasSub, _) <- atomically $ getSubscription >>= newSub + r@(hasSub, _) <- getSubscription unless hasSub $ atomically writeSub pure r where - getSubscription = TM.lookup entId $ clientSubs clnt - newSub = \case - Just sub -> pure (True, Just sub) - Nothing -> do - sub <- mkSub - TM.insert entId sub $ clientSubs clnt - pure (False, Just sub) + getSubscription = either ((False,) . Just) ((True,) . Just) <$> TM.lookupOrInsertIO entId mkSub (clientSubs clnt) subscribeServiceMessages :: ServiceId -> (Int64, IdsHash) -> M s BrokerMsg subscribeServiceMessages serviceId expected = @@ -1943,22 +1925,15 @@ client Right qMsg_ -> case qMsg_ of Nothing -> pure (qCnt + 1, msgCnt, dupCnt, evts) Just (qr, msg) -> - atomically (getSubscription rId) >>= \case - Nothing -> pure (qCnt + 1, msgCnt, dupCnt + 1, evts) - Just sub -> do + getSubscription rId >>= \case + -- delivered via the existing subscription, e.g. made by SEND after SUBS + Right _sub -> pure (qCnt + 1, msgCnt, dupCnt + 1, evts) + Left sub -> do ts <- getSystemSeconds atomically $ setDelivered sub msg ts atomically $ writeTBQueue msgQ [(NoCorrId, rId, MSG (encryptMsg qr msg))] pure (qCnt + 1, msgCnt + 1, dupCnt, evts) - getSubscription rId = - TM.lookup rId (subscriptions clnt) >>= \case - -- If delivery subscription already exists, then there is no need to deliver message. - -- It may have been created when the message is sent after service subscription is created. - Just _sub -> pure Nothing - Nothing -> do - sub <- newSubscription NoSub - TM.insert rId sub $ subscriptions clnt - pure $ Just sub + getSubscription rId = TM.lookupOrInsertIO rId (newSubscription NoSub) (subscriptions clnt) subscribeServiceNotifications :: ServiceId -> (Int64, IdsHash) -> M s BrokerMsg subscribeServiceNotifications serviceId expected = diff --git a/src/Simplex/Messaging/Session.hs b/src/Simplex/Messaging/Session.hs index 16ff3c934..6857e0bde 100644 --- a/src/Simplex/Messaging/Session.hs +++ b/src/Simplex/Messaging/Session.hs @@ -5,6 +5,7 @@ module Simplex.Messaging.Session ( SessionVar (..), getSessVar, + getSessVarIO, removeSessVar, withGetSessVar, withGetSessVar', @@ -32,12 +33,19 @@ getSessVar sessSeq sessKey vs sessionVarTs = maybe (Left <$> newSessionVar) (pur where newSessionVar :: STM (SessionVar a) newSessionVar = do - sessionVar <- newEmptyTMVar - sessionVarId <- stateTVar sessSeq $ \next -> (next, next + 1) - let v = SessionVar {sessionVar, sessionVarId, sessionVarTs} + v <- createSessionVar sessSeq sessionVarTs TM.insert sessKey v vs pure v +getSessVarIO :: Ord k => TVar Int -> k -> TMap k (SessionVar a) -> UTCTime -> IO (Either (SessionVar a) (SessionVar a)) +getSessVarIO sessSeq sessKey vs ts = TM.lookupOrInsertIO sessKey (createSessionVar sessSeq ts) vs + +createSessionVar :: TVar Int -> UTCTime -> STM (SessionVar a) +createSessionVar sessSeq sessionVarTs = do + sessionVar <- newEmptyTMVar + sessionVarId <- stateTVar sessSeq $ \next -> (next, next + 1) + pure SessionVar {sessionVar, sessionVarId, sessionVarTs} + removeSessVar :: Ord k => SessionVar a -> k -> TMap k (SessionVar a) -> STM () removeSessVar v sessKey vs = TM.lookup sessKey vs >>= \case @@ -64,7 +72,7 @@ withGetSessVar' :: (SessionVar a -> m b) -> (SessionVar a -> m b) -> m b withGetSessVar' sessSeq sessKey vs ts onNew onExisting = bracketOnError - (liftIO $ atomically $ getSessVar sessSeq sessKey vs ts) + (liftIO $ getSessVarIO sessSeq sessKey vs ts) (either (liftIO . atomically . dropEmptySessVar) (\_ -> pure ())) (either onNew onExisting) where diff --git a/src/Simplex/Messaging/TMap.hs b/src/Simplex/Messaging/TMap.hs index b09158d5f..b8d447118 100644 --- a/src/Simplex/Messaging/TMap.hs +++ b/src/Simplex/Messaging/TMap.hs @@ -14,6 +14,7 @@ module Simplex.Messaging.TMap insertM, delete, lookupInsert, + lookupOrInsertIO, lookupDelete, adjust, update, @@ -77,6 +78,13 @@ lookupInsert :: Ord k => k -> a -> TMap k a -> STM (Maybe a) lookupInsert k v m = stateTVar m $ M.alterF (,Just v) k {-# INLINE lookupInsert #-} +-- double lookup avoids STM transaction with a shared map in most cases +lookupOrInsertIO :: Ord k => k -> STM a -> TMap k a -> IO (Either a a) +lookupOrInsertIO k new m = lookupIO k m >>= maybe (atomically lookupOrInsert) (pure . Right) + where + lookupOrInsert = M.lookup k <$> readTVar m >>= maybe (new >>= \v -> Left v <$ insert k v m) (pure . Right) +{-# INLINE lookupOrInsertIO #-} + lookupDelete :: Ord k => k -> TMap k a -> STM (Maybe a) lookupDelete k m = stateTVar m $ M.alterF (,Nothing) k {-# INLINE lookupDelete #-}