export lock function

src/Simplex/Messaging/Agent/Client.hs

@@ -785,24 +785,13 @@
 
 withConnLock :: MonadUnliftIO m => AgentClient -> ConnId -> String -> m a -> m a
 withConnLock _ "" _ = id
-withConnLock AgentClient {connLocks} connId name = withLockMap_ connLocks connId name
+withConnLock AgentClient {connLocks} connId name = withLockMap connLocks connId name
 
 withInvLock :: MonadUnliftIO m => AgentClient -> ByteString -> String -> m a -> m a
-withInvLock AgentClient {invLocks} = withLockMap_ invLocks
+withInvLock AgentClient {invLocks} = withLockMap invLocks
 
 withConnLocks :: MonadUnliftIO m => AgentClient -> [ConnId] -> String -> m a -> m a
-withConnLocks AgentClient {connLocks} = withLocksMap_ connLocks . filter (not . B.null)
-
-withLockMap_ :: (Ord k, MonadUnliftIO m) => TMap k Lock -> k -> String -> m a -> m a
-withLockMap_ = withGetLock . getMapLock
-
-withLocksMap_ :: (Ord k, MonadUnliftIO m) => TMap k Lock -> [k] -> String -> m a -> m a
-withLocksMap_ = withGetLocks . getMapLock
-
-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
+withConnLocks AgentClient {connLocks} = withLocksMap connLocks . filter (not . B.null)
 
 withClient_ :: forall a m v err msg. (AgentMonad m, ProtocolServerClient v err msg) => AgentClient -> TransportSession msg -> ByteString -> (Client msg -> m a) -> m a
 withClient_ c tSess@(userId, srv, _) statCmd action = do
@@ -1519,7 +1508,7 @@
   where
     statsKey = AgentStatsKey {userId, host = strEncode $ clientTransportHost pc, clientTs = strEncode $ clientSessionTs pc, cmd, res}
 
 userServers :: forall p. (ProtocolTypeI p, UserProtocol p) => AgentClient -> TMap UserId (NonEmpty (ProtoServerWithAuth p))
 userServers c = case protocolTypeI @p of
   SPSMP -> smpServers c
   SPXFTP -> xftpServers c

src/Simplex/Messaging/Agent/Lock.hs

@@ -4,14 +4,20 @@ module Simplex.Messaging.Agent.Lock
   ( Lock,
     createLock,
     withLock,
+    waitForLock,
     withGetLock,
     withGetLocks,
+    withLockMap,
+    withLocksMap,
   )
 where
 
-import Control.Monad (void)
+import Control.Concurrent.STM (retry)
+import Control.Monad (unless, void)
 import Control.Monad.IO.Unlift
 import Data.Functor (($>))
+import Simplex.Messaging.TMap (TMap)
+import qualified Simplex.Messaging.TMap as TM
 import UnliftIO.Async (forConcurrently)
 import qualified UnliftIO.Exception as E
 import UnliftIO.STM
@@ -28,6 +34,9 @@ withLock lock name =
     (atomically $ putTMVar lock name)
     (void . atomically $ takeTMVar lock)
 
+waitForLock :: Lock -> STM ()
+waitForLock lock = isEmptyTMVar lock >>= (`unless` retry)
+
 withGetLock :: MonadUnliftIO m => (k -> STM Lock) -> k -> String -> m a -> m a
 withGetLock getLock key name a =
   E.bracket
@@ -47,3 +56,14 @@ withGetLocks getLock keys name = E.bracket holdLocks releaseLocks . const
 -- 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 -> String -> STM Lock
 getPutLock getLock key name = getLock key >>= \l -> putTMVar l name $> l
+
+withLockMap :: (Ord k, MonadUnliftIO m) => TMap k Lock -> k -> String -> m a -> m a
+withLockMap = withGetLock . getMapLock
+
+withLocksMap :: (Ord k, MonadUnliftIO m) => TMap k Lock -> [k] -> String -> m a -> m a
+withLocksMap = withGetLocks . getMapLock
+
+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