{-# LANGUAGE CPP #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE ScopedTypeVariables #-}
module Control.Concurrent.FVar (modifyFVar, newFVar, readFVar, FVar) where
import Control.Concurrent.MVar
import Control.Monad (replicateM_, when)
import Control.Monad.Catch (bracket, bracket_, MonadMask)
import Control.Monad.IO.Class (MonadIO, liftIO)
import Data.IORef
import Data.Maybe (isJust)
import System.IO.Unsafe (unsafePerformIO)
newtype FVar a = FVar { iRef :: IORef (RWVar a) } deriving (Eq)
data RWVar a = RWVar { readLock :: Maybe (MVar ()),
readMutex :: MVar (),
writeMutex :: MVar (),
readersWaiting :: Int,
writersWaiting :: Int,
readers :: Int,
writers :: Int,
unboxed :: a
}
#ifdef DEBUG
{-# NOINLINE out #-}
out :: MVar ()
out = unsafePerformIO $ newMVar ()
debug :: (MonadIO m) => String -> m ()
debug = liftIO . withMVar out . const . putStrLn
#else
debug :: (MonadIO m) => String -> m ()
debug = const $ return ()
#endif
newFVar :: (MonadIO m) => a -> m (FVar a)
newFVar val = liftIO $ do
rMVar <- newEmptyMVar
wMVar <- newEmptyMVar
ref <- newIORef RWVar { readLock = Nothing,
readersWaiting = 0,
writersWaiting = 0,
readers = 0,
writers = 0,
readMutex = rMVar,
writeMutex = wMVar,
unboxed = val }
return $ FVar ref
readFVar :: (MonadMask m, MonadIO m) => FVar a -> (a -> m b) -> m b
readFVar FVar {..} = bracket wantsRead readDone
where wantsRead = do
(x, mLock) <- liftIO $ atomicModifyIORef' iRef readerIn
maybe (access x) acquire mLock
access x = do
debug "reader reads fast"
return x
readDone _ = do
debug "reader done"
mLock <- liftIO $ atomicModifyIORef' iRef $ readerOut
mapM_ release mLock
acquire var = do
debug "reader waits"
liftIO $ do
takeMVar var
debug "reader acquired lock"
unboxed <$> readIORef iRef
release var = do
debug "reader released write lock"
liftIO $ putMVar var ()
modifyFVar :: (MonadMask m, MonadIO m) => FVar a -> (a -> m (a, b)) -> m b
modifyFVar FVar {..} mutator = bracket_ wantsWrite writeDone mutate
where wantsWrite = liftIO $ mapM_ acquire =<< atomicModifyIORef' iRef writerIn
writeDone = liftIO $ do
debug "writer done"
mLock <- atomicModifyIORef' iRef writerOut
mapM_ release mLock
acquire var = do
debug "writer waits"
takeMVar var
debug "writer acquired lock"
release (waiters, var) = do
debug "writer released lock"
liftIO $ replicateM_ waiters $ putMVar var ()
mutate = do
debug "writer writes"
(value, result) <- mutator =<< unboxed <$> liftIO (readIORef iRef)
liftIO . atomicModifyIORef' iRef $ \var -> (var { unboxed = value }, result)
readerIn :: RWVar a -> (RWVar a, (a, Maybe (MVar ())))
readerIn var@RWVar {..} =
(var { readers = readers + 1, readersWaiting = waitingNow }, (unboxed, readLock))
where waitingNow | locked = readersWaiting + 1
| otherwise = readersWaiting
locked = isJust readLock
readerOut :: RWVar a -> (RWVar a, Maybe (MVar ()))
readerOut var@RWVar {..} =
(var { readers = readersNow, writersWaiting = waitingNow}, withLock)
where withLock
| unlockWriter = return writeMutex
| otherwise = Nothing
readersNow = readers - 1
waitingNow
| unlockWriter = writersWaiting - 1
| otherwise = writersWaiting
unlockWriter = readersNow == readersWaiting && writersWaiting /= 0
writerIn :: RWVar a -> (RWVar a, Maybe (MVar ()))
writerIn var@RWVar {..} =
(var { readLock = return readMutex,
writers = writers + 1,
writersWaiting = waitingNow }, withLock)
where waitingNow
| stalled = writersWaiting + 1
| otherwise = writersWaiting
withLock
| stalled = return writeMutex
| otherwise = Nothing
stalled = writers /= 0 || readers /= readersWaiting
writerOut :: RWVar a -> (RWVar a, Maybe (Int, MVar ()))
writerOut var@RWVar {..} =
(var { readLock = rearmLock,
readersWaiting = 0,
writers = writers - 1,
writersWaiting = nowWaiting }, withLock)
where rearmLock
| stalledWriters = return readMutex
| otherwise = Nothing
withLock
| stalledReaders = return (readersWaiting, readMutex)
| stalledWriters = return (1, writeMutex)
| otherwise = Nothing
nowWaiting
| stalledWriters && not stalledReaders = writersWaiting - 1
| otherwise = writersWaiting
stalledReaders = readersWaiting /= 0
stalledWriters = writersWaiting /= 0