{-# LANGUAGE TupleSections, ConstraintKinds #-}

-- | Extra functions for "Control.Concurrent".
--
--   This module includes three new types of 'MVar', namely 'Lock' (no associated value),
--   'Var' (never empty) and 'Barrier' (filled at most once). See
--   <https://neilmitchell.blogspot.co.uk/2012/06/flavours-of-mvar_04.html this blog post>
--   for examples and justification.
--
--   If you need greater control of exceptions and threads
--   see the <https://hackage.haskell.org/package/slave-thread slave-thread> package.
--   If you need elaborate relationships between threads
--   see the <https://hackage.haskell.org/package/async async> package.
module Control.Concurrent.Extra(
    module Control.Concurrent,
    withNumCapabilities,
    once, onceFork,
    -- * Lock
    Lock, newLock, withLock, withLockTry,
    -- * Var
    Var, newVar, readVar,
    writeVar, writeVar',
    modifyVar, modifyVar',
    modifyVar_, modifyVar_',
    withVar,
    -- * Barrier
    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)


-- | On GHC 7.6 and above with the @-threaded@ flag, brackets a call to 'setNumCapabilities'.
--   On lower versions (which lack 'setNumCapabilities') this function just runs the argument action.
withNumCapabilities :: Int -> IO a -> IO a
withNumCapabilities :: forall a. Int -> IO a -> IO a
withNumCapabilities Int
new IO a
act | Bool
rtsSupportsBoundThreads = do
    Int
old <- IO Int
getNumCapabilities
    if Int
old forall a. Eq a => a -> a -> Bool
== Int
new then IO a
act else
        forall a b c. IO a -> IO b -> IO c -> IO c
bracket_ (Int -> IO ()
setNumCapabilities Int
new) (Int -> IO ()
setNumCapabilities Int
old) IO a
act
withNumCapabilities Int
_ IO a
act = IO a
act


-- | Given an action, produce a wrapped action that runs at most once.
--   If the function raises an exception, the same exception will be reraised each time.
--
-- > let x ||| y = do t1 <- onceFork x; t2 <- onceFork y; t1; t2
-- > \(x :: IO Int) -> void (once x) == pure ()
-- > \(x :: IO Int) -> join (once x) == x
-- > \(x :: IO Int) -> (do y <- once x; y; y) == x
-- > \(x :: IO Int) -> (do y <- once x; y ||| y) == x
once :: IO a -> IO (IO a)
once :: forall a. IO a -> IO (IO a)
once IO a
act = do
    Var (Once (Either SomeException a))
var <- forall a. a -> IO (Var a)
newVar forall a. Once a
OncePending
    let run :: Either SomeException a -> IO a
run = forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either forall e a. Exception e => e -> IO a
throwIO forall (f :: * -> *) a. Applicative f => a -> f a
pure
    forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall b. ((forall a. IO a -> IO a) -> IO b) -> IO b
mask forall a b. (a -> b) -> a -> b
$ \forall a. IO a -> IO a
unmask -> forall (m :: * -> *) a. Monad m => m (m a) -> m a
join forall a b. (a -> b) -> a -> b
$ forall a b. Var a -> (a -> IO (a, b)) -> IO b
modifyVar Var (Once (Either SomeException a))
var forall a b. (a -> b) -> a -> b
$ \Once (Either SomeException a)
v -> case Once (Either SomeException a)
v of
        OnceDone Either SomeException a
x -> forall (f :: * -> *) a. Applicative f => a -> f a
pure (Once (Either SomeException a)
v, forall a. IO a -> IO a
unmask forall a b. (a -> b) -> a -> b
$ forall {a}. Either SomeException a -> IO a
run Either SomeException a
x)
        OnceRunning Barrier (Either SomeException a)
x -> forall (f :: * -> *) a. Applicative f => a -> f a
pure (Once (Either SomeException a)
v, forall a. IO a -> IO a
unmask forall a b. (a -> b) -> a -> b
$ forall {a}. Either SomeException a -> IO a
run forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall a. Barrier a -> IO a
waitBarrier Barrier (Either SomeException a)
x)
        Once (Either SomeException a)
OncePending -> do
            Barrier (Either SomeException a)
b <- forall a. IO (Barrier a)
newBarrier
            forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ (forall a. Barrier a -> Once a
OnceRunning Barrier (Either SomeException a)
b,) forall a b. (a -> b) -> a -> b
$ do
                Either SomeException a
res <- forall a. IO a -> IO (Either SomeException a)
try_ forall a b. (a -> b) -> a -> b
$ forall a. IO a -> IO a
unmask IO a
act
                forall a. Partial => Barrier a -> a -> IO ()
signalBarrier Barrier (Either SomeException a)
b Either SomeException a
res
                forall a. Var a -> (a -> IO a) -> IO ()
modifyVar_ Var (Once (Either SomeException a))
var forall a b. (a -> b) -> a -> b
$ \Once (Either SomeException a)
_ -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall a. a -> Once a
OnceDone Either SomeException a
res
                forall {a}. Either SomeException a -> IO a
run Either SomeException a
res

data Once a = OncePending | OnceRunning (Barrier a) | OnceDone a


-- | Like 'once', but immediately starts running the computation on a background thread.
--
-- > \(x :: IO Int) -> join (onceFork x) == x
-- > \(x :: IO Int) -> (do a <- onceFork x; a; a) == x
onceFork :: IO a -> IO (IO a)
onceFork :: forall a. IO a -> IO (IO a)
onceFork IO a
act = do
    Barrier (Either SomeException a)
bar <- forall a. IO (Barrier a)
newBarrier
    forall a. IO a -> (Either SomeException a -> IO ()) -> IO ThreadId
forkFinally IO a
act forall a b. (a -> b) -> a -> b
$ forall a. Partial => Barrier a -> a -> IO ()
signalBarrier Barrier (Either SomeException a)
bar
    forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a c b.
Monad m =>
(a -> m c) -> (b -> m c) -> m (Either a b) -> m c
eitherM forall e a. Exception e => e -> IO a
throwIO forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall a. Barrier a -> IO a
waitBarrier Barrier (Either SomeException a)
bar


---------------------------------------------------------------------
-- LOCK

-- | Like an 'MVar', but has no value.
--   Used to guarantee single-threaded access, typically to some system resource.
--   As an example:
--
-- @
-- lock <- 'newLock'
-- let output = 'withLock' lock . putStrLn
-- forkIO $ do ...; output \"hello\"
-- forkIO $ do ...; output \"world\"
-- @
--
--   Here we are creating a lock to ensure that when writing output our messages
--   do not get interleaved. This use of 'MVar' never blocks on a put. It is permissible,
--   but rare, that a withLock contains a withLock inside it - but if so,
--   watch out for deadlocks.

newtype Lock = Lock (MVar ())

-- | Create a new 'Lock'.
newLock :: IO Lock
newLock :: IO Lock
newLock = MVar () -> Lock
Lock forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a. a -> IO (MVar a)
newMVar ()

-- | Perform some operation while holding 'Lock'. Will prevent all other
--   operations from using the 'Lock' while the action is ongoing.
withLock :: Lock -> IO a -> IO a
withLock :: forall a. Lock -> IO a -> IO a
withLock (Lock MVar ()
x) = forall a b. MVar a -> (a -> IO b) -> IO b
withMVar MVar ()
x forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. a -> b -> a
const

-- | Like 'withLock' but will never block. If the operation cannot be executed
--   immediately it will return 'Nothing'.
withLockTry :: Lock -> IO a -> IO (Maybe a)
withLockTry :: forall a. Lock -> IO a -> IO (Maybe a)
withLockTry (Lock MVar ()
m) IO a
act = forall a b c. IO a -> (a -> IO b) -> (a -> IO c) -> IO c
bracket
    (forall a. MVar a -> IO (Maybe a)
tryTakeMVar MVar ()
m)
    (\Maybe ()
v -> forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (forall a. Maybe a -> Bool
isJust Maybe ()
v) forall a b. (a -> b) -> a -> b
$ forall a. MVar a -> a -> IO ()
putMVar MVar ()
m ())
    (\Maybe ()
v -> if forall a. Maybe a -> Bool
isJust Maybe ()
v then forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall a. a -> Maybe a
Just IO a
act else forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a. Maybe a
Nothing)


---------------------------------------------------------------------
-- VAR

-- | Like an 'MVar', but must always be full.
--   Used to operate on a mutable variable in a thread-safe way.
--   As an example:
--
-- @
-- hits <- 'newVar' 0
-- forkIO $ do ...; 'modifyVar_' hits (+1); ...
-- i <- 'readVar' hits
-- print (\"HITS\",i)
-- @
--
--   Here we have a variable which we modify atomically, so modifications are
--   not interleaved. This use of 'MVar' never blocks on a put. No modifyVar
--   operation should ever block, and they should always complete in a reasonable
--   timeframe. A 'Var' should not be used to protect some external resource, only
--   the variable contained within. Information from a 'readVar' should not be subsequently
--   inserted back into the 'Var'.
newtype Var a = Var (MVar a)

-- | Create a new 'Var' with a value.
newVar :: a -> IO (Var a)
newVar :: forall a. a -> IO (Var a)
newVar = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall a. MVar a -> Var a
Var forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. a -> IO (MVar a)
newMVar

-- | Read the current value of the 'Var'.
readVar :: Var a -> IO a
readVar :: forall a. Var a -> IO a
readVar (Var MVar a
x) = forall a. MVar a -> IO a
readMVar MVar a
x

-- | Write a value to become the new value of 'Var'.
writeVar :: Var a -> a -> IO ()
writeVar :: forall a. Var a -> a -> IO ()
writeVar Var a
v a
x = forall a. Var a -> (a -> IO a) -> IO ()
modifyVar_ Var a
v forall a b. (a -> b) -> a -> b
$ forall a b. a -> b -> a
const forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) a. Applicative f => a -> f a
pure a
x

-- | Strict variant of 'writeVar'
writeVar' :: Var a -> a -> IO ()
writeVar' :: forall a. Var a -> a -> IO ()
writeVar' Var a
v a
x = forall a. Var a -> (a -> IO a) -> IO ()
modifyVar_' Var a
v forall a b. (a -> b) -> a -> b
$ forall a b. a -> b -> a
const forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) a. Applicative f => a -> f a
pure a
x

