module Simplex.Messaging.Agent.Lock
  ( Lock,
    createLock,
    createLockIO,
    withLock,
    withLock',
    withGetLock,
    withGetLocks,
    getPutLock,
  )
where

import Control.Monad (void)
import Control.Monad.Except (ExceptT (..), runExceptT)
import Control.Monad.IO.Unlift
import Data.Functor (($>))
import Data.Set (Set)
import qualified Data.Set as S
import Data.Text (Text)
import UnliftIO.Async (forConcurrently)
import qualified UnliftIO.Exception as E
import UnliftIO.STM

type Lock = TMVar Text

createLock :: STM Lock
createLock :: STM Lock
createLock = STM Lock
forall a. STM (TMVar a)
newEmptyTMVar
{-# INLINE createLock #-}

createLockIO :: IO Lock
createLockIO :: IO Lock
createLockIO = IO Lock
forall (m :: * -> *) a. MonadIO m => m (TMVar a)
newEmptyTMVarIO
{-# INLINE createLockIO #-}

withLock :: MonadUnliftIO m => Lock -> Text -> ExceptT e m a -> ExceptT e m a
withLock :: forall (m :: * -> *) e a.
MonadUnliftIO m =>
Lock -> Text -> ExceptT e m a -> ExceptT e m a
withLock Lock
lock Text
name = m (Either e a) -> ExceptT e m a
forall e (m :: * -> *) a. m (Either e a) -> ExceptT e m a
ExceptT (m (Either e a) -> ExceptT e m a)
-> (ExceptT e m a -> m (Either e a))
-> ExceptT e m a
-> ExceptT e m a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Lock -> Text -> m (Either e a) -> m (Either e a)
forall (m :: * -> *) a.
MonadUnliftIO m =>
Lock -> Text -> m a -> m a
withLock' Lock
lock Text
name (m (Either e a) -> m (Either e a))
-> (ExceptT e m a -> m (Either e a))
-> ExceptT e m a
-> m (Either e a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ExceptT e m a -> m (Either e a)
forall e (m :: * -> *) a. ExceptT e m a -> m (Either e a)
runExceptT
{-# INLINE withLock #-}

withLock' :: MonadUnliftIO m => Lock -> Text -> m a -> m a
withLock' :: forall (m :: * -> *) a.
MonadUnliftIO m =>
Lock -> Text -> m a -> m a
withLock' Lock
lock Text
name =
  m () -> m () -> m a -> m a
forall (m :: * -> *) a b c.
MonadUnliftIO m =>
m a -> m b -> m c -> m c
E.bracket_
    (STM () -> m ()
forall (m :: * -> *) a. MonadIO m => STM a -> m a
atomically (STM () -> m ()) -> STM () -> m ()
forall a b. (a -> b) -> a -> b
$ Lock -> Text -> STM ()
forall a. TMVar a -> a -> STM ()
putTMVar Lock
lock Text
name)
    (m Text -> m ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (m Text -> m ()) -> (STM Text -> m Text) -> STM Text -> m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. STM Text -> m Text
forall (m :: * -> *) a. MonadIO m => STM a -> m a
atomically (STM Text -> m ()) -> STM Text -> m ()
forall a b. (a -> b) -> a -> b
$ Lock -> STM Text
forall a. TMVar a -> STM a
takeTMVar Lock
lock)

withGetLock :: MonadUnliftIO m => (k -> STM Lock) -> k -> Text -> m a -> m a
withGetLock :: forall (m :: * -> *) k a.
MonadUnliftIO m =>
(k -> STM Lock) -> k -> Text -> m a -> m a
withGetLock k -> STM Lock
getLock k
key Text
name m a
a =
  m Lock -> (Lock -> m Text) -> (Lock -> m a) -> m a
forall (m :: * -> *) a b c.
MonadUnliftIO m =>
m a -> (a -> m b) -> (a -> m c) -> m c
E.bracket
    (STM Lock -> m Lock
forall (m :: * -> *) a. MonadIO m => STM a -> m a
atomically (STM Lock -> m Lock) -> STM Lock -> m Lock
forall a b. (a -> b) -> a -> b
$ (k -> STM Lock) -> k -> Text -> STM Lock
forall k. (k -> STM Lock) -> k -> Text -> STM Lock
getPutLock k -> STM Lock
getLock k
key Text
name)
    (STM Text -> m Text
forall (m :: * -> *) a. MonadIO m => STM a -> m a
atomically (STM Text -> m Text) -> (Lock -> STM Text) -> Lock -> m Text
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Lock -> STM Text
forall a. TMVar a -> STM a
takeTMVar)
    (m a -> Lock -> m a
forall a b. a -> b -> a
const m a
a)

withGetLocks :: MonadUnliftIO m => (k -> STM Lock) -> Set k -> Text -> m a -> m a
withGetLocks :: forall (m :: * -> *) k a.
MonadUnliftIO m =>
(k -> STM Lock) -> Set k -> Text -> m a -> m a
withGetLocks k -> STM Lock
getLock Set k
keys Text
name = m [Lock] -> ([Lock] -> m ()) -> ([Lock] -> m a) -> m a
forall (m :: * -> *) a b c.
MonadUnliftIO m =>
m a -> (a -> m b) -> (a -> m c) -> m c
E.bracket m [Lock]
holdLocks [Lock] -> m ()
forall {b}. [TMVar b] -> m ()
releaseLocks (([Lock] -> m a) -> m a) -> (m a -> [Lock] -> m a) -> m a -> m a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. m a -> [Lock] -> m a
forall a b. a -> b -> a
const
  where
    holdLocks :: m [Lock]
holdLocks = [k] -> (k -> m Lock) -> m [Lock]
forall (m :: * -> *) (t :: * -> *) a b.
(MonadUnliftIO m, Traversable t) =>
t a -> (a -> m b) -> m (t b)
forConcurrently (Set k -> [k]
forall a. Set a -> [a]
S.toList Set k
keys) ((k -> m Lock) -> m [Lock]) -> (k -> m Lock) -> m [Lock]
forall a b. (a -> b) -> a -> b
$ \k
key -> STM Lock -> m Lock
forall (m :: * -> *) a. MonadIO m => STM a -> m a
atomically (STM Lock -> m Lock) -> STM Lock -> m Lock
forall a b. (a -> b) -> a -> b
$ (k -> STM Lock) -> k -> Text -> STM Lock
forall k. (k -> STM Lock) -> k -> Text -> STM Lock
getPutLock k -> STM Lock
getLock k
key Text
name
    releaseLocks :: [TMVar b] -> m ()
releaseLocks = (TMVar b -> m b) -> [TMVar b] -> m ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (STM b -> m b
forall (m :: * -> *) a. MonadIO m => STM a -> m a
atomically (STM b -> m b) -> (TMVar b -> STM b) -> TMVar b -> m b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TMVar b -> STM b
forall a. TMVar a -> STM a
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 :: forall k. (k -> STM Lock) -> k -> Text -> STM Lock
getPutLock k -> STM Lock
getLock k
key Text
name = k -> STM Lock
getLock k
key STM Lock -> (Lock -> STM Lock) -> STM Lock
forall a b. STM a -> (a -> STM b) -> STM b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \Lock
l -> Lock -> Text -> STM ()
forall a. TMVar a -> a -> STM ()
putTMVar Lock
l Text
name STM () -> Lock -> STM Lock
forall (f :: * -> *) a b. Functor f => f a -> b -> f b
$> Lock
l