module HaskellWorks.CabalCache.Concurrent.Fork
  ( forkThreadsWait,
  ) where

import Control.Exception (finally)
import HaskellWorks.Prelude

import qualified Control.Concurrent     as IO
import qualified Control.Concurrent.STM as STM

forkThreadsWait :: Int -> IO () -> IO ()
forkThreadsWait :: Int -> IO () -> IO ()
forkThreadsWait Int
n IO ()
f = do
  TVar Int
tDone <- STM (TVar Int) -> IO (TVar Int)
forall a. STM a -> IO a
STM.atomically (STM (TVar Int) -> IO (TVar Int))
-> STM (TVar Int) -> IO (TVar Int)
forall a b. (a -> b) -> a -> b
$ Int -> STM (TVar Int)
forall a. a -> STM (TVar a)
STM.newTVar (Int
0 :: Int)
  [Int] -> (Int -> IO ThreadId) -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [Int
1 .. Int
n] ((Int -> IO ThreadId) -> IO ()) -> (Int -> IO ThreadId) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Int
_ -> IO () -> IO ThreadId
IO.forkIO do
    IO ()
f IO () -> IO () -> IO ()
forall a b. IO a -> IO b -> IO a
`finally` (STM () -> IO ()
forall a. STM a -> IO a
STM.atomically (STM () -> IO ()) -> STM () -> IO ()
forall a b. (a -> b) -> a -> b
$ TVar Int -> (Int -> Int) -> STM ()
forall a. TVar a -> (a -> a) -> STM ()
STM.modifyTVar TVar Int
tDone (Int -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1))

  STM () -> IO ()
forall a. STM a -> IO a
STM.atomically do
    Int
done <- TVar Int -> STM Int
forall a. TVar a -> STM a
STM.readTVar TVar Int
tDone
    Bool -> STM () -> STM ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
done Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
n) STM ()
forall a. STM a
STM.retry