-- | Modify a 'Var' producing a new value and a return result.
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 = forall a b. MVar a -> (a -> IO (a, b)) -> IO b
modifyMVar MVar a
x a -> IO (a, b)
f

-- | Strict variant of 'modifyVar''
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
    (a
newContents, b
res) <- forall a b. Var a -> (a -> IO (a, b)) -> IO b
modifyVar Var a
x forall a b. (a -> b) -> a -> b
$ \a
v -> do
        (a
newContents, b
res) <- a -> IO (a, b)
f a
v
        forall (f :: * -> *) a. Applicative f => a -> f a
pure (a
newContents, (a
newContents, b
res))
    forall a. a -> IO a
evaluate a
newContents
    forall (f :: * -> *) a. Applicative f => a -> f a
pure b
res

-- | Modify a 'Var', a restricted version of 'modifyVar'.
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 = forall a. MVar a -> (a -> IO a) -> IO ()
modifyMVar_ MVar a
x a -> IO a
f

-- | Strict variant of 'modifyVar_'
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
    a
newContents <- forall a b. Var a -> (a -> IO (a, b)) -> IO b
modifyVar Var a
x (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall a. a -> (a, a)
dupe forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> IO a
f)
    a
_ <- forall a. a -> IO a
evaluate a
newContents
    forall (f :: * -> *) a. Applicative f => a -> f a
