From a0cc40bfe8e3b20e645fd40008efe3b91482962d Mon Sep 17 00:00:00 2001
From: Anders Jensen <andersandjensen@gmail.com>
Date: Sun, 3 May 2020 02:36:16 +0200
Subject: [PATCH] Added support for replacement_nodes pruning in KBuckets for
 #79 (#80)

---
 kademlia/routing.py            |  6 +++++-
 kademlia/tests/test_routing.py | 13 +++++++++++++
 2 files changed, 18 insertions(+), 1 deletion(-)

diff --git a/kademlia/routing.py b/kademlia/routing.py
index 8bb5b8c..eaa356a 100644
--- a/kademlia/routing.py
+++ b/kademlia/routing.py
@@ -9,12 +9,13 @@ from kademlia.utils import shared_prefix, bytes_to_bit_string
 
 
 class KBucket:
-    def __init__(self, rangeLower, rangeUpper, ksize):
+    def __init__(self, rangeLower, rangeUpper, ksize, replacementNodeFactor=5):
         self.range = (rangeLower, rangeUpper)
         self.nodes = OrderedDict()
         self.replacement_nodes = OrderedDict()
         self.touch_last_updated()
         self.ksize = ksize
+        self.replacement_node_factor = replacementNodeFactor
 
     def touch_last_updated(self):
         self.last_updated = time.monotonic()
@@ -67,6 +68,9 @@ class KBucket:
             if node.id in self.replacement_nodes:
                 del self.replacement_nodes[node.id]
             self.replacement_nodes[node.id] = node
+            max_replacement_nodes = self.ksize * self.replacement_node_factor
+            while len(self.replacement_nodes) > max_replacement_nodes:
+                self.replacement_nodes.popitem(last=False)
             return False
         return True
 
diff --git a/kademlia/tests/test_routing.py b/kademlia/tests/test_routing.py
index 212a95f..bbdb062 100644
--- a/kademlia/tests/test_routing.py
+++ b/kademlia/tests/test_routing.py
@@ -65,6 +65,19 @@ class TestKBucket:
         assert bucket.has_in_range(mknode(intid=10)) is True
         assert bucket.has_in_range(mknode(intid=0)) is True
 
+    def test_replacement_factor(self, mknode):  # pylint: disable=no-self-use
+        k = 3
+        factor = 2
+        bucket = KBucket(0, 10, k, replacementNodeFactor=factor)
+        nodes = [mknode() for _ in range(10)]
+        for node in nodes:
+            bucket.add_node(node)
+
+        replacement_nodes = bucket.replacement_nodes
+        assert len(list(replacement_nodes.values())) == k * factor
+        assert list(replacement_nodes.values()) == nodes[k + 1:]
+        assert nodes[k] not in list(replacement_nodes.values())
+
 
 # pylint: disable=too-few-public-methods
 class TestRoutingTable: