{-# LANGUAGE TupleSections, ConstraintKinds #-}
module Control.Concurrent.Extra(
module Control.Concurrent,
withNumCapabilities,
once, onceFork,
Lock, newLock, withLock, withLockTry,
Var, newVar, readVar,
writeVar, writeVar',
modifyVar, modifyVar',
modifyVar_, modifyVar_',
withVar,
Barrier, newBarrier, signalBarrier, waitBarrier, waitBarrierMaybe,
) where
import Control.Concurrent
import Control.Exception.Extra
import Control.Monad.Extra
import Data.Maybe
import Data.Either.Extra
import Data.Functor
import Prelude
import Data.Tuple.Extra (dupe)
withNumCapabilities :: Int -> IO a -> IO a
withNumCapabilities :: forall a. Int -> IO a -> IO a
withNumCapabilities Int
new IO a
act | Bool
rtsSupportsBoundThreads = do
old <- IO Int
getNumCapabilities
if old == new then act else
bracket_ (setNumCapabilities new) (setNumCapabilities old) act
withNumCapabilities Int
_ IO a
act = IO a
act
once :: IO a -> IO (IO a)
once :: forall a. IO a -> IO (IO a)
once IO a
act = do
var <- Once (Either SomeException a)
-> IO (Var (Once (Either SomeException a)))
forall a. a -> IO (Var a)
newVar Once (Either SomeException a)
forall a. Once a
OncePending
let run = (SomeException -> IO a)
-> (a -> IO a) -> Either SomeException a -> IO a
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either SomeException -> IO a
forall e a. (Partial, Exception e) => e -> IO a
throwIO a -> IO a
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure
pure $ mask $ \forall a. IO a -> IO a
unmask -> IO (IO a) -> IO a
forall (m :: * -> *) a. Monad m => m (m a) -> m a
join (IO (IO a) -> IO a) -> IO (IO a) -> IO a
forall a b. (a -> b) -> a -> b
$ Var (Once (Either SomeException a))
-> (Once (Either SomeException a)
-> IO (Once (Either SomeException a), IO a))
-> IO (IO a)
forall a b. Var a -> (a -> IO (a, b)) -> IO b
modifyVar Var (Once (Either SomeException a))
var ((Once (Either SomeException a)
-> IO (Once (Either SomeException a), IO a))
-> IO (IO a))
-> (Once (Either SomeException a)
-> IO (Once (Either SomeException a), IO a))
-> IO (IO a)
forall a b. (a -> b) -> a -> b
$ \Once (Either SomeException a)
v -> case Once (Either SomeException a)
v of
OnceDone Either SomeException a
x -> (Once (Either SomeException a), IO a)
-> IO (Once (Either SomeException a), IO a)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Once (Either SomeException a)
v, IO a -> IO a
forall a. IO a -> IO a
unmask (IO a -> IO a) -> IO a -> IO a
forall a b. (a -> b) -> a -> b
$ Either SomeException a -> IO a
forall {a}. Either SomeException a -> IO a
run Either SomeException a
x)
OnceRunning Barrier (Either SomeException a)
x -> (Once (Either SomeException a), IO a)
-> IO (Once (Either SomeException a), IO a)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Once (Either SomeException a)
v, IO a -> IO a
forall a. IO a -> IO a
unmask (IO a -> IO a) -> IO a -> IO a
forall a b. (a -> b) -> a -> b
$ Either SomeException a -> IO a
forall {a}. Either SomeException a -> IO a
run (Either SomeException a -> IO a)
-> IO (Either SomeException a) -> IO a
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Barrier (Either SomeException a) -> IO (Either SomeException a)
forall a. Barrier a -> IO a
waitBarrier Barrier (Either SomeException a)
x)
Once (Either SomeException a)
OncePending -> do
b <- IO (Barrier (Either SomeException a))
forall a. IO (Barrier a)
newBarrier
pure $ (OnceRunning b,) $ do
res <- try_ $ unmask act
signalBarrier b res
modifyVar_ var $ \Once (Either SomeException a)
_ -> Once (Either SomeException a) -> IO (Once (Either SomeException a))
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Once (Either SomeException a)
-> IO (Once (Either SomeException a)))
-> Once (Either SomeException a)
-> IO (Once (Either SomeException a))
forall a b. (a -> b) -> a -> b
$ Either SomeException a -> Once (Either SomeException a)
forall a. a -> Once a
OnceDone Either SomeException a
res
run res
data Once a = OncePending | OnceRunning (Barrier a) | OnceDone a
onceFork :: IO a -> IO (IO a)
onceFork :: forall a. IO a -> IO (IO a)
onceFork IO a
act = do
bar <- IO (Barrier (Either SomeException a))
forall a. IO (Barrier a)
newBarrier
forkFinally act $ signalBarrier bar
pure $ eitherM throwIO pure $ waitBarrier bar
newtype Lock = Lock (MVar ())
newLock :: IO Lock
newLock :: IO Lock
newLock = MVar () -> Lock
Lock (MVar () -> Lock) -> IO (MVar ()) -> IO Lock
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> () -> IO (MVar ())
forall a. a -> IO (MVar a)
newMVar ()
withLock :: Lock -> IO a -> IO a
withLock :: forall a. Lock -> IO a -> IO a
withLock (Lock MVar ()
x) = MVar () -> (() -> IO a) -> IO a
forall a b. MVar a -> (a -> IO b) -> IO b
withMVar MVar ()
x ((() -> IO a) -> IO a) -> (IO a -> () -> IO a) -> IO a -> IO a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IO a -> () -> IO a
forall a b. a -> b -> a
const
withLockTry :: Lock -> IO a -> IO (Maybe a)
withLockTry :: forall a. Lock -> IO a -> IO (Maybe a)
withLockTry (Lock MVar ()
m) IO a
act = IO (Maybe ())
-> (Maybe () -> IO ())
-> (Maybe () -> IO (Maybe a))
-> IO (Maybe a)
forall a b c. IO a -> (a -> IO b) -> (a -> IO c) -> IO c
bracket
(MVar () -> IO (Maybe ())
forall a. MVar a -> IO (Maybe a)
tryTakeMVar MVar ()
m)
(\Maybe ()
v -> Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Maybe () -> Bool
forall a. Maybe a -> Bool
isJust Maybe ()
v) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ MVar () -> () -> IO ()
forall a. MVar a -> a -> IO ()
putMVar MVar ()
m ())
(\Maybe ()
v -> if Maybe () -> Bool
forall a. Maybe a -> Bool
isJust Maybe ()
v then (a -> Maybe a) -> IO a -> IO (Maybe a)
forall a b. (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> Maybe a
forall a. a -> Maybe a
Just IO a
act else Maybe a -> IO (Maybe a)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe a
forall a. Maybe a
Nothing)
newtype Var a = Var (MVar a)
newVar :: a -> IO (Var a)
newVar :: forall a. a -> IO (Var a)
newVar = (MVar a -> Var a) -> IO (MVar a) -> IO (Var a)
forall a b. (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap MVar a -> Var a
forall a. MVar a -> Var a
Var (IO (MVar a) -> IO (Var a))
-> (a -> IO (MVar a)) -> a -> IO (Var a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> IO (MVar a)
forall a. a -> IO (MVar a)
newMVar
readVar :: Var a -> IO a
readVar :: forall a. Var a -> IO a
readVar (Var MVar a
x) = MVar a -> IO a
forall a. MVar a -> IO a
readMVar MVar a
x
writeVar :: Var a -> a -> IO ()
writeVar :: forall a. Var a -> a -> IO ()
writeVar Var a
v a
x = Var a -> (a -> IO a) -> IO ()
forall a. Var a -> (a -> IO a) -> IO ()
modifyVar_ Var a
v ((a -> IO a) -> IO ()) -> (a -> IO a) -> IO ()
forall a b. (a -> b) -> a -> b
$ IO a -> a -> IO a
forall a b. a -> b -> a
const (IO a -> a -> IO a) -> IO a -> a -> IO a
forall a b. (a -> b) -> a -> b
$ a -> IO a
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
x
writeVar' :: Var a -> a -> IO ()
writeVar' :: forall a. Var a -> a -> IO ()
writeVar' Var a
v a
x = Var a -> (a -> IO a) -> IO ()
forall a. Var a -> (a -> IO a) -> IO ()
modifyVar_' Var a
v ((a -> IO a) -> IO ()) -> (a -> IO a) -> IO ()
forall a b. (a -> b) -> a -> b
$ IO a -> a -> IO a
forall a b. a -> b -> a
const (IO a -> a -> IO a) -> IO a -> a -> IO a
forall a b. (a -> b) -> a -> b
$ a -> IO a
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
x
modifyVar :: Var a -> (a -> IO (a, b)) -> IO b
modifyVar :: forall a b. Var a -> (a -> IO (a, b)) -> IO b
modifyVar (Var MVar a
x) a -> IO (a, b)
f = MVar a -> (a -> IO (a, b)) -> IO b
forall a b. MVar a -> (a -> IO (a, b)) -> IO b
modifyMVar MVar a
x a -> IO (a, b)
f
modifyVar' :: Var a -> (a -> IO (a, b)) -> IO b
modifyVar' :: forall a b. Var a -> (a -> IO (a, b)) -> IO b
modifyVar' Var a
x a -> IO (a, b)
f = do
(newContents, res) <- Var a -> (a -> IO (a, (a, b))) -> IO (a, b)
forall a b. Var a -> (a -> IO (a, b)) -> IO b
modifyVar Var a
x ((a -> IO (a, (a, b))) -> IO (a, b))
-> (a -> IO (a, (a, b))) -> IO (a, b)
forall a b. (a -> b) -> a -> b
$ \a
v -> do
(newContents, res) <- a -> IO (a, b)
f a
v
pure (newContents, (newContents, res))
evaluate newContents
pure res
modifyVar_ :: Var a -> (a -> IO a) -> IO ()
modifyVar_ :: forall a. Var a -> (a -> IO a) -> IO ()
modifyVar_ (Var MVar a
x) a -> IO a
f = MVar a -> (a -> IO a) -> IO ()
forall a. MVar a -> (a -> IO a) -> IO ()
modifyMVar_ MVar a
x a -> IO a
f
modifyVar_' :: Var a -> (a -> IO a) -> IO ()
modifyVar_' :: forall a. Var a -> (a -> IO a) -> IO ()
modifyVar_' Var a
x a -> IO a
f = do
newContents <- Var a -> (a -> IO (a, a)) -> IO a
forall a b. Var a -> (a -> IO (a, b)) -> IO b
modifyVar Var a
x ((a -> (a, a)) -> IO a -> IO (a, a)
forall a b. (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> (a, a)
forall a. a -> (a, a)
dupe (IO a -> IO (a, a)) -> (a -> IO a) -> a -> IO (a, a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> IO a
f)
_ <- evaluate newContents
pure ()
withVar :: Var a -> (a -> IO b) -> IO b
withVar :: forall a b. Var a -> (a -> IO b) -> IO b
withVar (Var MVar a
x) a -> IO b
f = MVar a -> (a -> IO b) -> IO b
forall a b. MVar a -> (a -> IO b) -> IO b
withMVar MVar a
x a -> IO b
f
newtype Barrier a = Barrier (MVar a)
newBarrier :: IO (Barrier a)
newBarrier :: forall a. IO (Barrier a)
newBarrier = MVar a -> Barrier a
forall a. MVar a -> Barrier a
Barrier (MVar a -> Barrier a) -> IO (MVar a) -> IO (Barrier a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IO (MVar a)
forall a. IO (MVar a)
newEmptyMVar
signalBarrier :: Partial => Barrier a -> a -> IO ()
signalBarrier :: forall a. Partial => Barrier a -> a -> IO ()
signalBarrier (Barrier MVar a
var) a
v = do
b <- MVar a -> a -> IO Bool
forall a. MVar a -> a -> IO Bool
tryPutMVar MVar a
var a
v
unless b $ errorIO "Control.Concurrent.Extra.signalBarrier, attempt to signal a barrier that has already been signaled"
waitBarrier :: Barrier a -> IO a
waitBarrier :: forall a. Barrier a -> IO a
waitBarrier (Barrier MVar a
var) = MVar a -> IO a
forall a. MVar a -> IO a
readMVar MVar a
var
waitBarrierMaybe :: Barrier a -> IO (Maybe a)
waitBarrierMaybe :: forall a. Barrier a -> IO (Maybe a)
waitBarrierMaybe (Barrier MVar a
bar) = MVar a -> IO (Maybe a)
forall a. MVar a -> IO (Maybe a)
tryReadMVar MVar a
bar