pure ()

-- | Perform some operation using the value in the 'Var',
--   a restricted version of 'modifyVar'.
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 = forall a b. MVar a -> (a -> IO b) -> IO b
withMVar MVar a
x a -> IO b
f


---------------------------------------------------------------------
-- BARRIER

-- | Starts out empty, then is filled exactly once. As an example:
--
-- @
-- bar <- 'newBarrier'
-- forkIO $ do ...; val <- ...; 'signalBarrier' bar val
-- print =<< 'waitBarrier' bar
-- @
--
--   Here we create a barrier which will contain some computed value.
--   A thread is forked to fill the barrier, while the main thread waits
--   for it to complete. A barrier has similarities to a future or promise
--   from other languages, has been known as an IVar in other Haskell work,
--   and in some ways is like a manually managed thunk.
newtype Barrier a = Barrier (MVar a)

-- | Create a new 'Barrier'.
newBarrier :: IO (Barrier a)
newBarrier :: forall a. IO (Barrier a)
newBarrier = forall a. MVar a -> Barrier a
Barrier forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a. IO (MVar a)
newEmptyMVar

-- | Write a value into the Barrier, releasing anyone at 'waitBarrier'.
--   Any subsequent attempts to signal the 'Barrier' will throw an exception.
signalBarrier :: Partial => Barrier a -> a -> IO ()
signalBarrier :: forall a. Partial => Barrier a -> a -> IO ()
signalBarrier (Barrier MVar a
var) a
v = do
    Bool
b <- forall a. MVar a -> a -> IO Bool
tryPutMVar MVar a
var a
v
    forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless Bool
b forall a b. (a -> b) -> a -> b
$ forall a. Partial => String -> IO a
errorIO String
"Control.Concurrent.Extra.signalBarrier, attempt to signal a barrier that has already been signaled"


-- | Wait until a barrier has been signaled with 'signalBarrier'.
waitBarrier :: Barrier a -> IO a
waitBarrier :: forall a. Barrier a -> IO a
waitBarrier (Barrier MVar a
var) = forall a. MVar a -> IO a
readMVar MVar a
var


-- | A version of 'waitBarrier' that never blocks, returning 'Nothing'
--   if the barrier has not yet been signaled.
waitBarrierMaybe :: Barrier a -> IO (Maybe a)
waitBarrierMaybe :: forall a. Barrier a -> IO (Maybe a)
waitBarrierMaybe (Barrier MVar a
bar) = forall a. MVar a -> IO (Maybe a)
tryReadMVar MVar a
bar