Moved BucketRefresher construction responsibility for greater

encapsulation.

1 files changed, 23 insertions, 25 deletions

diff --git a/src/Network/BitTorrent/MainlineDHT.hs b/src/Network/BitTorrent/MainlineDHT.hs
index 268cacfb..3e7a0eda 100644
--- a/src/Network/BitTorrent/MainlineDHT.hs
+++ b/src/Network/BitTorrent/MainlineDHT.hs
@@ -530,15 +530,19 @@ traced (TableMethods ins del lkup)
 
 type MainlineClient = Client String Method TransactionId NodeInfo (Message BValue)
 
+-- | Like 'nodeInfo' but falls back to 'iNADDR_ANY' for nodeIP' and 'nodePort'.
+mkNodeInfo :: NodeId -> SockAddr -> NodeInfo
+mkNodeInfo nid addr = NodeInfo
+    { nodeId   = nid
+    , nodeIP   = fromMaybe (toEnum 0) $ fromSockAddr addr
+    , nodePort = fromMaybe 0 $ sockAddrPort addr
+    }
+
 newClient :: SwarmsDatabase -> SockAddr -> IO (MainlineClient, Routing)
 newClient swarms addr = do
     udp <- udpTransport addr
     nid <- NodeId <$> getRandomBytes 20
-    let tentative_info = NodeInfo
-                { nodeId   = nid
-                , nodeIP   = fromMaybe (toEnum 0) $ fromSockAddr addr
-                , nodePort = fromMaybe 0 $ sockAddrPort addr
-                }
+    let tentative_info = mkNodeInfo nid addr
     tentative_info6 <-
         maybe tentative_info
               (\ip6 -> tentative_info { nodeId = fromMaybe (nodeId tentative_info)
@@ -549,37 +553,31 @@ newClient swarms addr = do
     addr4 <- atomically $ newTChan
     addr6 <- atomically $ newTChan
     mkrouting <- atomically $ do
-        let nobkts = R.defaultBucketCount :: Int
-        tbl4 <- newTVar $ R.nullTable (comparing nodeId) (\s -> hashWithSalt s . nodeId) tentative_info nobkts
-        tbl6 <- newTVar $ R.nullTable (comparing nodeId) (\s -> hashWithSalt s . nodeId) tentative_info6 nobkts
-        let updateIPVote tblvar addrvar a = do
+        -- We defer initializing the refreshSearch and refreshPing until we
+        -- have a client to send queries with.
+        let nullPing = const $ return False
+            nullSearch = mainlineSearch $ \_ _ -> return Nothing
+        refresher4 <- newBucketRefresher mainlineSpace tentative_info nullSearch nullPing
+        refresher6 <- newBucketRefresher mainlineSpace tentative_info6 nullSearch nullPing
+        let tbl4 = refreshBuckets refresher4
+            tbl6 = refreshBuckets refresher6
+            updateIPVote tblvar addrvar a = do
                 bkts <- readTVar tblvar
                 case bep42 a (nodeId $ R.thisNode bkts) of
                     Just nid -> do
                         let tbl = R.nullTable (comparing nodeId)
                                               (\s -> hashWithSalt s . nodeId)
-                                              (NodeInfo nid
-                                                        (fromMaybe (toEnum 0) $ fromSockAddr a)
-                                                        (fromMaybe 0 $ sockAddrPort a))
-                                              nobkts
+                                              (mkNodeInfo nid a)
+                                              (R.defaultBucketCount)
                         writeTVar tblvar tbl
                         writeTChan addrvar (a,map fst $ concat $ R.toList bkts)
                     Nothing -> return ()
         committee4 <- newTriadCommittee $ updateIPVote tbl4 addr4
         committee6 <- newTriadCommittee $ updateIPVote tbl6 addr6
-        sched4 <- newTVar Int.empty
-        sched6 <- newTVar Int.empty
         return $ \client ->
-            let refresher sched bkts = BucketRefresher
-                    { refreshInterval = 15 * 60
-                    , refreshQueue    = sched
-                    , refreshSearch   = nodeSearch client
-                    , refreshBuckets  = bkts
-                    , refreshPing     = ping client
-                    }
-                refresher4 = refresher sched4 tbl4
-                refresher6 = refresher sched6 tbl6
-            in Routing tentative_info committee4 committee6 refresher4 refresher6
+            -- Now we have a client, so tell the BucketRefresher how to search and ping.
+            let updIO r = updateRefresherIO (nodeSearch client) (ping client) r
+            in Routing tentative_info committee4 committee6 (updIO refresher4) (updIO refresher6)
     map_var <- atomically $ newTVar (0, mempty)
 
     let routing = mkrouting outgoingClient