    {-# LANGUAGE TupleSections #-}

    -- | Thread pool implementation. The three names correspond to the following
    --   priority levels (highest to lowest):
    --
    -- * 'addPoolException' - things that probably result in a build error,
    --   so kick them off quickly.
    --
    -- * 'addPoolResume' - things that started, blocked, and may have open
    --   resources in their closure.
    --
    -- * 'addPoolStart' - rules that haven't yet started.
    --
    -- * 'addPoolBatch' - rules that might batch if other rules start first.
    module General.Pool(
        Pool, runPool,
        addPool, PoolPriority(..),
        increasePool, keepAlivePool
        ) where

    import Control.Concurrent.Extra
    import System.Time.Extra
    import Control.Exception
    import Control.Monad.Extra
    import General.Timing
    import General.Thread
    import qualified Data.Heap as Heap
    import qualified Data.HashSet as Set
    import Data.IORef.Extra
    import System.Random

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

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

    data S = S
        {S -> Bool
    alive :: !Bool -- True until there's an exception, after which don't spawn more tasks
        ,S -> HashSet Thread
    threads :: !(Set.HashSet Thread) -- IMPORTANT: Must be strict or we leak thread stacks
        ,S -> Int
    threadsLimit :: {-# UNPACK #-} !Int -- user supplied thread limit, Set.size threads <= threadsLimit
        ,S -> Int
    threadsCount :: {-# UNPACK #-} !Int -- Set.size threads, but in O(1)
        ,S -> Int
    threadsMax :: {-# UNPACK #-} !Int -- high water mark of Set.size threads (accounting only)
        ,S -> Int
    threadsSum :: {-# UNPACK #-} !Int -- number of threads we have been through (accounting only)
        ,S -> IO Int
    rand :: IO Int -- operation to give us the next random Int
        ,S -> Heap (Entry (PoolPriority, Int) (IO ()))
    todo :: !(Heap.Heap (Heap.Entry (PoolPriority, Int) (IO ()))) -- operations waiting a thread
        }

    emptyS :: Int -> Bool -> IO S
    emptyS :: Int -> Bool -> IO S
    emptyS Int
    n Bool
    deterministic = do
        rand <- if Bool -> Bool
    not Bool
    deterministic then IO Int -> IO (IO Int)
    forall a. a -> IO a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure IO Int
    forall a (m :: * -> *). (Random a, MonadIO m) => m a
    randomIO else do
            ref <- Int -> IO (IORef Int)
    forall a. a -> IO (IORef a)
    newIORef Int
    0
            -- no need to be thread-safe - if two threads race they were basically the same time anyway
            pure $ do i <- readIORef ref; writeIORef' ref (i+1); pure i
        pure $ S True Set.empty n 0 0 0 rand Heap.empty

    data Pool = Pool
        !(Var S) -- Current state, 'alive' = False to say we are aborting
        !(Barrier (Either SomeException S)) -- Barrier to signal that we are finished

    withPool :: Pool -> (S -> IO (S, IO ())) -> IO ()
    withPool :: Pool -> (S -> IO (S, IO ())) -> IO ()
    withPool (Pool Var S
    var Barrier (Either SomeException S)
    _) S -> IO (S, IO ())
    f = IO (IO ()) -> IO ()
    forall (m :: * -> *) a. Monad m => m (m a) -> m a
    join (IO (IO ()) -> IO ()) -> IO (IO ()) -> IO ()
    forall a b. (a -> b) -> a -> b
    $ Var S -> (S -> IO (S, IO ())) -> IO (IO ())
    forall a b. Var a -> (a -> IO (a, b)) -> IO b
    modifyVar Var S
    var ((S -> IO (S, IO ())) -> IO (IO ()))
    -> (S -> IO (S, IO ())) -> IO (IO ())
    forall a b. (a -> b) -> a -> b
    $ \S
    s ->
        if S -> Bool
    alive S
    s then S -> IO (S, IO ())
    f S
    s else (S, IO ()) -> IO (S, IO ())
    forall a. a -> IO a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (S
    s, () -> IO ()
    forall a. a -> IO a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure ())

    withPool_ :: Pool -> (S -> IO S) -> IO ()
    withPool_ :: Pool -> (S -> IO S) -> IO ()
    withPool_ Pool
    pool S -> IO S
    act = Pool -> (S -> IO (S, IO ())) -> IO ()
    withPool Pool
    pool ((S -> IO (S, IO ())) -> IO ()) -> (S -> IO (S, IO ())) -> IO ()
    forall a b. (a -> b) -> a -> b
    $ (S -> (S, IO ())) -> IO S -> IO (S, IO ())
    forall a b. (a -> b) -> IO a -> IO b
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    fmap (, () -> IO ()
    forall a. a -> IO a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure()) (IO S -> IO (S, IO ())) -> (S -> IO S) -> S -> IO (S, IO ())
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . S -> IO S
    act

    worker :: Pool -> IO ()
    worker :: Pool -> IO ()
    worker Pool
    pool = Pool -> (S -> IO (S, IO ())) -> IO ()
    withPool Pool
    pool ((S -> IO (S, IO ())) -> IO ()) -> (S -> IO (S, IO ())) -> IO ()
    forall a b. (a -> b) -> a -> b
    $ \S
    s -> (S, IO ()) -> IO (S, IO ())
    forall a. a -> IO a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure ((S, IO ()) -> IO (S, IO ())) -> (S, IO ()) -> IO (S, IO ())
    forall a b. (a -> b) -> a -> b
    $ case Heap (Entry (PoolPriority, Int) (IO ()))
    -> Maybe
         (Entry (PoolPriority, Int) (IO ()),
          Heap (Entry (PoolPriority, Int) (IO ())))
    forall a. Heap a -> Maybe (a, Heap a)
    Heap.uncons (Heap (Entry (PoolPriority, Int) (IO ()))
     -> Maybe
          (Entry (PoolPriority, Int) (IO ()),
           Heap (Entry (PoolPriority, Int) (IO ()))))
    -> Heap (Entry (PoolPriority, Int) (IO ()))
    -> Maybe
         (Entry (PoolPriority, Int) (IO ()),
          Heap (Entry (PoolPriority, Int) (IO ())))
    forall a b. (a -> b) -> a -> b
    $ S -> Heap (Entry (PoolPriority, Int) (IO ()))
    todo S
    s of
        Maybe
      (Entry (PoolPriority, Int) (IO ()),
       Heap (Entry (PoolPriority, Int) (IO ())))
    Nothing -> (S
    s, () -> IO ()
    forall a. a -> IO a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure ())
        Just (Heap.Entry (PoolPriority, Int)
    _ IO ()
    now, Heap (Entry (PoolPriority, Int) (IO ()))
    todo2) -> (S
    s{todo = todo2}, IO ()
    now IO () -> IO () -> IO ()
    forall a b. IO a -> IO b -> IO b
    forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
    >> Pool -> IO ()
    worker Pool
    pool)

    -- | Given a pool, and a function that breaks the S invariants, restore them.
    --   They are only allowed to touch threadsLimit or todo.
    --   Assumes only requires spawning a most one job (e.g. can't increase the pool by more than one at a time)
    step :: Pool -> (S -> IO S) -> IO ()
    -- mask_ is so we don't spawn and not record it
    step :: Pool -> (S -> IO S) -> IO ()
    step pool :: Pool
    pool@(Pool Var S
    _ Barrier (Either SomeException S)
    done) S -> IO S
    op = IO () -> IO ()
    forall a. IO a -> IO a
    mask_ (IO () -> IO ()) -> IO () -> IO ()
    forall a b. (a -> b) -> a -> b
    $ Pool -> (S -> IO S) -> IO ()
    withPool_ Pool
    pool ((S -> IO S) -> IO ()) -> (S -> IO S) -> IO ()
    forall a b. (a -> b) -> a -> b
    $ \S
    s -> do
        s <- S -> IO S
    op S
    s
        case Heap.uncons $ todo s of
            Just (Heap.Entry (PoolPriority, Int)
    _ IO ()
    now, Heap (Entry (PoolPriority, Int) (IO ()))
    todo2) | S -> Int
    threadsCount S
    s Int -> Int -> Bool
    forall a. Ord a => a -> a -> Bool
    < S -> Int
    threadsLimit S
    s -> do
                -- spawn a new worker
                t <- IO () -> (Thread -> Either SomeException () -> IO ()) -> IO Thread
    forall a.
    IO a -> (Thread -> Either SomeException a -> IO ()) -> IO Thread
    newThreadFinally (IO ()
    now IO () -> IO () -> IO ()
    forall a b. IO a -> IO b -> IO b
    forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
    >> Pool -> IO ()
    worker Pool
    pool) ((Thread -> Either SomeException () -> IO ()) -> IO Thread)
    -> (Thread -> Either SomeException () -> IO ()) -> IO Thread
    forall a b. (a -> b) -> a -> b
    $ \Thread
    t Either SomeException ()
    res -> case Either SomeException ()
    res of
                    Left SomeException
    e -> Pool -> (S -> IO S) -> IO ()
    withPool_ Pool
    pool ((S -> IO S) -> IO ()) -> (S -> IO S) -> IO ()
    forall a b. (a -> b) -> a -> b
    $ \S
    s -> do
                        Barrier (Either SomeException S) -> Either SomeException S -> IO ()
    forall a. Partial => Barrier a -> a -> IO ()
    signalBarrier Barrier (Either SomeException S)
    done (Either SomeException S -> IO ())
    -> Either SomeException S -> IO ()
    forall a b. (a -> b) -> a -> b
    $ SomeException -> Either SomeException S
    forall a b. a -> Either a b
    Left SomeException
    e
                        S -> IO S
    forall a. a -> IO a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (Thread -> S -> S
    remThread Thread
    t S
    s){alive = False}
                    Right ()
    _ ->
                        Pool -> (S -> IO S) -> IO ()
    step Pool
    pool ((S -> IO S) -> IO ()) -> (S -> IO S) -> IO ()
    forall a b. (a -> b) -> a -> b
    $ S -> IO S
    forall a. a -> IO a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (S -> IO S) -> (S -> S) -> S -> IO S
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Thread -> S -> S
    remThread Thread
    t
                pure (addThread t s){todo = todo2}
            Maybe
      (Entry (PoolPriority, Int) (IO ()),
       Heap (Entry (PoolPriority, Int) (IO ())))
    Nothing | S -> Int
    threadsCount S
    s Int -> Int -> Bool
    forall a. Eq a => a -> a -> Bool
    == Int
    0 -> do
                Barrier (Either SomeException S) -> Either SomeException S -> IO ()
    forall a. Partial => Barrier a -> a -> IO ()
    signalBarrier Barrier (Either SomeException S)
    done (Either SomeException S -> IO ())
    -> Either SomeException S -> IO ()
    forall a b. (a -> b) -> a -> b
    $ S -> Either SomeException S
    forall a b. b -> Either a b
    Right S
    s
                S -> IO S
    forall a. a -> IO a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure S
    s{alive = False}
            Maybe
      (Entry (PoolPriority, Int) (IO ()),
       Heap (Entry (PoolPriority, Int) (IO ())))
    _ -> S -> IO S
    forall a. a -> IO a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure S
    s
        where
            addThread :: Thread -> S -> S
    addThread Thread
    t S
    s = S
    s{threads = Set.insert t $ threads s, threadsCount = threadsCount s + 1
                             ,threadsSum = threadsSum s + 1, threadsMax = threadsMax s `max` (threadsCount s + 1)}
            remThread :: Thread -> S -> S
    remThread Thread
    t S
    s = S
    s{threads = Set.delete t $ threads s, threadsCount = threadsCount s - 1}

    -- | Add a new task to the pool. See the top of the module for the relative ordering
    --   and semantics.
    addPool :: PoolPriority -> Pool -> IO a -> IO ()
    addPool :: forall a. PoolPriority -> Pool -> IO a -> IO ()
    addPool PoolPriority
    priority Pool
    pool IO a
    act = Pool -> (S -> IO S) -> IO ()
    step Pool
    pool ((S -> IO S) -> IO ()) -> (S -> IO S) -> IO ()
    forall a b. (a -> b) -> a -> b
    $ \S
    s -> do
        i <- S -> IO Int
    rand S
    s
        pure s{todo = Heap.insert (Heap.Entry (priority, i) $ void act) $ todo s}

    data PoolPriority
        = PoolException
        | PoolResume
        | PoolStart
        | PoolBatch
        | PoolDeprioritize Double
          deriving (PoolPriority -> PoolPriority -> Bool
    (PoolPriority -> PoolPriority -> Bool)
    -> (PoolPriority -> PoolPriority -> Bool) -> Eq PoolPriority
    forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
    $c== :: PoolPriority -> PoolPriority -> Bool
    == :: PoolPriority -> PoolPriority -> Bool
    $c/= :: PoolPriority -> PoolPriority -> Bool
    /= :: PoolPriority -> PoolPriority -> Bool
    Eq,Eq PoolPriority
    Eq PoolPriority =>
    (PoolPriority -> PoolPriority -> Ordering)
    -> (PoolPriority -> PoolPriority -> Bool)
    -> (PoolPriority -> PoolPriority -> Bool)
    -> (PoolPriority -> PoolPriority -> Bool)
    -> (PoolPriority -> PoolPriority -> Bool)
    -> (PoolPriority -> PoolPriority -> PoolPriority)
    -> (PoolPriority -> PoolPriority -> PoolPriority)
    -> Ord PoolPriority
    PoolPriority -> PoolPriority -> Bool
    PoolPriority -> PoolPriority -> Ordering
    PoolPriority -> PoolPriority -> PoolPriority
    forall a.
    Eq a =>
    (a -> a -> Ordering)
    -> (a -> a -> Bool)
    -> (a -> a -> Bool)
    -> (a -> a -> Bool)
    -> (a -> a -> Bool)
    -> (a -> a -> a)
    -> (a -> a -> a)
    -> Ord a
    $ccompare :: PoolPriority -> PoolPriority -> Ordering
    compare :: PoolPriority -> PoolPriority -> Ordering
    $c< :: PoolPriority -> PoolPriority -> Bool
    < :: PoolPriority -> PoolPriority -> Bool
    $c<= :: PoolPriority -> PoolPriority -> Bool
    <= :: PoolPriority -> PoolPriority -> Bool
    $c> :: PoolPriority -> PoolPriority -> Bool
    > :: PoolPriority -> PoolPriority -> Bool
    $c>= :: PoolPriority -> PoolPriority -> Bool
    >= :: PoolPriority -> PoolPriority -> Bool
    $cmax :: PoolPriority -> PoolPriority -> PoolPriority
    max :: PoolPriority -> PoolPriority -> PoolPriority
    $cmin :: PoolPriority -> PoolPriority -> PoolPriority
    min :: PoolPriority -> PoolPriority -> PoolPriority
    Ord)

    -- | Temporarily increase the pool by 1 thread. Call the cleanup action to restore the value.
    --   After calling cleanup you should requeue onto a new thread.
    increasePool :: Pool -> IO (IO ())
    increasePool :: Pool -> IO (IO ())
    increasePool Pool
    pool = do
        Pool -> (S -> IO S) -> IO ()
    step Pool
    pool ((S -> IO S) -> IO ()) -> (S -> IO S) -> IO ()
    forall a b. (a -> b) -> a -> b
    $ \S
    s -> S -> IO S
    forall a. a -> IO a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure S
    s{threadsLimit = threadsLimit s + 1}
        IO () -> IO (IO ())
    forall a. a -> IO a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (IO () -> IO (IO ())) -> IO () -> IO (IO ())
    forall a b. (a -> b) -> a -> b
    $ Pool -> (S -> IO S) -> IO ()
    step Pool
    pool ((S -> IO S) -> IO ()) -> (S -> IO S) -> IO ()
    forall a b. (a -> b) -> a -> b
    $ \S
    s -> S -> IO S
    forall a. a -> IO a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure S
    s{threadsLimit = threadsLimit s - 1}

    -- | Make sure the pool cannot run out of tasks (and thus everything finishes) until after the cancel is called.
    --   Ensures that a pool that will requeue in time doesn't go idle.
    keepAlivePool :: Pool -> IO (IO ())
    keepAlivePool :: Pool -> IO (IO ())
    keepAlivePool Pool
    pool = do
        bar <- IO (Barrier ())
    forall a. IO (Barrier a)
    newBarrier
        addPool PoolResume pool $ do
            cancel <- increasePool pool
            waitBarrier bar
            cancel
        pure $ signalBarrier bar ()

    -- | 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 :: Bool -> Int -> (Pool -> IO ()) -> IO () -- run all tasks in the pool
    runPool :: Bool -> Int -> (Pool -> IO ()) -> IO ()
    runPool Bool
    deterministic Int
    n Pool -> IO ()
    act = do
        s <- S -> IO (Var S)
    forall a. a -> IO (Var a)
    newVar (S -> IO (Var S)) -> IO S -> IO (Var S)
    forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
    =<< Int -> Bool -> IO S
    emptyS Int
    n Bool
    deterministic
        done <- newBarrier
        let pool = Var S -> Barrier (Either SomeException S) -> Pool
    Pool Var S
    s Barrier (Either SomeException S)
    done

        -- if someone kills our thread, make sure we kill our child threads
        let cleanup = IO (IO ()) -> IO ()
    forall (m :: * -> *) a. Monad m => m (m a) -> m a
    join (IO (IO ()) -> IO ()) -> IO (IO ()) -> IO ()
    forall a b. (a -> b) -> a -> b
    $ Var S -> (S -> IO (S, IO ())) -> IO (IO ())
    forall a b. Var a -> (a -> IO (a, b)) -> IO b
    modifyVar Var S
    s ((S -> IO (S, IO ())) -> IO (IO ()))
    -> (S -> IO (S, IO ())) -> IO (IO ())
    forall a b. (a -> b) -> a -> b
    $ \S
    s -> (S, IO ()) -> IO (S, IO ())
    forall a. a -> IO a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (S
    s{alive=False}, [Thread] -> IO ()
    stopThreads ([Thread] -> IO ()) -> [Thread] -> IO ()
    forall a b. (a -> b) -> a -> b
    $ HashSet Thread -> [Thread]
    forall a. HashSet a -> [a]
    Set.toList (HashSet Thread -> [Thread]) -> HashSet Thread -> [Thread]
    forall a b. (a -> b) -> a -> b
    $ S -> HashSet Thread
    threads S
    s)

        let ghc10793 = do
                -- if this thread dies because it is blocked on an MVar there's a chance we have
                -- a better error in the done barrier, and GHC raised the exception wrongly, see:
                -- https://ghc.haskell.org/trac/ghc/ticket/10793
                Double -> IO ()
    sleep Double
    1 -- give it a little bit of time for the finally to run
                        -- no big deal, since the blocked indefinitely takes a while to fire anyway
                res <- Barrier (Either SomeException S)
    -> IO (Maybe (Either SomeException S))
    forall a. Barrier a -> IO (Maybe a)
    waitBarrierMaybe Barrier (Either SomeException S)
    done
                case res of
                    Just (Left SomeException
    e) -> SomeException -> IO b
    forall e a. (Partial, Exception e) => e -> IO a
    throwIO SomeException
    e
                    Maybe (Either SomeException S)
    _ -> BlockedIndefinitelyOnMVar -> IO b
    forall e a. (Partial, Exception e) => e -> IO a
    throwIO BlockedIndefinitelyOnMVar
    BlockedIndefinitelyOnMVar
        flip finally cleanup $ handle (\BlockedIndefinitelyOnMVar
    BlockedIndefinitelyOnMVar -> IO ()
    forall {b}. IO b
    ghc10793) $ do
            addPool PoolStart pool $ act pool
            res <- waitBarrier done
            case res of
                Left SomeException
    e -> SomeException -> IO ()
    forall e a. (Partial, Exception e) => e -> IO a
    throwIO SomeException
    e
                Right S
    s -> String -> IO ()
    addTiming (String -> IO ()) -> String -> IO ()
    forall a b. (a -> b) -> a -> b
    $ String
    "Pool finished (" String -> String -> String
    forall a. [a] -> [a] -> [a]
    ++ Int -> String
    forall a. Show a => a -> String
    show (S -> Int
    threadsSum S
    s) String -> String -> String
    forall a. [a] -> [a] -> [a]
    ++ String
    " threads, " String -> String -> String
    forall a. [a] -> [a] -> [a]
    ++ Int -> String
    forall a. Show a => a -> String
    show (S -> Int
    threadsMax S
    s) String -> String -> String
    forall a. [a] -> [a] -> [a]
    ++ String
    " max)"