1 files changed, 13 insertions, 3 deletions

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 ()