1 files changed, 42 insertions, 19 deletions

diff --git a/src/Network/Tox/DHT/Handlers.hs b/src/Network/Tox/DHT/Handlers.hs
index 494e319b..d3d36525 100644
--- a/src/Network/Tox/DHT/Handlers.hs
+++ b/src/Network/Tox/DHT/Handlers.hs
@@ -1,7 +1,8 @@
+{-# LANGUAGE CPP                        #-}
 {-# LANGUAGE GeneralizedNewtypeDeriving #-}
+{-# LANGUAGE NamedFieldPuns             #-}
 {-# LANGUAGE PatternSynonyms            #-}
 {-# LANGUAGE TupleSections              #-}
-{-# LANGUAGE CPP                        #-}
 module Network.Tox.DHT.Handlers where
 
 import Network.Tox.DHT.Transport as DHTTransport
@@ -11,7 +12,7 @@ import Crypto.Tox
 import Network.Kademlia.Search
 import qualified Data.Wrapper.PSQInt as Int
 import Network.Kademlia
-import Network.Kademlia.Bootstrap    (touchBucket)
+import Network.Kademlia.Bootstrap
 import Network.Address               (WantIP (..), ipFamily, testIdBit,fromSockAddr, sockAddrPort)
 import qualified Network.Kademlia.Routing as R
 import Control.TriadCommittee
@@ -111,18 +112,28 @@ classify client msg = fromMaybe (IsUnknown "unknown")
 
 data Routing = Routing
     { tentativeId :: NodeInfo
-    , sched4      :: !( TVar (Int.PSQ POSIXTime) )
-    , routing4    :: !( TVar (R.BucketList NodeInfo) )
     , committee4  :: TriadCommittee NodeId SockAddr
-    , sched6      :: !( TVar (Int.PSQ POSIXTime) )
-    , routing6    :: !( TVar (R.BucketList NodeInfo) )
     , committee6  :: TriadCommittee NodeId SockAddr
+    , refresher4  :: BucketRefresher NodeId NodeInfo
+    , refresher6  :: BucketRefresher NodeId NodeInfo
     }
 
+sched4 :: Routing -> TVar (Int.PSQ POSIXTime)
+sched4 Routing { refresher4 = BucketRefresher { refreshQueue } } = refreshQueue
+
+sched6 :: Routing -> TVar (Int.PSQ POSIXTime)
+sched6 Routing { refresher6 = BucketRefresher { refreshQueue } } = refreshQueue
+
+routing4 :: Routing -> TVar (R.BucketList NodeInfo)
+routing4 Routing { refresher4 = BucketRefresher { refreshBuckets } } = refreshBuckets
+
+routing6 :: Routing -> TVar (R.BucketList NodeInfo)
+routing6 Routing { refresher6 = BucketRefresher { refreshBuckets } } = refreshBuckets
+
 newRouting :: SockAddr -> TransportCrypto
                        -> (TVar (R.BucketList NodeInfo) -> SockAddr -> STM ()) -- ^ invoked on IPv4 change
                        -> (TVar (R.BucketList NodeInfo) -> SockAddr -> STM ()) -- ^ invoked on IPv6 change
-                       -> IO Routing
+                       -> IO (Client -> Routing)
 newRouting addr crypto update4 update6 = do
     let tentative_ip4 = fromMaybe (IPv4 $ toEnum 0) (IPv4 <$> fromSockAddr addr)
         tentative_ip6 = fromMaybe (IPv6 $ toEnum 0) (IPv6 <$> fromSockAddr addr)
@@ -146,7 +157,17 @@ newRouting addr crypto update4 update6 = do
         committee6 <- newTriadCommittee (update6 tbl6) -- updateIPVote tbl6 addr6
         sched4 <- newTVar Int.empty
         sched6 <- newTVar Int.empty
-        return $ Routing tentative_info sched4 tbl4 committee4 sched6 tbl6 committee6
+        return $ \client ->
+            let refresher sched bkts = BucketRefresher
+                    { refreshInterval = 15 * 60
+                    , refreshQueue    = sched
+                    , refreshSearch   = nodeSearch client
+                    , refreshBuckets  = bkts
+                    , refreshPing     = ping client
+                    }
+                refresher4 = refresher sched4 tbl6
+                refresher6 = refresher sched6 tbl6
+            in Routing tentative_info committee4 committee6 refresher4 refresher6
 
 
 -- TODO: This should cover more cases
@@ -322,26 +343,28 @@ updateRouting client routing orouter naddr msg
     | PacketKind 0x21 <- msgType msg = return () -- ignore lan discovery
     | otherwise                      = do
     case prefer4or6 naddr Nothing of
-        Want_IP4 -> updateTable client naddr orouter (routing4 routing) (committee4 routing) (sched4 routing)
-        Want_IP6 -> updateTable client naddr orouter (routing6 routing) (committee6 routing) (sched6 routing)
+        Want_IP4 -> updateTable client naddr orouter (committee4 routing) (refresher4 routing)
+        Want_IP6 -> updateTable client naddr orouter (committee6 routing) (refresher4 routing)
         Want_Both -> error $ "BUG:unreachable at " ++ __FILE__ ++ ":" ++ show __LINE__
 
-updateTable :: Client -> NodeInfo -> OnionRouter -> TVar (R.BucketList NodeInfo) -> TriadCommittee NodeId SockAddr -> TVar (Int.PSQ POSIXTime) -> IO ()
-updateTable client naddr orouter tbl committee sched = do
-        self <- atomically $ R.thisNode <$> readTVar tbl
+updateTable :: Client -> NodeInfo -> OnionRouter -> TriadCommittee NodeId SockAddr -> BucketRefresher NodeId NodeInfo -> IO ()
+updateTable client naddr orouter committee refresher = do
+        self <- atomically $ R.thisNode <$> readTVar (refreshBuckets refresher)
         when (nodeIP self /= nodeIP naddr) $ do
             -- TODO: IP address vote?
-            insertNode (toxKademlia client committee orouter tbl sched) naddr
+            insertNode (toxKademlia client committee orouter refresher) naddr
 
-toxKademlia :: Client -> TriadCommittee NodeId SockAddr -> OnionRouter -> TVar (R.BucketList NodeInfo) -> TVar (Int.PSQ POSIXTime) -> Kademlia NodeId NodeInfo
-toxKademlia client committee orouter var sched
+toxKademlia :: Client -> TriadCommittee NodeId SockAddr -> OnionRouter
+                -> BucketRefresher NodeId NodeInfo
+                -> Kademlia NodeId NodeInfo
+toxKademlia client committee orouter refresher
     = Kademlia quietInsertions
                toxSpace
-               (vanillaIO var $ ping client)
+               (vanillaIO (refreshBuckets refresher) $ ping client)
                    { tblTransition = \tr -> do
                         io1 <- transitionCommittee committee tr
-                        io2 <- touchBucket toxSpace (15*60) var sched tr
-                        hookBucketList toxSpace var orouter tr
+                        io2 <- touchBucket refresher tr -- toxSpace (15*60) var sched tr
+                        hookBucketList toxSpace (refreshBuckets refresher) orouter tr
                         return $ do
                             io1 >> io2
                             {-