<commit_before>
-- | Thread pool implementation.
module Development.Shake.Pool(Pool, addPool, blockPool, runPool) where

import Control.Concurrent
import Control.Exception
import Control.Monad
import Data.IORef


data Pool = Pool (IORef (Maybe SomeException)) (MVar ()) (MVar Int)


-- | Add a new task to the pool
addPool :: Pool -> IO a -> IO ()
addPool (Pool err done count) act = do
    modifyMVar_ count $ \count -> do
        when (count == 0) $ takeMVar done
        return $ count + 1
    forkIO $ do
        res <- try act
        modifyMVar_ count $ \count -> do
            when (count == 1) $ putMVar done ()
            case res of
                Left e -> writeIORef err $ Just e
                _ -> return ()
            return $ count - 1
    return ()


-- | A blocking action is being run while on the pool, yeild your thread.
blockPool :: Pool -> IO a -> IO a
blockPool pool act = act


-- | Run all the tasks in the pool on the given number of works.
--   If any thread throws an exception, the exception will be reraised.
runPool :: Int -> (Pool -> IO ()) -> IO () -- run all tasks in the pool
runPool n act = do
    count <- newMVar 0
    done <- newMVar ()
    err <- newIORef Nothing
    act $ Pool err done count
    takeMVar done
    e <- readIORef err
    case e of
        Nothing -> return ()
        Just err -> throw err
<commit_msg>Rewrite the Pool module, now actually passes the tests<commit_after>{-# LANGUAGE RecordWildCards #-}

-- | Thread pool implementation.
module Development.Shake.Pool(Pool, addPool, blockPool, runPool) where

import Control.Concurrent
import Control.Exception hiding (blocked)
import Control.Monad
import Development.Shake.Locks
import qualified Data.HashSet as Set


---------------------------------------------------------------------
-- SUPER QUEUE

data SuperQueue a = SuperQueue [a] [a]

newSuperQueue :: SuperQueue a
newSuperQueue = SuperQueue [] []

enqueuePriority :: a -> SuperQueue a -> SuperQueue a
enqueuePriority x (SuperQueue p n) = SuperQueue (x:p) n

enqueue :: a -> SuperQueue a -> SuperQueue a
enqueue x (SuperQueue p n) = SuperQueue p (x:n)

dequeue :: SuperQueue a -> Maybe (a, SuperQueue a)
dequeue (SuperQueue (p:ps) ns) = Just (p, SuperQueue ps ns)
dequeue (SuperQueue [] (n:ns)) = Just (n, SuperQueue [] ns)
dequeue (SuperQueue [] []) = Nothing


---------------------------------------------------------------------
-- THREAD POOL

{-
Must keep a list of active threads, so can raise exceptions in a timely manner
Must spawn a fresh thread to do blockPool
If any worker throws an exception, must signal to all the other workers
-}

data Pool = Pool Int (Var (Maybe S)) (MVar (Maybe SomeException))

data S = S
    {threads :: Set.HashSet ThreadId
    ,working :: Int -- threads which are actively working
    ,blocked :: Int -- threads which are blocked
    ,todo :: SuperQueue (IO ())
    }


emptyS :: S
emptyS = S Set.empty 0 0 newSuperQueue


-- | Given a pool, and a function that breaks the S invariants, restore them
--   They are only allowed to touch working or todo
step :: Pool -> (S -> S) -> IO ()
step pool@(Pool n var done) op = do
    let onVar act = modifyVar_ var $ maybe (return Nothing) act
    onVar $ \s -> do
        s <- return $ op s
        case dequeue (todo s) of
            Just (now, todo2) | working s < n -> do
                -- spawn a new worker
                t <- forkIO $ do
                    t <- myThreadId
                    res <- try now
                    case res of
                        Left e -> onVar $ \s -> do
                            mapM_ killThread $ Set.toList $ Set.delete t $ threads s
                            putMVar done $ Just e
                            return Nothing
                        Right _ -> step pool $ \s -> s{working = working s - 1, threads = Set.delete t $ threads s}
                return $ Just s{working = working s + 1, todo = todo2, threads = Set.insert t $ threads s}
            Nothing | working s == 0 && blocked s == 0 -> do
                putMVar done Nothing
                return Nothing
            _ -> return $ Just s


-- | Add a new task to the pool
addPool :: Pool -> IO a -> IO ()
addPool pool act = step pool $ \s -> s{todo = enqueue (act >> return ()) (todo s)}


-- | A blocking action is being run while on the pool, yeild your thread.
--   Should only be called by an action under addPool.
blockPool :: Pool -> IO a -> IO a
blockPool pool act = do
    step pool $ \s -> s{working = working s - 1, blocked = blocked s + 1}
    res <- act
    var <- newEmptyMVar
    let act = do
            step pool $ \s -> s{working = working s + 1, blocked = blocked s - 1}
            putMVar var ()
    step pool $ \s -> s{todo = enqueuePriority act $ todo s}
    takeMVar var
    return res


-- | Run all the tasks in the pool on the given number of works.
--   If any thread throws an exception, the exception will be reraised.
runPool :: Int -> (Pool -> IO ()) -> IO () -- run all tasks in the pool
runPool n act = do
    s <- newVar $ Just emptyS
    res <- newEmptyMVar
    let pool = Pool n s res
    addPool pool $ act pool
    res <- readMVar res
    case res of
        Nothing -> return ()
        Just e -> throw e
{'fraction_non_alphanumeric': 0.12208189820130119, 'fraction_numerical': 0.0032529659395331037, 'mean_word_length': 1.9184812953657175, 'pattern_counts': {'":': 0, '<': 18, '<?xml version=': 0, '>': 47, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 3, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '629927', 'n_tokens_mistral': 1522, 'n_tokens_neox': 1454, 'n_words': 658}