module Development.Shake.Pool(
Pool, runPool,
addPool, addPoolPriority,
increasePool
) where
import Control.Concurrent.Extra
import System.Time.Extra
import Control.Exception
import Control.Monad
import General.Timing
import qualified Data.HashSet as Set
import qualified Data.HashMap.Strict as Map
import System.Random
type NonDet a = IO a
data Queue a = Queue [a] (Either [a] (Tree a))
newQueue :: Bool -> Queue a
newQueue deterministic = Queue [] $ if deterministic then Left [] else Right emptyTree
enqueuePriority :: a -> Queue a -> Queue a
enqueuePriority x (Queue p t) = Queue (x:p) t
enqueue :: a -> Queue a -> Queue a
enqueue x (Queue p (Left xs)) = Queue p $ Left $ x:xs
enqueue x (Queue p (Right t)) = Queue p $ Right $ insertTree x t
dequeue :: Queue a -> NonDet (Maybe (a, Queue a))
dequeue (Queue (p:ps) t) = return $ Just (p, Queue ps t)
dequeue (Queue [] (Left (x:xs))) = return $ Just (x, Queue [] $ Left xs)
dequeue (Queue [] (Left [])) = return Nothing
dequeue (Queue [] (Right t)) = do
bs <- randomIO
return $ case removeTree bs t of
Nothing -> Nothing
Just (x,t) -> Just (x, Queue [] $ Right t)
data Tree a = Tree {-# UNPACK #-} !Int (Map.HashMap Int a)
emptyTree :: Tree a
emptyTree = Tree 0 Map.empty
insertTree :: a -> Tree a -> Tree a
insertTree x (Tree n mp) = Tree (n+1) $ Map.insert n x mp
removeTree :: Int -> Tree a -> Maybe (a, Tree a)
removeTree rnd (Tree n mp)
| n == 0 = Nothing
| n == 1 = Just (mp Map.! 0, emptyTree)
| i == n-1 = Just (mp Map.! i, Tree (n-1) $ Map.delete i mp)
| otherwise = Just (mp Map.! i, Tree (n-1) $ Map.insert i (mp Map.! (n-1)) $ Map.delete (n-1) mp)
where
i = abs rnd `mod` n
data Pool = Pool
!(Var (Maybe S))
!(Barrier (Either SomeException S))
data S = S
{threads :: !(Set.HashSet ThreadId)
,threadsLimit :: {-# UNPACK #-} !Int
,threadsMax :: {-# UNPACK #-} !Int
,threadsSum :: {-# UNPACK #-} !Int
,todo :: !(Queue (IO ()))
}
emptyS :: Int -> Bool -> S
emptyS n deterministic = S Set.empty n 0 0 $ newQueue deterministic
worker :: Pool -> IO ()
worker pool@(Pool var done) = do
let onVar act = modifyVar var $ maybe (return (Nothing, return ())) act
join $ onVar $ \s -> do
res <- dequeue $ todo s
case res of
Nothing -> return (Just s, return ())
Just (now, todo2) -> return (Just s{todo = todo2}, now >> worker pool)
step :: Pool -> (S -> NonDet S) -> IO ()
step pool@(Pool var done) op = do
let onVar act = modifyVar_ var $ maybe (return Nothing) act
onVar $ \s -> do
s <- op s
res <- dequeue $ todo s
case res of
Just (now, todo2) | Set.size (threads s) < threadsLimit s -> do
t <- forkFinally (now >> worker pool) $ \res -> case res of
Left e -> onVar $ \s -> do
t <- myThreadId
mapM_ killThread $ Set.toList $ Set.delete t $ threads s
signalBarrier done $ Left e
return Nothing
Right _ -> do
t <- myThreadId
step pool $ \s -> return s{threads = Set.delete t $ threads s}
let threads2 = Set.insert t $ threads s
return $ Just s{todo = todo2, threads = threads2
,threadsSum = threadsSum s + 1, threadsMax = threadsMax s `max` Set.size threads2}
Nothing | Set.null $ threads s -> do
signalBarrier done $ Right s
return Nothing
_ -> return $ Just s
addPool :: Pool -> IO a -> IO ()
addPool pool act = step pool $ \s -> do
todo <- return $ enqueue (void act) (todo s)
return s{todo = todo}
addPoolPriority :: Pool -> IO a -> IO ()
addPoolPriority pool act = step pool $ \s -> do
todo <- return $ enqueuePriority (void act) (todo s)
return s{todo = todo}
increasePool :: Pool -> IO (IO ())
increasePool pool = do
step pool $ \s -> return s{threadsLimit = threadsLimit s + 1}
return $ step pool $ \s -> return s{threadsLimit = threadsLimit s - 1}
runPool :: Bool -> Int -> (Pool -> IO ()) -> IO ()
runPool deterministic n act = do
s <- newVar $ Just $ emptyS n deterministic
done <- newBarrier
let cleanup = modifyVar_ s $ \s -> do
case s of
Just s -> mapM_ killThread $ Set.toList $ threads s
Nothing -> return ()
return Nothing
let ghc10793 = do
sleep 1
res <- waitBarrierMaybe done
case res of
Just (Left e) -> throwIO e
_ -> throwIO BlockedIndefinitelyOnMVar
handle (\BlockedIndefinitelyOnMVar -> ghc10793) $ flip onException cleanup $ do
let pool = Pool s done
addPool pool $ act pool
res <- waitBarrier done
case res of
Left e -> throwIO e
Right s -> addTiming $ "Pool finished (" ++ show (threadsSum s) ++ " threads, " ++ show (threadsMax s) ++ " max)"