[Haskell-cafe] Re: Waiting for thread to finish

[Matthew Brecknell]

Brad Clow:
> If you would like to wait on multiple threads, you can use STM like so:
> 
> import Control.Concurrent
> import Control.Concurrent.STM
> import Control.Exception
> 
> main = do
>   tc <- atomically $ newTVar 2
>   run tc (print (last [1..100000000]))
>   run tc (print (last [1..110000000]))
>   print "Waiting...."
>   atomically $ readTVar tc >>= \x -> if x == 0 then return () else retry
>   print "OK."
>   where
>     run tc f = forkIO (f `finally` atomReplace (\x -> x - 1) tc)
> 
> atomReplace fn x = atomically $ readTVar x >>= writeTVar x . fn

Nice! Although, to wait for all of a set of threads, you really only
need to wait for each in turn, so you could do this with plain MVars.
The real power of STM becomes apparent when you need to wait for any of
a set of results, for example:

> import Control.Arrow
> import Control.Concurrent
> import Control.Concurrent.STM
> import Control.Concurrent.STM.TVar
> 
> newtype Wait a = Wait (TVar (Maybe a))
> 
> fork :: IO a -> IO (Wait a)
> fork m = do
>   w <- atomically (newTVar Nothing)
>   forkIO (m >>= atomically . writeTVar w . Just)
>   return (Wait w)
> 
> wait :: Wait a -> IO a
> wait (Wait w) = atomically $ do
>   r <- readTVar w
>   case r of
>     Just a -> return a
>     Nothing -> retry
> 
> wait_all :: [Wait a] -> IO [a]
> wait_all [] = return []
> wait_all (w:ws) = do
>   r <- wait w
>   t <- wait_all ws
>   return (r:t)
> 
> wait_first :: [Wait a] -> IO (a, [Wait a])
> wait_first [] = error "wait_first: nothing to wait for"
> wait_first ws = atomically (do_wait ws) where
>   do_wait [] = retry
>   do_wait (Wait w : ws) = do
>     r <- readTVar w
>     case r of
>       Nothing -> fmap (second (Wait w:)) (do_wait ws)
>       Just s -> return (s,ws)
> 
> main = do
>   w1 <- fork (test 50000000)
>   w2 <- fork (test 10000000)
>   w3 <- fork (test 100000000)
>   (r,ws) <- wait_first [w1,w2,w3]
>   putStrLn ("First result: " ++ show r)
>   rs <- wait_all ws
>   putStrLn ("Remaining results: " ++ show rs)
> 
> test :: Integer -> IO Integer
> test i = do
>   let r = last [1..i]
>   putStrLn ("Result " ++ show r)
>   return r

You might recognise the Wait type as being identical to TMVar, although
I use a slightly different set of operations. Throw
Control.Concurrent.STM.TChan into the mix, and you have some very rich
possibilities indeed.