Merge pull request #2 from timthelion/master

Support for grandchild threads.

Test/GrandChildTest.lhs

@@ -0,0 +1,35 @@
+What happens if I create a tree of threads like:
+
+ThreadMain
+ ChildThread
+  GrandChildThread
+
+And ChildThread ends?  Normally GrandChildThread dies.  This module tests whether ThreadManager successfully flattens this thread hierarchy to:
+
+ThreadMain
+ ChildThread
+ GrandChildThread
+
+thus saving GrandChildThread in the case of ChildThread's death.
+
+>import qualified Control.Concurrent.ThreadManager as TM
+>import Control.Concurrent
+
+>main :: IO ()
+>main = do
+> tm <- TM.make
+> TM.fork tm (childThread tm)
+> TM.waitForAll tm
+> print "If you see this message, your grandChild has died."
+
+
+>childThread :: TM.ThreadManager -> IO ()
+>childThread tm = do
+> TM.fork tm (grandChildThread 0)
+
+> return ()
+
+>grandChildThread :: Int ->  IO ()
+>grandChildThread count = 
+> print count >> threadDelay 1000 >> (grandChildThread $ count+1)
+

src/Control/Concurrent/ThreadManager.hs

@@ -15,9 +15,9 @@ module Control.Concurrent.ThreadManager
   ) where
 
 import Control.Concurrent      (ThreadId, forkIO)
-import Control.Concurrent.MVar (MVar, modifyMVar, newEmptyMVar, newMVar, putMVar, takeMVar, tryTakeMVar)
+import Control.Concurrent.MVar (MVar, modifyMVar, newEmptyMVar, newMVar, putMVar, takeMVar, tryTakeMVar,readMVar)
 import Control.Exception       (SomeException, try)
-import Control.Monad           (join, replicateM)
+import Control.Monad           (join, replicateM, when)
 import qualified Data.Map as M
 
 data ThreadStatus =
@@ -70,7 +70,15 @@ waitFor (TM tm) tid =
 
 -- | Block until all managed threads terminate.
 waitForAll :: ThreadManager -> IO ()
-waitForAll (TM tm) =
-    modifyMVar tm elems >>= mapM_ takeMVar
+waitForAll tm@(TM tmMvar) = do
+    threadMap <- readMVar tmMvar
+    let threads = M.keys threadMap
+    statuses <- mapM (getStatus tm) threads
+    _ <- mapM (waitFor tm) threads
+    Control.Monad.when (foldr checkStatus False statuses) $
+        waitForAll tm
   where
-    elems m = return (M.empty, M.elems m)
+    checkStatus :: Maybe ThreadStatus -> Bool -> Bool
+    checkStatus _ True = True
+    checkStatus (Just Running) False = True
+    checkStatus _ False = False