3 files changed, 24 insertions, 11 deletions

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.