7 files changed, 129 insertions, 16 deletions

diff --git a/src/Control/Concurrent/Async/Lifted/Instrument.hs b/src/Control/Concurrent/Async/Lifted/Instrument.hs
new file mode 100644
index 00000000..eab0fadc
--- /dev/null
+++ b/src/Control/Concurrent/Async/Lifted/Instrument.hs
@@ -0,0 +1,5 @@
+module Control.Concurrent.Async.Lifted.Instrument
+    ( module Control.Concurrent.Async.Lifted
+    ) where
+
+import Control.Concurrent.Async.Lifted
diff --git a/src/Control/Concurrent/Lifted/Instrument.hs b/src/Control/Concurrent/Lifted/Instrument.hs
new file mode 100644
index 00000000..9ec5deef
--- /dev/null
+++ b/src/Control/Concurrent/Lifted/Instrument.hs
@@ -0,0 +1,78 @@
+{-# LANGUAGE FlexibleContexts #-}
+module Control.Concurrent.Lifted.Instrument
+    ( module Control.Concurrent.Lifted
+    , forkIO
+    , fork
+    , labelThread
+    , threadsInformation
+    , PerThread(..)
+    ) where
+
+import qualified Control.Concurrent.Lifted as Raw
+import Control.Concurrent.Lifted hiding (fork)
+import Control.Monad.Trans.Control
+import System.IO.Unsafe
+import System.Mem.Weak
+import qualified Data.Map as Map
+import qualified Data.IntMap as IntMap
+import Control.Exception.Lifted
+import Control.Monad.Base
+-- import Control.Monad.IO.Class
+import qualified GHC.Conc as GHC
+import Data.Time.Clock
+
+data PerThread = PerThread
+ { -- wkid :: Weak ThreadId
+   lbl :: String
+ , startTime :: UTCTime
+ }
+ deriving (Eq,Ord,Show) -- ,Data,Generic)
+
+data GlobalState = GlobalState
+ { threads :: Map.Map ThreadId PerThread -- IntMap.IntMap PerThread
+ -- , uniqSource :: Int
+ }
+
+globals :: MVar GlobalState
+globals = unsafePerformIO $ newMVar $ GlobalState
+    { threads = Map.empty
+    -- , uniqSource = 0
+    }
+{-# NOINLINE globals #-}
+
+
+forkIO :: IO () -> IO ThreadId
+forkIO = fork
+
+fork :: MonadBaseControl IO m => m () -> m ThreadId
+fork action = do
+    t <- Raw.fork $ do
+        -- wkid <- myThreadId >>= liftBase . mkWeakThreadId
+        -- tid <- newUniq
+        tid <- myThreadId
+        tm <- liftBase getCurrentTime
+        bracket_ (modifyThreads $ Map.insert tid (PerThread "" tm))
+                 (modifyThreads $ Map.delete tid)
+                 action
+    return t
+
+labelThread :: ThreadId -> String -> IO ()
+labelThread tid s = do
+    GHC.labelThread tid s
+    putStrLn $ "labelThread "++s++" "++show tid
+    modifyThreads $ Map.adjust (\pt -> pt { lbl = s }) tid
+
+threadsInformation :: IO [PerThread]
+threadsInformation = do
+    m <- threads <$> readMVar globals
+    return $ Map.elems m
+
+-- newUniq :: MonadBaseControl IO m => m Int
+-- newUniq = do
+--     modifyMVar globals (\st -> return (st { uniqSource = succ (uniqSource st) }, uniqSource st))
+ 
+modifyThreads :: MonadBase IO m => (Map.Map ThreadId PerThread -> Map.Map ThreadId PerThread) -> m ()
+modifyThreads f = do
+    g <- takeMVar globals
+    let f' st = st { threads = f (threads st) }
+    putMVar globals (f' g)
diff --git a/src/Network/BitTorrent/DHT/Query.hs b/src/Network/BitTorrent/DHT/Query.hs
index 73b3d492..533068c6 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
@@ -225,8 +231,8 @@ refreshNodes nid = do
 
 -- | This operation do not block but acquire exclusive access to
 --   routing table.
-insertNode :: forall ip. Address ip => NodeInfo ip -> Maybe ReflectedIP -> DHT ip ThreadId
-insertNode info witnessed_ip0 = fork $ do
+insertNode :: forall ip. Address ip => NodeInfo ip -> Maybe ReflectedIP -> DHT ip ()
+insertNode info witnessed_ip0 = do
   var <- asks routingInfo
   tm <- getTimestamp
   let showTable = do
@@ -286,8 +292,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/Routing.hs b/src/Network/BitTorrent/DHT/Routing.hs
index 8a6849a1..d64e415e 100644
--- a/src/Network/BitTorrent/DHT/Routing.hs
+++ b/src/Network/BitTorrent/DHT/Routing.hs
@@ -529,11 +529,11 @@ splitTip nid n i bucket
 -- k nodes in them.  Which subtrees I mean is illustrated in Fig 1. of Kademlia
 -- paper.  The rule requiring additional splits is in section 2.4.
 modifyBucket
-  :: forall f ip xs. (Alternative f, Eq ip) =>
-     NodeId -> (Bucket ip -> f (xs, Bucket ip)) -> Table ip -> f (xs,Table ip)
+  :: forall ip xs. (Eq ip) =>
+     NodeId -> (Bucket ip -> Maybe (xs, Bucket ip)) -> Table ip -> Maybe (xs,Table ip)
 modifyBucket nodeId f = go (0 :: BitIx)
   where
-    go :: BitIx -> Table ip -> f (xs, Table ip)
+    go :: BitIx -> Table ip -> Maybe (xs, Table ip)
     go i (Zero table  bucket)
       | testIdBit nodeId i  = second (Zero  table) <$> f bucket
       |     otherwise       = second (`Zero` bucket) <$> go (succ i) table
@@ -562,8 +562,8 @@ data CheckPing ip = CheckPing [NodeInfo ip]
 
 
 -- | Atomic 'Table' update
-insert :: (Alternative m, Eq ip) => Timestamp -> Event ip -> Table ip -> m ([CheckPing ip], Table ip)
-insert tm event = modifyBucket (eventId event) (insertBucket tm event)
+insert :: (Eq ip, Applicative m) => Timestamp -> Event ip -> Table ip -> m ([CheckPing ip], Table ip)
+insert tm event tbl = pure $ fromMaybe ([],tbl) $ modifyBucket (eventId event) (insertBucket tm event) tbl
 
 
 {-----------------------------------------------------------------------
diff --git a/src/Network/BitTorrent/DHT/Session.hs b/src/Network/BitTorrent/DHT/Session.hs
index c08021c7..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
@@ -395,8 +399,9 @@ routableAddress = do
 myNodeIdAccordingTo :: NodeAddr ip -> DHT ip NodeId
 myNodeIdAccordingTo _ = do
   info <- asks routingInfo >>= liftIO . atomically . readTVar
-  fallback <- asks tentativeNodeId
-  return $ maybe fallback myNodeId info
+  maybe (asks tentativeNodeId)
+        (return . myNodeId)
+        info
 
 -- | Get current routing table. Normally you don't need to use this
 -- function, but it can be usefull for debugging and profiling purposes.
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 ()