refactor STM queues (#1447)

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

@@ -26,22 +26,18 @@ import Simplex.Messaging.Server.MsgStore.Types
 import Simplex.Messaging.Server.QueueStore
 import Simplex.Messaging.Server.QueueStore.STM
 import Simplex.Messaging.Server.StoreLog
-import Simplex.Messaging.TMap (TMap)
-import qualified Simplex.Messaging.TMap as TM
 import Simplex.Messaging.Util ((<$$>), ($>>=))
 import System.IO (IOMode (..))
 
 data STMMsgStore = STMMsgStore
   { storeConfig :: STMStoreConfig,
-    queues :: TMap RecipientId STMQueue,
-    senders :: TMap SenderId RecipientId,
-    notifiers :: TMap NotifierId RecipientId,
-    storeLog :: TVar (Maybe (StoreLog 'WriteMode))
+    queueStore :: STMQueueStore STMQueue
   }
 
 data STMQueue = STMQueue
   { -- To avoid race conditions and errors when restoring queues,
     -- Nothing is written to TVar when queue is deleted.
+    recipientId :: RecipientId,
     queueRec :: TVar (Maybe QueueRec),
     msgQueue_ :: TVar (Maybe STMMsgQueue)
   }
@@ -57,12 +53,9 @@ data STMStoreConfig = STMStoreConfig
     quota :: Int
   }
 
-instance STMQueueStore STMMsgStore where
-  queues' = queues
-  senders' = senders
-  notifiers' = notifiers
-  storeLog' = storeLog
-  mkQueue _ qr = STMQueue <$> newTVar (Just qr) <*> newTVar Nothing
+instance STMStoreClass STMMsgStore where
+  stmQueueStore = queueStore
+  mkQueue _ rId qr = STMQueue rId <$> newTVar (Just qr) <*> newTVar Nothing
   msgQueue_' = msgQueue_
 
 instance MsgStoreClass STMMsgStore where
@@ -73,32 +66,31 @@ instance MsgStoreClass STMMsgStore where
 
   newMsgStore :: STMStoreConfig -> IO STMMsgStore
   newMsgStore storeConfig = do
-    queues <- TM.emptyIO
-    senders <- TM.emptyIO
-    notifiers <- TM.emptyIO
-    storeLog <- newTVarIO Nothing
-    pure STMMsgStore {storeConfig, queues, senders, notifiers, storeLog}
+    queueStore <- newQueueStore
+    pure STMMsgStore {storeConfig, queueStore}
 
   setStoreLog :: STMMsgStore -> StoreLog 'WriteMode -> IO ()
-  setStoreLog st sl = atomically $ writeTVar (storeLog st) (Just sl)
+  setStoreLog st sl = atomically $ writeTVar (storeLog $ queueStore st) (Just sl)
 
-  closeMsgStore st = readTVarIO (storeLog st) >>= mapM_ closeStoreLog
-
-  activeMsgQueues = queues
-  {-# INLINE activeMsgQueues #-}
+  closeMsgStore st = readTVarIO (storeLog $ queueStore st) >>= mapM_ closeStoreLog
 
   withAllMsgQueues _ = withActiveMsgQueues
   {-# INLINE withAllMsgQueues #-}
 
   logQueueStates _ = pure ()
+  {-# INLINE logQueueStates #-}
 
   logQueueState _ = pure ()
+  {-# INLINE logQueueState #-}
+
+  recipientId' = recipientId
+  {-# INLINE recipientId' #-}
 
   queueRec' = queueRec
   {-# INLINE queueRec' #-}
 
-  getMsgQueue :: STMMsgStore -> RecipientId -> STMQueue -> STM STMMsgQueue
-  getMsgQueue _ _ STMQueue {msgQueue_} = readTVar msgQueue_ >>= maybe newQ pure
+  getMsgQueue :: STMMsgStore -> STMQueue -> STM STMMsgQueue
+  getMsgQueue _ STMQueue {msgQueue_} = readTVar msgQueue_ >>= maybe newQ pure
     where
       newQ = do
         msgQueue <- newTQueue
@@ -108,23 +100,23 @@ instance MsgStoreClass STMMsgStore where
         writeTVar msgQueue_ (Just q)
         pure q
 
-  getPeekMsgQueue :: STMMsgStore -> RecipientId -> STMQueue -> STM (Maybe (STMMsgQueue, Message))
-  getPeekMsgQueue _ _ q@STMQueue {msgQueue_} = readTVar msgQueue_ $>>= \mq -> (mq,) <$$> tryPeekMsg_ q mq
+  getPeekMsgQueue :: STMMsgStore -> STMQueue -> STM (Maybe (STMMsgQueue, Message))
+  getPeekMsgQueue _ q@STMQueue {msgQueue_} = readTVar msgQueue_ $>>= \mq -> (mq,) <$$> tryPeekMsg_ q mq
 
   -- does not create queue if it does not exist, does not delete it if it does (can't just close in-memory queue)
-  withIdleMsgQueue :: Int64 -> STMMsgStore -> RecipientId -> STMQueue -> (STMMsgQueue -> STM a) -> STM (Maybe a, Int)
-  withIdleMsgQueue _ _ _ STMQueue {msgQueue_} action = readTVar msgQueue_ >>= \case
+  withIdleMsgQueue :: Int64 -> STMMsgStore -> STMQueue -> (STMMsgQueue -> STM a) -> STM (Maybe a, Int)
+  withIdleMsgQueue _ _ STMQueue {msgQueue_} action = readTVar msgQueue_ >>= \case
     Just q -> do
       r <- action q
       sz <- getQueueSize_ q
       pure (Just r, sz)
     Nothing -> pure (Nothing, 0)
 
-  deleteQueue :: STMMsgStore -> RecipientId -> STMQueue -> IO (Either ErrorType QueueRec)
-  deleteQueue ms rId q = fst <$$> deleteQueue' ms rId q
+  deleteQueue :: STMMsgStore -> STMQueue -> IO (Either ErrorType QueueRec)
+  deleteQueue ms q = fst <$$> deleteQueue' ms q
 
-  deleteQueueSize :: STMMsgStore -> RecipientId -> STMQueue -> IO (Either ErrorType (QueueRec, Int))
-  deleteQueueSize ms rId q = deleteQueue' ms rId q >>= mapM (traverse getSize)
+  deleteQueueSize :: STMMsgStore -> STMQueue -> IO (Either ErrorType (QueueRec, Int))
+  deleteQueueSize ms q = deleteQueue' ms q >>= mapM (traverse getSize)
     -- traverse operates on the second tuple element
     where
       getSize = maybe (pure 0) (\STMMsgQueue {size} -> readTVarIO size)
@@ -137,9 +129,9 @@ instance MsgStoreClass STMMsgStore where
         mapM_ (writeTQueue q) msgs
         pure msgs
 
-  writeMsg :: STMMsgStore -> RecipientId -> STMQueue -> Bool -> Message -> ExceptT ErrorType IO (Maybe (Message, Bool))
-  writeMsg ms rId q' _logState msg = liftIO $ atomically $ do
-    STMMsgQueue {msgQueue = q, canWrite, size} <- getMsgQueue ms rId q'
+  writeMsg :: STMMsgStore -> STMQueue -> Bool -> Message -> ExceptT ErrorType IO (Maybe (Message, Bool))
+  writeMsg ms q' _logState msg = liftIO $ atomically $ do
+    STMMsgQueue {msgQueue = q, canWrite, size} <- getMsgQueue ms q'
     canWrt <- readTVar canWrite
     empty <- isEmptyTQueue q
     if canWrt || empty
@@ -171,5 +163,5 @@ instance MsgStoreClass STMMsgStore where
       Just _ -> modifyTVar' size (subtract 1)
       _ -> pure ()
 
-  isolateQueue :: RecipientId -> STMQueue -> String -> STM a -> ExceptT ErrorType IO a
-  isolateQueue _ _ _ = liftIO . atomically
+  isolateQueue :: STMQueue -> String -> STM a -> ExceptT ErrorType IO a
+  isolateQueue _ _ = liftIO . atomically