-- | A helper module which takes care of parallelism
{-# 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

-- | What to do about an 'Action'?
data ActionStatus
  = ActionReady
    -- ^ the action is ready to be executed
  | ActionSkip
    -- ^ the action should be skipped
  | ActionWait
    -- ^ not sure what to do yet; wait
  deriving ActionStatus -> ActionStatus -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: ActionStatus -> ActionStatus -> Bool
$c/= :: ActionStatus -> ActionStatus -> Bool
== :: ActionStatus -> ActionStatus -> Bool
$c== :: ActionStatus -> ActionStatus -> Bool
Eq

data Action = Action
  { Action -> STM ActionStatus
actionStatus :: STM ActionStatus
  , Action -> IO ()
actionRun :: IO ()
  , Action -> STM ()
actionSkip :: STM ()
  }

-- | Take a list of actions and execute them in parallel, no more than @n@
-- at the same time.
--
-- The action itself is asynchronous, ie. it returns immediately and does
-- the work in new threads. It returns an action which aborts tests and
-- cleans up.
runInParallel
  :: Int -- ^ maximum number of parallel threads
  -> [Action] -- ^ list of actions to execute.
    -- The first action in the pair tells if the second action is ready to run.
  -> IO (IO ())
-- This implementation tries its best to ensure that exceptions are
-- properly propagated to the caller and threads are not left running.
--
-- Note that exceptions inside tests are already caught by the test
-- actions themselves. Any exceptions that reach this function or its
-- threads are by definition unexpected.
runInParallel :: Int -> [Action] -> IO (IO ())
runInParallel Int
nthreads [Action]
actions = do
  ThreadId
callingThread <- IO ThreadId
myThreadId

  -- Don't let the main thread be garbage-collected
  -- Otherwise we may get a "thread blocked indefinitely in an STM
  -- transaction" exception when a child thread is blocked and GC'd.
  -- (See e.g. https://github.com/UnkindPartition/tasty/issues/15)
  -- FIXME is this still needed?
  StablePtr ThreadId
_ <- forall a. a -> IO (StablePtr a)
newStablePtr ThreadId
callingThread

  TMVar [Action]
actionsVar <- forall a. STM a -> IO a
atomically forall a b. (a -> b) -> a -> b
$ forall a. a -> STM (TMVar a)
newTMVar [Action]
actions

  [Async ()]
pids <- forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM Int
nthreads (forall a. IO a -> IO (Async a)
async forall a b. (a -> b) -> a -> b
$ TMVar [Action] -> IO ()
work TMVar [Action]
actionsVar)

  forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ do
    -- Tell worker threads there is no more work after their current task.
    -- 'cancel' below by itself is not sufficient because if an exception
    -- is thrown in the middle of a test, the worker thread simply marks
    -- the test as failed and moves on to their next task. We also need to
    -- make it clear that there are no further tasks.
    [Action]
_ <- forall a. STM a -> IO a
atomically forall a b. (a -> b) -> a -> b
$ forall a. TMVar a -> a -> STM a
swapTMVar TMVar [Action]
actionsVar []
    -- Cancel all the current tasks, waiting for workers to clean up.
    -- The waiting part is important (see #249), that's why we use cancel
    -- instead of killThread.
    forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ forall a. Async a -> IO ()
cancel [Async ()]
pids

work :: TMVar [Action] -> IO ()
work :: TMVar [Action] -> IO ()
work TMVar [Action]
actionsVar = IO ()
go
  where
    go :: IO ()
go = do
      forall (m :: * -> *) a. Monad m => m (m a) -> m a
join forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. STM a -> IO a
atomically forall a b. (a -> b) -> a -> b
$ do
        Maybe (Action, [Action])
mb_ready <- [Action] -> STM (Maybe (Action, [Action]))
findBool forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall a. TMVar a -> STM a
takeTMVar TMVar [Action]
actionsVar
        case Maybe (Action, [Action])
mb_ready of
          Maybe (Action, [Action])
Nothing -> do
            -- Nothing left to do. Put back the TMVar so that other threads
            -- do not block on an empty TMVar (see #249) and return.
            forall a. TMVar a -> a -> STM ()
putTMVar TMVar [Action]
actionsVar []
            forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a. Monad m => a -> m a
return ()
          Just (Action
this, [Action]
rest) -> do
            forall a. TMVar a -> a -> STM ()
putTMVar TMVar [Action]
actionsVar [Action]
rest
            forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ Action -> IO ()
actionRun Action
this forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> IO ()
go

-- | Find a ready-to-run item. Filter out the items that will never be
-- ready to run.
--
-- Return the ready item and the remaining ones.
--
-- This action may block if no items are ready to run just yet.
--
-- Return 'Nothing' if there are no runnable items left.
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 [] [] =
      -- nothing to do
      forall (m :: * -> *) a. Monad m => a -> m a
return forall a. Maybe a
Nothing
    go [Action]
_ [] =
      -- nothing ready yet
      forall a. STM a
retry
    go [Action]
past (Action
this : [Action]
rest) = do
      ActionStatus
status <- Action -> STM ActionStatus
actionStatus Action
this
      case ActionStatus
status of
        ActionStatus
ActionReady -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall a. a -> Maybe a
Just (Action
this, forall a. [a] -> [a]
reverse [Action]
past forall a. [a] -> [a] -> [a]
++ [Action]
rest)
        ActionStatus
ActionWait -> [Action] -> [Action] -> STM (Maybe (Action, [Action]))
go (Action
this 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