3 files changed, 18 insertions, 18 deletions

diff --git a/dht/Presence/DNSCache.hs b/dht/Presence/DNSCache.hs
index e28655c5..14581fee 100644
--- a/dht/Presence/DNSCache.hs
+++ b/dht/Presence/DNSCache.hs
@@ -22,12 +22,7 @@ module DNSCache
     , withPort
     ) where
 
-#ifdef THREAD_DEBUG
-import Control.Concurrent.Lifted.Instrument
-#else
-import Control.Concurrent.Lifted
-import GHC.Conc                  (labelThread)
-#endif
+import Control.Concurrent.ThreadUtil
 import Control.Arrow
 import Control.Concurrent.STM
 import Data.Text       ( Text )
@@ -108,10 +103,10 @@ make6mapped4 :: SockAddr -> SockAddr
 make6mapped4 addr@(SockAddrInet6 {}) = addr
 make6mapped4 addr@(SockAddrInet port a) = SockAddrInet6 port 0 (0,0,0xFFFF,fromBE32 a) 0
 
-tryForkOS :: IO () -> IO ThreadId
-tryForkOS action = catchIOError (forkOS action) $ \e -> do
+tryForkOS :: String -> IO () -> IO ThreadId
+tryForkOS lbl action = catchIOError (forkOSLabeled lbl action) $ \e -> do
     dput XMisc $ "DNSCache: Link with -threaded to avoid excessively long time-out."
-    forkIO action
+    forkLabeled lbl action
 
 
 -- Attempt to resolve the given domain name.  Returns an empty list if the
@@ -126,8 +121,7 @@ rawForwardResolve ::
 rawForwardResolve dns onFail timeout addrtext = do
     r <- atomically newEmptyTMVar
     mvar <- interruptibleDelay
-    rt <- tryForkOS $ do
-        myThreadId >>= flip labelThread ("resolve."++show addrtext)
+    rt <- tryForkOS ("resolve."++show addrtext) $ do
         resolver r mvar
     startDelay mvar timeout
     did <- atomically $ tryPutTMVar r []
diff --git a/dht/src/Control/Concurrent/ThreadUtil.hs b/dht/src/Control/Concurrent/ThreadUtil.hs
index 2888e899..a258d933 100644
--- a/dht/src/Control/Concurrent/ThreadUtil.hs
+++ b/dht/src/Control/Concurrent/ThreadUtil.hs
@@ -21,4 +21,11 @@ forkLabeled lbl action = do
     labelThread t lbl
     return t
 {-# INLINE forkLabeled #-}
+
+forkOSLabeled :: String -> IO () -> IO ThreadId
+forkOSLabeled lbl action = do
+    t <- forkOS action
+    labelThread t lbl
+    return t
+{-# INLINE forkOSLabeled #-}
 #endif
diff --git a/dht/src/Network/QueryResponse/TCP.hs b/dht/src/Network/QueryResponse/TCP.hs
index de6e1eb7..e4831fb2 100644
--- a/dht/src/Network/QueryResponse/TCP.hs
+++ b/dht/src/Network/QueryResponse/TCP.hs
@@ -19,6 +19,7 @@ import Data.Maybe
 import Data.Ord
 import Data.Time.Clock.POSIX
 import Data.Word
+import Data.String (IsString(..))
 import Network.BSD
 import Network.Socket
 import System.Timeout
@@ -68,6 +69,7 @@ killSession :: TCPSession st -> IO ()
 killSession PendingTCPSession       = return ()
 killSession TCPSession{tcpThread=t} = killThread t
 
+showStat :: IsString p => TCPSession st -> p
 showStat r = case r of PendingTCPSession -> "pending."
                        TCPSession {}     -> "established."
 
@@ -105,7 +107,8 @@ acquireConnection mvar tcpcache stream addr bDoCon = do
                 st <- streamHello stream addr h
                 dput XTCP $ "TCP Connected! " ++ show (streamAddr stream addr)
                 signal <- newTVarIO False
-                rthread <- forkIO $ do
+                let showAddr a = show (streamAddr stream a)
+                rthread <- forkLabeled ("tcp:"++showAddr addr) $ do
                   atomically (readTVar signal >>= check)
                   fix $ \loop -> do
                     x <- streamDecode st
@@ -127,8 +130,6 @@ acquireConnection mvar tcpcache stream addr bDoCon = do
                             forM_ (zip [1..] $ MM.toList c) $ \(i,MM.Binding (TCPAddress addr) r (Down tm)) -> do
                                 dput XTCP $ unwords [show i ++ ".", "Still connected:",  show addr, show (now - tm), showStat r]
                             hClose h
-                let showAddr a = show (streamAddr stream a)
-                labelThread rthread ("tcp:"++showAddr addr)
                 let v = TCPSession
                         { tcpHandle = h
                         , tcpState  = st
@@ -142,8 +143,7 @@ acquireConnection mvar tcpcache stream addr bDoCon = do
                     writeTVar (lru tcpcache) c'
                     writeTVar signal True
                     return rs
-                forM_ retires $ \(MM.Binding (TCPAddress k) r _) -> void $ forkIO $ do
-                    myThreadId >>= flip labelThread ("tcp-close:"++show k)
+                forM_ retires $ \(MM.Binding (TCPAddress k) r _) -> void $ forkLabeled ("tcp-close:"++show k) $ do
                     dput XTCP $ "TCP dropped: " ++ show k
                     killSession r
                     case r of TCPSession {tcpState=st,tcpHandle=h} -> do
@@ -184,10 +184,9 @@ tcpTransport maxcon stream = do
     return $ (,) tcpcache Transport
         { awaitMessage   = \f -> takeMVar msgvar >>= \x -> f x `catchIOError` (\e -> dput XTCP ("TCP transport stopped. " ++ show e) >> f Nothing)
         , sendMessage    = \addr (bDoCon,y) -> do
-                                t <- forkIO $ do
+                                void . forkLabeled "tcp-send" $ do
                                     msock <- acquireConnection msgvar tcpcache stream addr bDoCon
                                     mapM_ ($ y) msock
                                   `catchIOError` \e -> dput XTCP $ "TCP-send: " ++ show e
-                                labelThread t "tcp-send"
         , closeTransport = closeAll tcpcache stream >> putMVar msgvar Nothing
         }