content: module Test.Framework.Runners.ThreadPool (
          executeOnPool
      ) where
  
  import Control.Concurrent
  import Control.Monad
  
  import qualified Data.IntMap as IM
  
  import Foreign.StablePtr
  
  
  data WorkerEvent token a = WorkerTermination
                           | WorkerItem token a
  
  -- | Execute IO actions on several threads and return their results in the original
  -- order.  It is guaranteed that no action from the input list is executed unless all
  -- the items that precede it in the list have been executed or are executing at that
  -- moment.
  executeOnPool :: Int    -- ^ Number of threads to use
                -> [IO a] -- ^ Actions to execute: these will be scheduled left to right
                -> IO [a] -- ^ Ordered results of executing the given IO actions in parallel
  executeOnPool n actions = do
      -- Prepare the channels
      input_chan <- newChan
      output_chan <- newChan
      
      -- Write the actions as items to the channel followed by one termination per thread
      -- that indicates they should terminate. We do this on another thread for
      -- maximum laziness (in case one the actions we are going to run depend on the
      -- output from previous actions..)
      _ <- forkIO $ writeList2Chan input_chan (zipWith WorkerItem [0..] actions ++ replicate n WorkerTermination)
      
      -- Spawn workers
      forM_ [1..n] (const $ forkIO $ poolWorker input_chan output_chan)
  
      -- Short version: make sure we do the right thing if a test blocks on dead
      -- MVars or TVars.
      -- Long version: GHC is clever enough to throw an exception (BlockedOnDeadMVar
      -- or BlockedIndefinitely) when a thread is waiting for a MVar or TVar that can't
      -- be written to. However, it doesn't know anything about the handlers for those
      -- exceptions. Therefore, when a worker runs a test that causes this exception,
      -- since the main thread is blocking on the worker, the main thread gets the
      -- exception too despite the fact that the main thread will be runnable as soon
      -- as the worker catches its own exception. The below makes sure the main thread
      -- is always reachable by the GC, which is the mechanism for finding threads
      -- that are unrunnable.
      -- See also the ticket where SimonM (semi-cryptically) explains this:
      -- http://hackage.haskell.org/trac/ghc/ticket/3291
      --
      -- NB: this actually leaks stable pointers. We could prevent this by making
      -- takeWhileWorkersExist do |unsafePerformIO newStablePtr| when returning the
      -- lazily-demanded tail of the list, but its a bit of a pain. For now, just
      -- grit our teeth and accept the leak.
      _stablePtr <- myThreadId >>= newStablePtr
      
      -- Return the results generated by the worker threads lazily and in
      -- the same order as we got the inputs
      fmap (reorderFrom 0 . takeWhileWorkersExist n) $ getChanContents output_chan
  
  poolWorker :: Chan (WorkerEvent token (IO a)) -> Chan (WorkerEvent token a) -> IO ()
  poolWorker input_chan output_chan = do
      -- Read an action and work out whether we should continue or stop
      action_item <- readChan input_chan
      case action_item of
          WorkerTermination -> writeChan output_chan WorkerTermination -- Must have run out of real actions to execute
          WorkerItem token action -> do
              -- Do the action then loop
              result <- action
              writeChan output_chan (WorkerItem token result)
              poolWorker input_chan output_chan
  
  -- | Keep grabbing items out of the infinite list of worker outputs until we have
  -- recieved word that all of the workers have shut down.  This lets us turn a possibly
  -- infinite list of outputs into a certainly finite one suitable for use with reorderFrom.
  takeWhileWorkersExist :: Int -> [WorkerEvent token a] -> [(token, a)]
  takeWhileWorkersExist worker_count events
    | worker_count <= 0 = []
    | otherwise         = case events of
                              (WorkerTermination:events')  -> takeWhileWorkersExist (worker_count - 1) events'
                              (WorkerItem token x:events') -> (token, x) : takeWhileWorkersExist worker_count events'
                              []                           -> []
  
  -- | This function carefully shuffles the input list so it in the total order
  -- defined by the integers paired with the elements.  If the list is @xs@ and
  -- the supplied initial integer is @n@, it must be the case that:
  --
  -- > sort (map fst xs) == [n..n + (length xs - 1)]
  --
  -- This function returns items in the lazy result list as soon as it is sure
  -- it has the right item for that position.
  reorderFrom :: Int -> [(Int, a)] -> [a]
  reorderFrom from initial_things = go from initial_things IM.empty False
    where go next [] buf _
            | IM.null buf = []    -- If the buffer and input list is empty, we're done
            | otherwise   = go next (IM.toList buf) IM.empty False    -- Make sure we check the buffer even if the list is done
          go next all_things@((token, x):things) buf buf_useful
            | token == next                       -- If the list token matches the one we were expecting we can just take the item
            = x : go (next + 1) things buf True   -- Always worth checking the buffer now because the expected item has changed
            | buf_useful                                                              -- If it's worth checking the buffer, it's possible the token we need is in it
            , (Just x', buf') <- IM.updateLookupWithKey (const $ const Nothing) next buf  -- Delete the found item from the map (if we find it) to save space
            = x' : go (next + 1) all_things buf' True                                     -- Always worth checking the buffer now because the expected item has changed
            | otherwise                                       -- Token didn't match, buffer unhelpful: it must be in the tail of the list
            = go next things (IM.insert token x buf) False    -- Since we've already checked the buffer, stop bothering to do so until something changes -}