Implemented bucket refresh for Mainline.

1 files changed, 32 insertions, 6 deletions

diff --git a/Kademlia.hs b/Kademlia.hs
index 9316f135..bcddfab1 100644
--- a/Kademlia.hs
+++ b/Kademlia.hs
@@ -18,7 +18,6 @@ import GHC.Conc (labelThread)
 import Control.Concurrent.Lifted
 #endif
 import Data.Bits
-
 import Text.PrettyPrint as PP hiding ((<>), ($$))
 import Text.PrettyPrint.HughesPJClass hiding ((<>),($$))
 import Data.IP
@@ -28,6 +27,11 @@ import Data.Monoid
 import Data.Time.Clock.POSIX (POSIXTime)
 import Data.Wrapper.PSQInt ( pattern (:->) )
 import qualified Data.Wrapper.PSQInt as Int
+import Network.BitTorrent.DHT.Search
+import Network.DatagramServer.Types (genBucketSample)
+import Network.Address (bucketRange)
+import Data.Serialize (Serialize)
+import Data.Hashable
 
 -- | The status of a given node with respect to a given routint table.
 data RoutingStatus
@@ -206,15 +210,16 @@ insertNode (Kademlia reporter space io) node = do
 -- It might also be better to pass the timestamp of the transition here and
 -- keep the refresh queue in better sync with the routing table by updating it
 -- within the STM monad.
-touchBucket :: POSIXTime -> (ni -> Int) -> TVar (Int.PSQ POSIXTime) -> RoutingTransition ni -> STM (IO ())
-touchBucket interval bktnum psq tr
+touchBucket :: KademliaSpace nid ni -> POSIXTime -> TVar (BucketList ni) -> TVar (Int.PSQ POSIXTime) -> RoutingTransition ni -> STM (IO ())
+touchBucket space interval bkts psq tr
     | (transitionedTo tr == Applicant)
                 = return $ return ()
     | otherwise = return $ do
          now <- getPOSIXTime
-         atomically $ modifyTVar' psq
-                    $ Int.insert (bktnum $ transitioningNode tr)
-                                 (now + interval)
+         atomically $ do
+            let nid = kademliaLocation space (transitioningNode tr)
+            num <- R.bucketNumber space nid <$> readTVar bkts
+            modifyTVar' psq $ Int.insert num (now + interval)
 
 -- | > pollForRefresh interval queue refresh
 --
@@ -250,3 +255,24 @@ pollForRefresh interval psq refresh = do
                              return ()
             seconds    -> threadDelay ( seconds * 1000000 )
         again
+
+refreshBucket :: forall nid ni addr.
+                 ( FiniteBits nid, Serialize nid, Ord nid, Ord ni, Hashable nid, Hashable ni, Ord addr ) =>
+                 Search nid addr ni ni -> TVar (BucketList ni) -> nid -> Int -> IO ()
+refreshBucket sch var nid n = do
+    tbl <- atomically (readTVar var)
+    sample <- genBucketSample nid (bucketRange (n) (n + 1 < bktCount tbl))
+    resultCounter <- atomically $ newTVar 0
+    let fullcount = R.defaultBucketSize
+    let checkBucketFull :: ni -> STM Bool
+        checkBucketFull found_node = do
+            tbl <- readTVar var
+            let counts = R.shape tbl
+            when (n == R.bucketNumber (searchSpace sch) (kademliaLocation (searchSpace sch) found_node) tbl)
+                 $ modifyTVar resultCounter (+ 1)
+            resultCount <- readTVar resultCounter
+            case drop (n - 1) counts of
+                (cnt:_) | cnt < fullcount       -> return True
+                _ | resultCount + 3 < fullcount -> return True -- +3 because maybe duplicates.
+                _                               -> return False
+    search sch tbl sample checkBucketFull