Add bucket count option

2 files changed, 15 insertions, 9 deletions

diff --git a/src/Network/BitTorrent/DHT/Routing.hs b/src/Network/BitTorrent/DHT/Routing.hs
index 4897a22e..f6926cb1 100644
--- a/src/Network/BitTorrent/DHT/Routing.hs
+++ b/src/Network/BitTorrent/DHT/Routing.hs
@@ -325,8 +325,10 @@ instance Pretty (Table ip) where
       ss = shape t
 
 -- | Empty table with specified /spine/ node id.
-nullTable :: Eq ip => NodeId -> Table ip
-nullTable nid = Tip nid defaultBucketCount PSQ.empty
+nullTable :: Eq ip => NodeId -> BucketCount -> Table ip
+nullTable nid n = Tip nid (bucketCount (pred n)) PSQ.empty
+  where
+    bucketCount x = max 0 (min 160 x)
 
 -- | Get the /spine/ node id.
 thisId :: Table ip -> NodeId
diff --git a/src/Network/BitTorrent/DHT/Session.hs b/src/Network/BitTorrent/DHT/Session.hs
index 55e66a7e..c47f7491 100644
--- a/src/Network/BitTorrent/DHT/Session.hs
+++ b/src/Network/BitTorrent/DHT/Session.hs
@@ -96,13 +96,16 @@ defaultAlpha = 3
 
 data Options = Options
   { -- | the degree of parallelism in 'find_node' queries.
-    optAlpha   :: {-# UNPACK #-} !Alpha
+    optAlpha       :: {-# UNPACK #-} !Alpha
 
     -- | number of nodes to return in 'find_node' responses.
-  , optK       :: {-# UNPACK #-} !K
+  , optK           :: {-# UNPACK #-} !K
+
+    -- | Number of buckets to maintain.
+  , optBucketCount :: {-# UNPACK #-} !BucketCount
 
     -- | RPC timeout.
-  , optTimeout ::                !NominalDiffTime
+  , optTimeout     ::                !NominalDiffTime
 
 --  , optReannounceInterval :: NominalDiffTime
 --  , optDataExpiredTimeout :: NominalDiffTime
@@ -110,9 +113,10 @@ data Options = Options
 
 instance Default Options where
   def = Options
-    { optAlpha   = defaultAlpha
-    , optK       = defaultK
-    , optTimeout = 5 -- seconds
+    { optAlpha       = defaultAlpha
+    , optK           = defaultK
+    , optBucketCount = defaultBucketCount
+    , optTimeout     = 5 -- seconds
     }
 
 microseconds :: NominalDiffTime -> Int
@@ -194,7 +198,7 @@ runDHT handlers opts naddr action = runResourceT $ do
     (_, m) <- allocate (newManager (toSockAddr naddr) handlers) closeManager
     myId   <- liftIO genNodeId
     node   <- liftIO $ Node opts m
-             <$> newMVar (nullTable myId)
+             <$> newMVar (nullTable myId (optBucketCount opts))
              <*> newTVarIO def
              <*> (newTVarIO =<< nullSessionTokens)
              <*> pure logger