diff --git a/protocols/kad/src/kbucket/sub_bucket.rs b/protocols/kad/src/kbucket/sub_bucket.rs index 086ba1ef..aa74b9a1 100644 --- a/protocols/kad/src/kbucket/sub_bucket.rs +++ b/protocols/kad/src/kbucket/sub_bucket.rs @@ -14,8 +14,6 @@ * limitations under the License. */ -use crate::K_VALUE; -use arrayvec::ArrayVec; use std::time::Instant; enum ChangePosition { @@ -93,15 +91,17 @@ pub struct Position(pub usize); #[derive(Debug, Clone)] pub struct SubBucket { - pub nodes: ArrayVec<[Node; K_VALUE.get()]>, + pub nodes: Vec, pub first_connected_pos: Option, + pub capacity: usize, } impl SubBucket { - pub fn new() -> Self { + pub fn new(capacity: usize) -> Self { Self { - nodes: ArrayVec::new(), + nodes: Vec::with_capacity(capacity + 1), first_connected_pos: None, + capacity, } } @@ -122,7 +122,7 @@ impl SubBucket { } pub fn is_full(&self) -> bool { - self.nodes.is_full() + self.nodes.len() >= self.capacity } pub fn all_nodes_connected(&self) -> bool { @@ -181,7 +181,16 @@ impl SubBucket { self.change_connected_pos(ChangePosition::RemoveDisconnected) } } - self.nodes.pop_at(position.0) + if position.0 >= self.nodes.len() { + println!( + "WARNING: tried to evict node at {} while there's only {} nodes", + position.0, + self.nodes.len() + ); + None + } else { + Some(self.nodes.remove(position.0)) + } } pub fn pop_node(&mut self) -> Option { diff --git a/protocols/kad/src/kbucket/swamp.rs b/protocols/kad/src/kbucket/swamp.rs index 04fa40b9..5e08d353 100644 --- a/protocols/kad/src/kbucket/swamp.rs +++ b/protocols/kad/src/kbucket/swamp.rs @@ -17,6 +17,7 @@ use crate::kbucket::{ AppliedPending, InsertResult, KeyBytes, Node, NodeStatus, PendingNode, SubBucket, }; +use crate::K_VALUE; use std::time::{Duration, Instant}; #[derive(Debug, Clone)] @@ -33,7 +34,7 @@ where { pub fn new(pending_timeout: Duration) -> Self { Self { - bucket: SubBucket::new(), + bucket: SubBucket::new(K_VALUE.get()), pending: None, pending_timeout, } diff --git a/protocols/kad/src/kbucket/weighted.rs b/protocols/kad/src/kbucket/weighted.rs index 50489a89..4ee511f6 100644 --- a/protocols/kad/src/kbucket/weighted.rs +++ b/protocols/kad/src/kbucket/weighted.rs @@ -141,7 +141,9 @@ where } fn get_bucket_mut(&mut self, weight: u32) -> &mut SubBucket> { - self.map.entry(weight).or_insert(SubBucket::new()) + self.map + .entry(weight) + .or_insert(SubBucket::new(W_VALUE.get())) } fn append_connected_node(&mut self, node: WeightedNode) {