mirror of
https://github.com/simplex-chat/simplexmq.git
synced 2026-04-26 13:07:25 +00:00
120 lines
3.4 KiB
Haskell
120 lines
3.4 KiB
Haskell
{-# LANGUAGE BangPatterns #-}
|
|
{-# LANGUAGE ConstraintKinds #-}
|
|
{-# LANGUAGE FlexibleContexts #-}
|
|
{-# LANGUAGE FlexibleInstances #-}
|
|
{-# LANGUAGE LambdaCase #-}
|
|
{-# LANGUAGE MultiParamTypeClasses #-}
|
|
{-# LANGUAGE NamedFieldPuns #-}
|
|
{-# LANGUAGE TupleSections #-}
|
|
|
|
module Simplex.Messaging.Server.MsgStore.STM
|
|
( STMMsgStore,
|
|
MsgQueue (..),
|
|
newMsgStore,
|
|
getMsgQueue,
|
|
delMsgQueue,
|
|
delMsgQueueSize,
|
|
writeMsg,
|
|
tryPeekMsg,
|
|
peekMsg,
|
|
tryDelMsg,
|
|
tryDelPeekMsg,
|
|
deleteExpiredMsgs,
|
|
)
|
|
where
|
|
|
|
import qualified Data.ByteString.Char8 as B
|
|
import Data.Functor (($>))
|
|
import Data.Int (Int64)
|
|
import Data.Time.Clock.System (SystemTime (systemSeconds))
|
|
import Simplex.Messaging.Protocol (Message (..), MsgId, RecipientId)
|
|
import Simplex.Messaging.TMap (TMap)
|
|
import qualified Simplex.Messaging.TMap as TM
|
|
import UnliftIO.STM
|
|
|
|
data MsgQueue = MsgQueue
|
|
{ msgQueue :: TQueue Message,
|
|
quota :: Int,
|
|
canWrite :: TVar Bool,
|
|
size :: TVar Int
|
|
}
|
|
|
|
type STMMsgStore = TMap RecipientId MsgQueue
|
|
|
|
newMsgStore :: IO STMMsgStore
|
|
newMsgStore = TM.emptyIO
|
|
|
|
getMsgQueue :: STMMsgStore -> RecipientId -> Int -> STM MsgQueue
|
|
getMsgQueue st rId quota = maybe newQ pure =<< TM.lookup rId st
|
|
where
|
|
newQ = do
|
|
msgQueue <- newTQueue
|
|
canWrite <- newTVar True
|
|
size <- newTVar 0
|
|
let q = MsgQueue {msgQueue, quota, canWrite, size}
|
|
TM.insert rId q st
|
|
pure q
|
|
|
|
delMsgQueue :: STMMsgStore -> RecipientId -> STM ()
|
|
delMsgQueue st rId = TM.delete rId st
|
|
|
|
delMsgQueueSize :: STMMsgStore -> RecipientId -> STM Int
|
|
delMsgQueueSize st rId = TM.lookupDelete rId st >>= maybe (pure 0) (\MsgQueue {size} -> readTVar size)
|
|
|
|
writeMsg :: MsgQueue -> Message -> STM (Maybe (Message, Bool))
|
|
writeMsg MsgQueue {msgQueue = q, quota, canWrite, size} !msg = do
|
|
canWrt <- readTVar canWrite
|
|
empty <- isEmptyTQueue q
|
|
if canWrt || empty
|
|
then do
|
|
canWrt' <- (quota >) <$> readTVar size
|
|
writeTVar canWrite $! canWrt'
|
|
modifyTVar' size (+ 1)
|
|
if canWrt'
|
|
then writeTQueue q msg $> Just (msg, empty)
|
|
else (writeTQueue q $! msgQuota) $> Nothing
|
|
else pure Nothing
|
|
where
|
|
msgQuota = MessageQuota {msgId = msgId msg, msgTs = msgTs msg}
|
|
|
|
tryPeekMsg :: MsgQueue -> STM (Maybe Message)
|
|
tryPeekMsg = tryPeekTQueue . msgQueue
|
|
{-# INLINE tryPeekMsg #-}
|
|
|
|
peekMsg :: MsgQueue -> STM Message
|
|
peekMsg = peekTQueue . msgQueue
|
|
{-# INLINE peekMsg #-}
|
|
|
|
tryDelMsg :: MsgQueue -> MsgId -> STM (Maybe Message)
|
|
tryDelMsg mq msgId' =
|
|
tryPeekMsg mq >>= \case
|
|
msg_@(Just msg)
|
|
| msgId msg == msgId' || B.null msgId' -> tryDeleteMsg mq >> pure msg_
|
|
| otherwise -> pure Nothing
|
|
_ -> pure Nothing
|
|
|
|
-- atomic delete (== read) last and peek next message if available
|
|
tryDelPeekMsg :: MsgQueue -> MsgId -> STM (Maybe Message, Maybe Message)
|
|
tryDelPeekMsg mq msgId' =
|
|
tryPeekMsg mq >>= \case
|
|
msg_@(Just msg)
|
|
| msgId msg == msgId' || B.null msgId' -> (msg_,) <$> (tryDeleteMsg mq >> tryPeekMsg mq)
|
|
| otherwise -> pure (Nothing, msg_)
|
|
_ -> pure (Nothing, Nothing)
|
|
|
|
deleteExpiredMsgs :: MsgQueue -> Int64 -> STM Int
|
|
deleteExpiredMsgs mq old = loop 0
|
|
where
|
|
loop dc =
|
|
tryPeekMsg mq >>= \case
|
|
Just Message {msgTs}
|
|
| systemSeconds msgTs < old ->
|
|
tryDeleteMsg mq >> loop (dc + 1)
|
|
_ -> pure dc
|
|
|
|
tryDeleteMsg :: MsgQueue -> STM ()
|
|
tryDeleteMsg MsgQueue {msgQueue = q, size} =
|
|
tryReadTQueue q >>= \case
|
|
Just _ -> modifyTVar' size (subtract 1)
|
|
_ -> pure ()
|