1{-# LANGUAGE TupleSections, ConstraintKinds #-}
2
3-- | Extra functions for "Control.Concurrent".
4--
5--   This module includes three new types of 'MVar', namely 'Lock' (no associated value),
6--   'Var' (never empty) and 'Barrier' (filled at most once). See
7--   <https://neilmitchell.blogspot.co.uk/2012/06/flavours-of-mvar_04.html this blog post>
8--   for examples and justification.
9--
10--   If you need greater control of exceptions and threads
11--   see the <https://hackage.haskell.org/package/slave-thread slave-thread> package.
12--   If you need elaborate relationships between threads
13--   see the <https://hackage.haskell.org/package/async async> package.
14module Control.Concurrent.Extra(
15    module Control.Concurrent,
16    withNumCapabilities,
17    once, onceFork,
18    -- * Lock
19    Lock, newLock, withLock, withLockTry,
20    -- * Var
21    Var, newVar, readVar, writeVar, modifyVar, modifyVar_, withVar,
22    -- * Barrier
23    Barrier, newBarrier, signalBarrier, waitBarrier, waitBarrierMaybe,
24    ) where
25
26import Control.Concurrent
27import Control.Exception.Extra
28import Control.Monad.Extra
29import Data.Maybe
30import Data.Either.Extra
31import Data.Functor
32import Prelude
33
34
35-- | On GHC 7.6 and above with the @-threaded@ flag, brackets a call to 'setNumCapabilities'.
36--   On lower versions (which lack 'setNumCapabilities') this function just runs the argument action.
37withNumCapabilities :: Int -> IO a -> IO a
38withNumCapabilities new act | rtsSupportsBoundThreads = do
39    old <- getNumCapabilities
40    if old == new then act else
41        bracket_ (setNumCapabilities new) (setNumCapabilities old) act
42withNumCapabilities _ act = act
43
44
45-- | Given an action, produce a wrapped action that runs at most once.
46--   If the function raises an exception, the same exception will be reraised each time.
47--
48-- > let x ||| y = do t1 <- onceFork x; t2 <- onceFork y; t1; t2
49-- > \(x :: IO Int) -> void (once x) == pure ()
50-- > \(x :: IO Int) -> join (once x) == x
51-- > \(x :: IO Int) -> (do y <- once x; y; y) == x
52-- > \(x :: IO Int) -> (do y <- once x; y ||| y) == x
53once :: IO a -> IO (IO a)
54once act = do
55    var <- newVar OncePending
56    let run = either throwIO pure
57    pure $ mask $ \unmask -> join $ modifyVar var $ \v -> case v of
58        OnceDone x -> pure (v, unmask $ run x)
59        OnceRunning x -> pure (v, unmask $ run =<< waitBarrier x)
60        OncePending -> do
61            b <- newBarrier
62            pure $ (OnceRunning b,) $ do
63                res <- try_ $ unmask act
64                signalBarrier b res
65                modifyVar_ var $ \_ -> pure $ OnceDone res
66                run res
67
68data Once a = OncePending | OnceRunning (Barrier a) | OnceDone a
69
70
71-- | Like 'once', but immediately starts running the computation on a background thread.
72--
73-- > \(x :: IO Int) -> join (onceFork x) == x
74-- > \(x :: IO Int) -> (do a <- onceFork x; a; a) == x
75onceFork :: IO a -> IO (IO a)
76onceFork act = do
77    bar <- newBarrier
78    forkFinally act $ signalBarrier bar
79    pure $ eitherM throwIO pure $ waitBarrier bar
80
81
82---------------------------------------------------------------------
83-- LOCK
84
85-- | Like an 'MVar', but has no value.
86--   Used to guarantee single-threaded access, typically to some system resource.
87--   As an example:
88--
89-- @
90-- lock <- 'newLock'
91-- let output = 'withLock' lock . putStrLn
92-- forkIO $ do ...; output \"hello\"
93-- forkIO $ do ...; output \"world\"
94-- @
95--
96--   Here we are creating a lock to ensure that when writing output our messages
97--   do not get interleaved. This use of 'MVar' never blocks on a put. It is permissible,
98--   but rare, that a withLock contains a withLock inside it - but if so,
99--   watch out for deadlocks.
100
101newtype Lock = Lock (MVar ())
102
103-- | Create a new 'Lock'.
104newLock :: IO Lock
105newLock = Lock <$> newMVar ()
106
107-- | Perform some operation while holding 'Lock'. Will prevent all other
108--   operations from using the 'Lock' while the action is ongoing.
109withLock :: Lock -> IO a -> IO a
110withLock (Lock x) = withMVar x . const
111
112-- | Like 'withLock' but will never block. If the operation cannot be executed
113--   immediately it will return 'Nothing'.
114withLockTry :: Lock -> IO a -> IO (Maybe a)
115withLockTry (Lock m) act = bracket
116    (tryTakeMVar m)
117    (\v -> when (isJust v) $ putMVar m ())
118    (\v -> if isJust v then fmap Just act else pure Nothing)
119
120
121---------------------------------------------------------------------
122-- VAR
123
124-- | Like an 'MVar', but must always be full.
125--   Used to operate on a mutable variable in a thread-safe way.
126--   As an example:
127--
128-- @
129-- hits <- 'newVar' 0
130-- forkIO $ do ...; 'modifyVar_' hits (+1); ...
131-- i <- 'readVar' hits
132-- print (\"HITS\",i)
133-- @
134--
135--   Here we have a variable which we modify atomically, so modifications are
136--   not interleaved. This use of 'MVar' never blocks on a put. No modifyVar
137--   operation should ever block, and they should always complete in a reasonable
138--   timeframe. A 'Var' should not be used to protect some external resource, only
139--   the variable contained within. Information from a 'readVar' should not be subsequently
140--   inserted back into the 'Var'.
141newtype Var a = Var (MVar a)
142
143-- | Create a new 'Var' with a value.
144newVar :: a -> IO (Var a)
145newVar = fmap Var . newMVar
146
147-- | Read the current value of the 'Var'.
148readVar :: Var a -> IO a
149readVar (Var x) = readMVar x
150
151-- | Write a value to become the new value of 'Var'.
152writeVar :: Var a -> a -> IO ()
153writeVar v x = modifyVar_ v $ const $ pure x
154
155-- | Modify a 'Var' producing a new value and a return result.
156modifyVar :: Var a -> (a -> IO (a, b)) -> IO b
157modifyVar (Var x) f = modifyMVar x f
158
159-- | Modify a 'Var', a restricted version of 'modifyVar'.
160modifyVar_ :: Var a -> (a -> IO a) -> IO ()
161modifyVar_ (Var x) f = modifyMVar_ x f
162
163-- | Perform some operation using the value in the 'Var',
164--   a restricted version of 'modifyVar'.
165withVar :: Var a -> (a -> IO b) -> IO b
166withVar (Var x) f = withMVar x f
167
168
169---------------------------------------------------------------------
170-- BARRIER
171
172-- | Starts out empty, then is filled exactly once. As an example:
173--
174-- @
175-- bar <- 'newBarrier'
176-- forkIO $ do ...; val <- ...; 'signalBarrier' bar val
177-- print =<< 'waitBarrier' bar
178-- @
179--
180--   Here we create a barrier which will contain some computed value.
181--   A thread is forked to fill the barrier, while the main thread waits
182--   for it to complete. A barrier has similarities to a future or promise
183--   from other languages, has been known as an IVar in other Haskell work,
184--   and in some ways is like a manually managed thunk.
185newtype Barrier a = Barrier (MVar a)
186
187-- | Create a new 'Barrier'.
188newBarrier :: IO (Barrier a)
189newBarrier = Barrier <$> newEmptyMVar
190
191-- | Write a value into the Barrier, releasing anyone at 'waitBarrier'.
192--   Any subsequent attempts to signal the 'Barrier' will throw an exception.
193signalBarrier :: Partial => Barrier a -> a -> IO ()
194signalBarrier (Barrier var) v = do
195    b <- tryPutMVar var v
196    unless b $ errorIO "Control.Concurrent.Extra.signalBarrier, attempt to signal a barrier that has already been signaled"
197
198
199-- | Wait until a barrier has been signaled with 'signalBarrier'.
200waitBarrier :: Barrier a -> IO a
201waitBarrier (Barrier var) = readMVar var
202
203
204-- | A version of 'waitBarrier' that never blocks, returning 'Nothing'
205--   if the barrier has not yet been signaled.
206waitBarrierMaybe :: Barrier a -> IO (Maybe a)
207waitBarrierMaybe (Barrier bar) = tryReadMVar bar
208