Implemented bucket refresh for Mainline.

1 files changed, 38 insertions, 13 deletions

diff --git a/Mainline.hs b/Mainline.hs
index 8853e920..1e30b718 100644
--- a/Mainline.hs
+++ b/Mainline.hs
@@ -38,13 +38,16 @@ import Data.Monoid
 import Data.Ord
 import qualified Data.Serialize           as S
 import Data.Set                           (Set)
+import Data.Time.Clock.POSIX              (POSIXTime)
 import Data.Torrent
 import Data.Typeable
 import Data.Word
+import qualified Data.Wrapper.PSQInt      as Int
 import Kademlia
 import Network.Address                    (Address, fromSockAddr, setPort,
                                            sockAddrPort, testIdBit, toSockAddr)
 import Network.BitTorrent.DHT.ContactInfo as Peers
+import Network.BitTorrent.DHT.Search      (Search (..))
 import Network.BitTorrent.DHT.Token       as Token
 import qualified Network.DHT.Routing      as R
          ;import Network.DHT.Routing      (Info, Timestamp, getTimestamp)
@@ -54,6 +57,10 @@ import Network.Socket
 newtype NodeId = NodeId ByteString
  deriving (Eq,Ord,Show,ByteArrayAccess, BEncode, Bits, Hashable)
 
+instance S.Serialize NodeId where
+    get = NodeId <$> S.getBytes 20
+    put (NodeId bs) = S.putByteString bs
+
 instance FiniteBits NodeId where
     finiteBitSize _ = 160
 
@@ -62,6 +69,7 @@ data NodeInfo = NodeInfo
   , nodeIP   :: IP
   , nodePort :: PortNumber
   }
+ deriving (Eq,Ord)
 
 -- The Hashable instance depends only on the IP address and port number.  It is
 -- used to compute the announce token.
@@ -271,10 +279,12 @@ type RoutingInfo = Info NodeInfo NodeId
 
 data Routing = Routing
     { tentativeId :: NodeInfo
-    , routing4 :: !( TVar (R.BucketList NodeInfo) )
-    , committee4 :: TriadCommittee NodeId SockAddr
-    , routing6 :: !( TVar (R.BucketList NodeInfo) )
-    , committee6 :: TriadCommittee NodeId SockAddr
+    , 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
     }
 
 type MainlineClient = Client String Method TransactionId NodeInfo (Message BValue)
@@ -294,7 +304,12 @@ newClient addr = do
         tbl6 <- newTVar $ R.nullTable (comparing nodeId) (\s -> hashWithSalt s . nodeId) tenative_info nobkts
         committee4 <- newTriadCommittee (const $ return ()) -- TODO: update tbl4
         committee6 <- newTriadCommittee (const $ return ()) -- TODO: update tbl6
-        return $ Routing tenative_info tbl4 committee4 tbl6 committee6
+        sched4 <- newTVar Int.empty
+        sched6 <- newTVar Int.empty
+        return $ Routing tenative_info sched4 tbl4 committee4 sched6 tbl6 committee6
+    -- TODO: Provide some means of shutting down these two auxillary threads:
+    refresh_thread4 <- pollForRefresh (15*60) (sched4 routing) (refreshBucket nodeSearch (routing4 routing) (nodeId tenative_info))
+    refresh_thread6 <- pollForRefresh (15*60) (sched6 routing) (refreshBucket nodeSearch (routing6 routing) (nodeId tenative_info))
     swarms <- newSwarmsDatabase
     map_var <- atomically $ newTVar (0, mempty)
     let net = onInbound (updateRouting outgoingClient routing)
@@ -346,12 +361,16 @@ defaultHandler meth = MethodHandler decodePayload errorPayload returnError
     returnError :: NodeInfo -> BValue -> IO Error
     returnError _ _ = return $ Error MethodUnknown ("Unknown method " <> meth)
 
-mainlineKademlia :: MainlineClient -> TriadCommittee NodeId SockAddr -> TVar (R.BucketList NodeInfo) ->  Kademlia NodeId NodeInfo
-mainlineKademlia client committee var
+mainlineKademlia :: MainlineClient -> TriadCommittee NodeId SockAddr -> TVar (R.BucketList NodeInfo) -> TVar (Int.PSQ POSIXTime) -> Kademlia NodeId NodeInfo
+mainlineKademlia client committee var sched
     = Kademlia quietInsertions
                mainlineSpace
                (vanillaIO var $ ping client)
-                   { tblTransition = transitionCommittee committee }
+                   { tblTransition = \tr -> do
+                        io1 <- transitionCommittee committee tr
+                        io2 <- touchBucket mainlineSpace (15*60) var sched tr
+                        return $ io1 >> io2
+                   }
 
 
 mainlineSpace :: R.KademliaSpace NodeId NodeInfo
@@ -365,19 +384,20 @@ transitionCommittee :: TriadCommittee NodeId SockAddr -> RoutingTransition NodeI
 transitionCommittee committee (RoutingTransition ni Stranger) = do
     delVote committee (nodeId ni)
     return $ return ()
+transitionCommittee committee _ = return $ return ()
 
 updateRouting :: MainlineClient -> Routing -> NodeInfo -> Message BValue -> IO ()
 updateRouting client routing naddr msg = do
     case prefer4or6 naddr Nothing of
-        Want_IP4 -> go (routing4 routing) (committee4 routing)
-        Want_IP6 -> go (routing6 routing) (committee6 routing)
+        Want_IP4 -> go (routing4 routing) (committee4 routing) (sched4 routing)
+        Want_IP6 -> go (routing6 routing) (committee6 routing) (sched6 routing)
  where
-    go tbl committee = do
+    go tbl committee sched = do
         case msg of
             R { rspReflectedIP = Just sockaddr }
                 -> atomically $ addVote committee (nodeId naddr) sockaddr
             _   -> return ()
-        insertNode (mainlineKademlia client committee tbl) naddr
+        insertNode (mainlineKademlia client committee tbl sched) naddr
 
 data Ping = Ping deriving Show
 
@@ -478,7 +498,7 @@ findNodeH routing addr (FindNode node iptyp) = do
     ks6 <- bool (return []) (go $ routing6 routing) (preferred /= Want_IP4)
     return $ NodeFound ks ks6
  where
-    go var = R.kclosest nodeId k node <$> atomically (readTVar var)
+    go var = R.kclosest mainlineSpace k node <$> atomically (readTVar var)
     k = R.defaultK
 
 
@@ -754,3 +774,8 @@ delVote triad voter = do
         writeTVar (triadSlot slot triad) Nothing
         triadCountVotes triad
 
+nodeSearch = Search
+    { searchSpace = mainlineSpace
+    , searchNodeAddress = nodeIP &&& nodePort
+    , searchQuery = error "searchQuery"
+    }