-- minimal STM taskpool example which behaves strange: sometimes it hangs
-- after "finished" has been printed, i.e. after all worker are finished 
-- (and the forked threads should be terminated anyway since the main
-- thread has finished).
--
-- I'm testing the implementation in bash with
--
--   i=0; while true;do 
--      printf "%4d:\n" $((i=$i+1))
--      pool 8 64 1 +RTS -N8
--   done
--
-- on an 8-core machine. The parameters are
--
--   - the number of forked (light-weighted) threads
--   - the number of tasks in the pool
--   - the size of the tasks. for easy testing, this value should be
--     very small, e.g. 1 .. 10
--
-- This is just a simplified version of a more complex problem.

module Main where
import Control.Concurrent
import GHC.Conc
import Control.Monad
import qualified Data.Set as Set
import Data.Set (Set)
import Control.Concurrent.STM
import Data.Number.CReal
import System.Environment


main :: IO ()
main = do
    [threads,num,len] <- map read `fmap` getArgs 

    -- generate data structures and fill channel with tasks
    chan    <- newTChanIO
    working <- newTVarIO Set.empty
    atomically $ mapM (writeTChan chan) $ replicate num len

    -- fork threads and WAIT until all threads are finished. 
    replicateM_ threads (forkIO $ thread chan working)
    wait chan working 
    putStrLn "finished"


-- execute a task until the channel is empty and all workers are finished.
thread chan working = do
    task <- get chan working
    case task of
        Nothing -> return ()
        Just t  -> do 
            calcPi t 
            thread chan working
  where calcPi digits = showCReal (fromEnum digits) pi `pseq` return ()


wait chan working = do 
    atomically $ do
        work  <- readTVar working
        empty <- isEmptyTChan chan
        check (Set.null work && empty)


get chan working = do
    tid <- myThreadId

    -- atomically commit that this thread is not working anymore (since we 
    -- try to get a task we must be quasi-idle!
    atomically $ do
        work  <- Set.delete tid `fmap` readTVar working
        writeTVar working work

    -- waits for a new task. if all threads are idle and the pool is empty,
    -- return.
    atomically $ do
        empty <- isEmptyTChan chan
        work  <- readTVar working

        if (not empty) 
            then do 
                task <- readTChan chan
                writeTVar working (Set.insert tid work)
                return (Just task)
            else if Set.null work
                    then return Nothing
                    else retry