hgraph-1.10.0.0: src/HGraph/Parallel.hs

{-# Language CPP #-}

module HGraph.Parallel
       ( produceConsumeExhaustive
       , processJobList
       )
where

import HGraph.Debugging

import Control.Concurrent
import Data.Maybe
import Control.Monad
import Control.Concurrent.MVar
import Control.Exception
import System.IO (hPutStrLn, stderr)

-- | Runs the producer until it output `Nothing`.
-- Each successful product is sent to the consumer.
-- Returns an IO action which takes the next product (or `Nothing` if there are no more products)
produceConsumeExhaustive :: (IO (Maybe a)) -> (a -> IO b) -> IO (IO (Maybe b))
produceConsumeExhaustive produce consume = do
  numForks <- getNumCapabilities
  mProduct <- newEmptyMVar
  workerDone <- replicateM numForks $ newEmptyMVar
  producers <- forM workerDone $ \doneM -> forkFinally
    ( let loop = do
              g <- hasLocked (__FILE__ ++ " on line " ++ show __LINE__) $ produce
              when (isJust g) $ do
                hasLocked (__FILE__ ++ " on line " ++ show __LINE__) $ putMVar mProduct g
                loop
      in loop
    )
    (\e -> do
      hasLocked (__FILE__ ++ " on line " ++ show __LINE__) $ putMVar doneM ()
      case e of
        Left err -> (sayErrString (__FILE__ ++ " on line " ++ show __LINE__
                                   ++ " produce-consume error (producer):\n\t -"
                                   ++ show (err :: SomeException)))
        Right res -> do
          return ()
    )
  consumerDone <- replicateM numForks $ newEmptyMVar
  consumerOutput <- newEmptyMVar
  consumers <- forM consumerDone $ \doneM -> forkFinally
    ( let loop = do
                  g <- hasLocked (__FILE__ ++ " on line " ++ show __LINE__) $ takeMVar mProduct
                  case g of
                    Just p -> do
                      out <- hasLocked (__FILE__ ++ " on line " ++ show __LINE__) $ consume p
                      hasLocked (__FILE__ ++ " on line " ++ show __LINE__) $ putMVar consumerOutput (Just out)
                      loop
                    Nothing -> hasLocked (__FILE__ ++ " on line " ++ show __LINE__) $ putMVar mProduct Nothing
      in loop
    )
    (\e -> do
      hasLocked (__FILE__ ++ " on line " ++ show __LINE__) $ putMVar doneM ()
      case e of
        Left err -> (sayErrString (__FILE__ ++ " on line " ++ show __LINE__ ++ ": produce-consume error (consumer):\n\t- " ++ show (err :: SomeException)))
        Right res -> do
          return ()
    )
  forkFinally (do
    hasLocked (__FILE__ ++ " on line " ++ show __LINE__) $ mapM_ takeMVar workerDone
    hasLocked (__FILE__ ++ " on line " ++ show __LINE__) $ putMVar mProduct Nothing
    hasLocked (__FILE__ ++ " on line " ++ show __LINE__) $ mapM_ takeMVar consumerDone
    )
    (\res -> do
      hasLocked (__FILE__ ++ " on line " ++ show __LINE__) $ putMVar consumerOutput Nothing
      case res of
        Left err -> sayErrString (__FILE__ ++ " on line " ++ show __LINE__
                                  ++ ": produce-consume error (final):\n\t- "
                                  ++ show (err :: SomeException))
        Right _ -> return ()
    )

  return (takeMVar consumerOutput)

processJobList :: (a -> IO ([a], [b])) -> [a] -> IO (IO (Maybe b))
processJobList worker jobs = do
  numForks <- getNumCapabilities
  workerDone <- replicateM numForks $ newEmptyMVar
  workerOutput <- newEmptyMVar
  mReady <- newMVar ()
  mJobs  <- newMVar jobs
  mBusyCount <- newMVar 0

  _ <- forM workerDone $ \doneM -> forkFinally
    ( let loop = do
              hasLocked (__FILE__ ++ " on line " ++ show __LINE__) $ takeMVar mReady
              openJobs <- hasLocked (__FILE__ ++ " on line " ++ show __LINE__) $ takeMVar mJobs
              busy0 <- modifyMVar mBusyCount (\c -> return $ (c + 1, c + 1))
              case openJobs of
                [] -> do
                  modifyMVar_ mBusyCount (\c -> return $ c - 1)
                  if busy0 == 1 then do
                    hasLocked (__FILE__ ++ " on line " ++ show __LINE__) $ putMVar mReady ()
                    hasLocked (__FILE__ ++ " on line " ++ show __LINE__) $ putMVar mJobs []
                  else
                    loop
                (j : jobs') -> do
                  when (not $ null jobs') $ 
                    hasLocked (__FILE__ ++ " on line " ++ show __LINE__) $ putMVar mReady ()
                  hasLocked (__FILE__ ++ " on line " ++ show __LINE__) $ putMVar mJobs jobs'
                  (newJobs, newResults) <- worker j
                  openJobs'' <- modifyMVar mJobs $ \openJobs' -> do
                    return $ (newJobs ++ openJobs', openJobs')
                  busy <- modifyMVar mBusyCount (\c -> return $ (c - 1, c - 1))
                  when (null openJobs'' && ( (not $ null newJobs) || busy == 0)) $ do
                    hasLocked (__FILE__ ++ " on line " ++ show __LINE__) $ putMVar mReady ()
                  hasLocked (__FILE__ ++ " on line " ++ show __LINE__) $ mapM_ ((putMVar workerOutput) . Just) newResults
                  loop
      in loop
    )
    (\e -> do
      hasLocked (__FILE__ ++ " on line " ++ show __LINE__) $ putMVar doneM ()
      case e of
        Left err -> (sayErrString (__FILE__ ++ " on line " ++ show __LINE__
                                   ++ " job list error (worker):\n\t -"
                                   ++ show (err :: SomeException)))
        Right res -> do
          return ()
    )

  forkFinally (do
    hasLocked (__FILE__ ++ " on line " ++ show __LINE__) $ mapM_ takeMVar workerDone
    )
    (\res -> do
      hasLocked (__FILE__ ++ " on line " ++ show __LINE__) $ putMVar workerOutput Nothing
      case res of
        Left err -> sayErrString (__FILE__ ++ " on line " ++ show __LINE__
                                  ++ ": job list error (end):\n\t- "
                                  ++ show (err :: SomeException))
        Right _ -> return ()
    )

  return (takeMVar workerOutput)