module WorkGroup where

import Control.Monad (replicateM_)
import Control.Concurrent
import Control.Concurrent.QSem
import Control.Concurrent.MVar

-- | A WorkGroup is a manager for creating threads. When you create threads
-- using forkIOwg, you can wait for all them to exit using joinWorkGroup.
-- It's implemented because haskell doesn't have joinThread function.
data WorkGroup = WorkGroup {
      wgNum :: MVar Int
    , wgSem :: QSem
    }

newWorkGroup :: IO WorkGroup
newWorkGroup = do
  num <- newMVar 0
  sem <- newQSem 0
  return WorkGroup { wgNum=num, wgSem=sem }

-- | forkIO in workgroup. WorkGroup will keep track of started threads and
-- master thread can wait for all of them via joinWorkGroup wg.
forkIOwg :: WorkGroup -> IO () -> IO ThreadId
forkIOwg wg action = do
  modifyMVar_ (wgNum wg) (\n -> return $ n + 1)
  forkFinally action (\e -> signalQSem $ wgSem wg)

joinWorkGroup :: WorkGroup -> IO ()
joinWorkGroup wg = do
  n <- readMVar (wgNum wg)
  replicateM_ n $ waitQSem $ wgSem wg