4 files changed, 37 insertions, 7 deletions

diff --git a/src/Network/BitTorrent/DHT/Query.hs b/src/Network/BitTorrent/DHT/Query.hs
index 73b3d492..47d17622 100644
--- a/src/Network/BitTorrent/DHT/Query.hs
+++ b/src/Network/BitTorrent/DHT/Query.hs
@@ -9,6 +9,7 @@
 --   Normally, you don't need to import this module, use
 --   "Network.BitTorrent.DHT" instead.
 --
+{-# LANGUAGE CPP                 #-}
 {-# LANGUAGE FlexibleContexts    #-}
 {-# LANGUAGE ScopedTypeVariables #-}
 {-# LANGUAGE TemplateHaskell     #-}
@@ -49,7 +50,12 @@ module Network.BitTorrent.DHT.Query
        , (<@>)
        ) where
 
+#ifdef THREAD_DEBUG
+import Control.Concurrent.Lifted.Instrument hiding (yield)
+#else
+import GHC.Conc (labelThread)
 import Control.Concurrent.Lifted hiding (yield)
+#endif
 import Control.Exception.Lifted hiding (Handler)
 import Control.Monad.Reader
 import Control.Monad.Logger
@@ -227,6 +233,7 @@ refreshNodes nid = do
 --   routing table.
 insertNode :: forall ip. Address ip => NodeInfo ip -> Maybe ReflectedIP -> DHT ip ThreadId
 insertNode info witnessed_ip0 = fork $ do
+  -- myThreadId >>= liftIO . flip labelThread "DHT.insertNode"
   var <- asks routingInfo
   tm <- getTimestamp
   let showTable = do
@@ -286,8 +293,10 @@ insertNode info witnessed_ip0 = fork $ do
                            return ps
   ps <- join $ liftIO $ atomically $ atomicInsert arrival0 witnessed_ip0
   showTable
-  _ <- fork $ forM_ ps $ \(CheckPing ns)-> do
-    forM_ ns $ \n -> do
+  _ <- fork $ do
+    myThreadId >>= liftIO . flip labelThread "DHT.insertNode.pingResults"
+    forM_ ps $ \(CheckPing ns)-> do
+      forM_ ns $ \n -> do
         (b,mip) <- probeNode (nodeAddr n)
         let alive = PingResult n b
         $(logDebugS) "insertNode" $ T.pack ("PingResult "++show (nodeId n,b))
diff --git a/src/Network/BitTorrent/DHT/Session.hs b/src/Network/BitTorrent/DHT/Session.hs
index 8dc3f7ac..bad783a5 100644
--- a/src/Network/BitTorrent/DHT/Session.hs
+++ b/src/Network/BitTorrent/DHT/Session.hs
@@ -72,7 +72,11 @@ module Network.BitTorrent.DHT.Session
 import Prelude hiding (ioError)
 
 import Control.Concurrent.STM
+#ifdef THREAD_DEBUG
+import Control.Concurrent.Async.Lifted.Instrument
+#else
 import Control.Concurrent.Async.Lifted
+#endif
 import Control.Exception.Lifted hiding (Handler)
 import Control.Monad.Base
 import Control.Monad.Logger
diff --git a/src/Network/KRPC/Manager.hs b/src/Network/KRPC/Manager.hs
index 4852eb38..22d111e2 100644
--- a/src/Network/KRPC/Manager.hs
+++ b/src/Network/KRPC/Manager.hs
@@ -7,6 +7,7 @@
 --
 --   Normally, you don't need to import this module.
 --
+{-# LANGUAGE CPP                    #-}
 {-# LANGUAGE OverloadedStrings      #-}
 {-# LANGUAGE FlexibleInstances      #-}
 {-# LANGUAGE FlexibleContexts       #-}
@@ -41,8 +42,12 @@ module Network.KRPC.Manager
        ) where
 
 import Control.Applicative
-import Control.Concurrent
-import Control.Concurrent.Lifted (fork)
+#ifdef THREAD_DEBUG
+import Control.Concurrent.Lifted.Instrument
+#else
+import GHC.Conc (labelThread)
+import Control.Concurrent.Lifted
+#endif
 import Control.Exception hiding (Handler)
 import qualified Control.Exception.Lifted as E (Handler (..))
 import Control.Exception.Lifted as Lifted (catches, finally)
@@ -432,6 +437,7 @@ dispatchHandler q @ KQuery {..} addr = do
 --
 handleQuery :: MonadKRPC h m => BValue -> KQuery -> SockAddr -> m ()
 handleQuery raw q addr = void $ fork $ do
+  myThreadId >>= liftIO . flip labelThread "KRPC.handleQuery"
   Manager {..} <- getManager
   res <- dispatchHandler q addr
   let resbe = either toBEncode toBEncode res
@@ -481,6 +487,7 @@ listen :: MonadKRPC h m => m ()
 listen = do
   Manager {..} <- getManager
   tid <- fork $ do
+    myThreadId >>= liftIO . flip labelThread "KRPC.listen"
     listener `Lifted.finally`
       liftIO (takeMVar listenerThread)
   liftIO $ putMVar listenerThread tid
diff --git a/src/Network/StreamServer.hs b/src/Network/StreamServer.hs
index a6cead0e..34b9388e 100644
--- a/src/Network/StreamServer.hs
+++ b/src/Network/StreamServer.hs
@@ -1,4 +1,5 @@
 -- | This module implements a bare-bones TCP or Unix socket server.
+{-# LANGUAGE CPP #-}
 {-# LANGUAGE TypeFamilies #-}
 {-# LANGUAGE TypeOperators #-}
 {-# LANGUAGE OverloadedStrings #-}
@@ -34,7 +35,12 @@ import System.IO
     )
 import Control.Monad
 import Control.Monad.Fix (fix)
-import Control.Concurrent (forkIO, threadDelay, ThreadId, mkWeakThreadId)
+#ifdef THREAD_DEBUG
+import Control.Concurrent.Lifted.Instrument (forkIO, threadDelay, ThreadId, mkWeakThreadId, labelThread, myThreadId)
+#else
+import GHC.Conc (labelThread)
+import Control.Concurrent (forkIO, threadDelay, ThreadId, mkWeakThreadId, myThreadId)
+#endif
 import Control.Exception (catch,handle,try,finally)
 import System.IO.Error (tryIOError)
 import System.Mem.Weak
@@ -113,7 +119,9 @@ streamServer cfg addr = do
                                 threadDelay 5000000
                                 loop
     listen sock maxListenQueue
-    thread <- mkWeakThreadId <=< forkIO $ acceptLoop cfg sock 0
+    thread <- mkWeakThreadId <=< forkIO $ do
+        myThreadId >>= flip labelThread "StreamServer.acceptLoop"
+        acceptLoop cfg sock 0
     return (ServerHandle sock thread)
 
 -- | Not exported.  This, combined with 'acceptException' form a mutually recursive
@@ -124,7 +132,9 @@ acceptLoop cfg sock n = handle (acceptException cfg n sock) $ do
     con <- accept sock
     let conkey = n + 1
     h <- socketToHandle (fst con) ReadWriteMode
-    forkIO $ serverSession cfg (restrictHandleSocket h (fst con)) conkey h
+    forkIO $ do
+        myThreadId >>= flip labelThread "StreamServer.session"
+        serverSession cfg (restrictHandleSocket h (fst con)) conkey h
     acceptLoop cfg sock (n + 1)
 
 acceptException :: ServerConfig -> Int -> Socket -> IOError -> IO ()