{-# LANGUAGE DeriveDataTypeable #-}
module Test.Tasty.Parallel (ActionStatus(..), Action(..), runInParallel) where
import Control.Monad
import Control.Concurrent
import Control.Concurrent.Async
import Control.Concurrent.STM
import Foreign.StablePtr
data ActionStatus
= ActionReady
| ActionSkip
| ActionWait
deriving ActionStatus -> ActionStatus -> Bool
(ActionStatus -> ActionStatus -> Bool)
-> (ActionStatus -> ActionStatus -> Bool) -> Eq ActionStatus
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: ActionStatus -> ActionStatus -> Bool
== :: ActionStatus -> ActionStatus -> Bool
$c/= :: ActionStatus -> ActionStatus -> Bool
/= :: ActionStatus -> ActionStatus -> Bool
Eq
data Action = Action
{ Action -> STM ActionStatus
actionStatus :: STM ActionStatus
, Action -> IO ()
actionRun :: IO ()
, Action -> STM ()
actionSkip :: STM ()
}
runInParallel
:: Int
-> [Action]
-> IO (IO ())
runInParallel :: Int -> [Action] -> IO (IO ())
runInParallel Int
nthreads [Action]
actions = do
callingThread <- IO ThreadId
myThreadId
_ <- newStablePtr callingThread
actionsVar <- atomically $ newTMVar actions
pids <- replicateM nthreads (async $ work actionsVar)
return $ do
_ <- atomically $ swapTMVar actionsVar []
mapM_ cancel pids
work :: TMVar [Action] -> IO ()
work :: TMVar [Action] -> IO ()
work TMVar [Action]
actionsVar = IO ()
go
where
go :: IO ()
go = do
IO (IO ()) -> IO ()
forall (m :: * -> *) a. Monad m => m (m a) -> m a
join (IO (IO ()) -> IO ())
-> (STM (IO ()) -> IO (IO ())) -> STM (IO ()) -> IO ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. STM (IO ()) -> IO (IO ())
forall a. STM a -> IO a
atomically (STM (IO ()) -> IO ()) -> STM (IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ do
mb_ready <- [Action] -> STM (Maybe (Action, [Action]))
findBool ([Action] -> STM (Maybe (Action, [Action])))
-> STM [Action] -> STM (Maybe (Action, [Action]))
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< TMVar [Action] -> STM [Action]
forall a. TMVar a -> STM a
takeTMVar TMVar [Action]
actionsVar
case mb_ready of
Maybe (Action, [Action])
Nothing -> do
TMVar [Action] -> [Action] -> STM ()
forall a. TMVar a -> a -> STM ()
putTMVar TMVar [Action]
actionsVar []
IO () -> STM (IO ())
forall a. a -> STM a
forall (m :: * -> *) a. Monad m => a -> m a
return (IO () -> STM (IO ())) -> IO () -> STM (IO ())
forall a b. (a -> b) -> a -> b
$ () -> IO ()
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ()
Just (Action
this, [Action]
rest) -> do
TMVar [Action] -> [Action] -> STM ()
forall a. TMVar a -> a -> STM ()
putTMVar TMVar [Action]
actionsVar [Action]
rest
IO () -> STM (IO ())
forall a. a -> STM a
forall (m :: * -> *) a. Monad m => a -> m a
return (IO () -> STM (IO ())) -> IO () -> STM (IO ())
forall a b. (a -> b) -> a -> b
$ Action -> IO ()
actionRun Action
this IO () -> IO () -> IO ()
forall a b. IO a -> IO b -> IO b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> IO ()
go
findBool :: [Action] -> STM (Maybe (Action, [Action]))
findBool :: [Action] -> STM (Maybe (Action, [Action]))
findBool = [Action] -> [Action] -> STM (Maybe (Action, [Action]))
go []
where
go :: [Action] -> [Action] -> STM (Maybe (Action, [Action]))
go [] [] =
Maybe (Action, [Action]) -> STM (Maybe (Action, [Action]))
forall a. a -> STM a
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe (Action, [Action])
forall a. Maybe a
Nothing
go [Action]
_ [] =
STM (Maybe (Action, [Action]))
forall a. STM a
retry
go [Action]
past (Action
this : [Action]
rest) = do
status <- Action -> STM ActionStatus
actionStatus Action
this
case status of
ActionStatus
ActionReady -> Maybe (Action, [Action]) -> STM (Maybe (Action, [Action]))
forall a. a -> STM a
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe (Action, [Action]) -> STM (Maybe (Action, [Action])))
-> Maybe (Action, [Action]) -> STM (Maybe (Action, [Action]))
forall a b. (a -> b) -> a -> b
$ (Action, [Action]) -> Maybe (Action, [Action])
forall a. a -> Maybe a
Just (Action
this, [Action] -> [Action]
forall a. [a] -> [a]
reverse [Action]
past [Action] -> [Action] -> [Action]
forall a. [a] -> [a] -> [a]
++ [Action]
rest)
ActionStatus
ActionWait -> [Action] -> [Action] -> STM (Maybe (Action, [Action]))
go (Action
this Action -> [Action] -> [Action]
forall a. a -> [a] -> [a]
: [Action]
past) [Action]
rest
ActionStatus
ActionSkip -> do
Action -> STM ()
actionSkip Action
this
[Action] -> [Action] -> STM (Maybe (Action, [Action]))
go [Action]
past [Action]
rest