3 files changed, 46 insertions, 2 deletions

diff --git a/src/Control/Concurrent/Tasks.hs b/src/Control/Concurrent/Tasks.hs
new file mode 100644
index 00000000..287542ee
--- /dev/null
+++ b/src/Control/Concurrent/Tasks.hs
@@ -0,0 +1,44 @@
+{-# LANGUAGE CPP #-}
+module Control.Concurrent.Tasks where
+
+import Control.Concurrent.STM
+import Control.Exception
+import Data.Function
+import Data.List
+#ifdef THREAD_DEBUG
+import Control.Concurrent.Lifted.Instrument
+#else
+import Control.Concurrent.Lifted
+import GHC.Conc                  (labelThread)
+#endif
+
+newtype TaskGroup = TaskGroup
+    { taskQueue :: TChan (String,IO ())
+    }
+
+withTaskGroup :: String -> Int -> (TaskGroup -> IO ()) -> IO ()
+withTaskGroup glabel numslots action = do
+    tg <- atomically $ newTChan
+    cnt <- atomically $ newTVar 0
+    thread <- fork $ do
+        myThreadId >>= flip labelThread glabel
+        fix $ \again -> do
+            (slot, (lbl,task)) <- atomically $ do
+                slot <- readTVar cnt
+                check (slot < numslots)
+                writeTVar cnt (succ slot)
+                t <- readTChan tg
+                return (slot,t)
+            _ <- fork $ do
+                myThreadId >>= flip labelThread (intercalate "." [glabel,show slot,lbl])
+                task `catch` (\(SomeException e) -> return ())
+                atomically $ modifyTVar' cnt pred
+            again
+    action (TaskGroup tg) `onException` killThread thread
+    atomically $ do
+        isEmptyTChan tg >>= check
+        readTVar cnt >>= check . (== 0)
+    killThread thread
+
+forkTask :: TaskGroup -> String -> IO () -> IO ()
+forkTask (TaskGroup q) lbl action = atomically $ writeTChan q (lbl,action)
diff --git a/src/Network/BitTorrent/DHT/Search.hs b/src/Network/BitTorrent/DHT/Search.hs
index 16a20620..29e91633 100644
--- a/src/Network/BitTorrent/DHT/Search.hs
+++ b/src/Network/BitTorrent/DHT/Search.hs
@@ -6,7 +6,7 @@
 {-# LANGUAGE LambdaCase #-}
 module Network.BitTorrent.DHT.Search where
 
-import Tasks
+import Control.Concurrent.Tasks
 import Control.Concurrent.STM
 import Control.Exception
 import Control.Monad
diff --git a/src/Network/Kademlia.hs b/src/Network/Kademlia.hs
index c6c59ae6..53c37175 100644
--- a/src/Network/Kademlia.hs
+++ b/src/Network/Kademlia.hs
@@ -35,7 +35,7 @@ import System.Timeout
 import Text.PrettyPrint               as PP hiding (($$), (<>))
 import Text.PrettyPrint.HughesPJClass hiding (($$), (<>))
 import System.IO
-import Tasks
+import Control.Concurrent.Tasks
 
 -- | The status of a given node with respect to a given routint table.
 data RoutingStatus