libp2p_kad/kbucket/bucket.rs
1// Copyright 2019 Parity Technologies (UK) Ltd.
2//
3// Permission is hereby granted, free of charge, to any person obtaining a
4// copy of this software and associated documentation files (the "Software"),
5// to deal in the Software without restriction, including without limitation
6// the rights to use, copy, modify, merge, publish, distribute, sublicense,
7// and/or sell copies of the Software, and to permit persons to whom the
8// Software is furnished to do so, subject to the following conditions:
9//
10// The above copyright notice and this permission notice shall be included in
11// all copies or substantial portions of the Software.
12//
13// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
14// OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
15// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
16// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
17// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
18// FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
19// DEALINGS IN THE SOFTWARE.
20
21//! The internal API for a single `KBucket` in a `KBucketsTable`.
22//!
23//! > **Note**: Uniqueness of entries w.r.t. a `Key` in a `KBucket` is not
24//! > checked in this module. This is an invariant that must hold across all
25//! > buckets in a `KBucketsTable` and hence is enforced by the public API
26//! > of the `KBucketsTable` and in particular the public `Entry` API.
27
28use super::*;
29pub(crate) use crate::K_VALUE;
30/// A `PendingNode` is a `Node` that is pending insertion into a `KBucket`.
31#[derive(Debug, Clone)]
32pub(crate) struct PendingNode<TKey, TVal> {
33 /// The pending node to insert.
34 node: Node<TKey, TVal>,
35
36 /// The status of the pending node.
37 status: NodeStatus,
38
39 /// The instant at which the pending node is eligible for insertion into a bucket.
40 replace: Instant,
41}
42
43/// The status of a node in a bucket.
44///
45/// The status of a node in a bucket together with the time of the
46/// last status change determines the position of the node in a
47/// bucket.
48#[derive(PartialEq, Eq, Debug, Copy, Clone)]
49pub enum NodeStatus {
50 /// The node is considered connected.
51 Connected,
52 /// The node is considered disconnected.
53 Disconnected,
54}
55
56impl<TKey, TVal> PendingNode<TKey, TVal> {
57 pub(crate) fn status(&self) -> NodeStatus {
58 self.status
59 }
60
61 pub(crate) fn value_mut(&mut self) -> &mut TVal {
62 &mut self.node.value
63 }
64
65 pub(crate) fn is_ready(&self) -> bool {
66 Instant::now() >= self.replace
67 }
68
69 #[cfg(test)]
70 pub(crate) fn set_ready_at(&mut self, t: Instant) {
71 self.replace = t;
72 }
73
74 pub(crate) fn into_node(self) -> Node<TKey, TVal> {
75 self.node
76 }
77}
78
79/// A `Node` in a bucket, representing a peer participating
80/// in the Kademlia DHT together with an associated value (e.g. contact
81/// information).
82#[derive(Debug, Clone, PartialEq, Eq)]
83pub struct Node<TKey, TVal> {
84 /// The key of the node, identifying the peer.
85 pub key: TKey,
86 /// The associated value.
87 pub value: TVal,
88}
89
90/// The position of a node in a `KBucket`, i.e. a non-negative integer
91/// in the range `[0, bucket_size)`.
92#[derive(Copy, Clone, Debug, PartialEq, Eq, PartialOrd, Ord)]
93pub(crate) struct Position(usize);
94/// A `KBucket` is a list of up to `capacity` keys and associated values,
95/// ordered from least-recently connected to most-recently connected.
96#[derive(Debug, Clone)]
97pub(crate) struct KBucket<TKey, TVal> {
98 /// The nodes contained in the bucket.
99 nodes: Vec<Node<TKey, TVal>>,
100 /// The maximal number of nodes that a bucket can contain.
101 capacity: usize,
102
103 /// The position (index) in `nodes` that marks the first connected node.
104 ///
105 /// Since the entries in `nodes` are ordered from least-recently connected to
106 /// most-recently connected, all entries above this index are also considered
107 /// connected, i.e. the range `[0, first_connected_pos)` marks the sub-list of entries
108 /// that are considered disconnected and the range
109 /// `[first_connected_pos, capacity)` marks sub-list of entries that are
110 /// considered connected.
111 ///
112 /// `None` indicates that there are no connected entries in the bucket, i.e.
113 /// the bucket is either empty, or contains only entries for peers that are
114 /// considered disconnected.
115 first_connected_pos: Option<usize>,
116
117 /// A node that is pending to be inserted into a full bucket, should the
118 /// least-recently connected (and currently disconnected) node not be
119 /// marked as connected within `unresponsive_timeout`.
120 pending: Option<PendingNode<TKey, TVal>>,
121
122 /// The timeout window before a new pending node is eligible for insertion,
123 /// if the least-recently connected node is not updated as being connected
124 /// in the meantime.
125 pending_timeout: Duration,
126}
127
128/// The result of inserting an entry into a bucket.
129#[must_use]
130#[derive(Debug, Clone, PartialEq, Eq)]
131pub(crate) enum InsertResult<TKey> {
132 /// The entry has been successfully inserted.
133 Inserted,
134 /// The entry is pending insertion because the relevant bucket is currently full.
135 /// The entry is inserted after a timeout elapsed, if the status of the
136 /// least-recently connected (and currently disconnected) node in the bucket
137 /// is not updated before the timeout expires.
138 Pending {
139 /// The key of the least-recently connected entry that is currently considered
140 /// disconnected and whose corresponding peer should be checked for connectivity
141 /// in order to prevent it from being evicted. If connectivity to the peer is
142 /// re-established, the corresponding entry should be updated with
143 /// [`NodeStatus::Connected`].
144 disconnected: TKey,
145 },
146 /// The entry was not inserted because the relevant bucket is full.
147 Full,
148}
149
150/// The result of applying a pending node to a bucket, possibly
151/// replacing an existing node.
152#[derive(Debug, Clone, PartialEq, Eq)]
153pub(crate) struct AppliedPending<TKey, TVal> {
154 /// The key of the inserted pending node.
155 pub(crate) inserted: Node<TKey, TVal>,
156 /// The node that has been evicted from the bucket to make room for the
157 /// pending node, if any.
158 pub(crate) evicted: Option<Node<TKey, TVal>>,
159}
160
161impl<TKey, TVal> Default for KBucket<TKey, TVal> {
162 fn default() -> Self {
163 KBucket {
164 nodes: Vec::with_capacity(K_VALUE.get()),
165 capacity: K_VALUE.get(),
166 first_connected_pos: None,
167 pending: None,
168 pending_timeout: Duration::from_secs(60),
169 }
170 }
171}
172
173impl<TKey, TVal> KBucket<TKey, TVal>
174where
175 TKey: Clone + AsRef<KeyBytes>,
176 TVal: Clone,
177{
178 /// Creates a new `KBucket` with the given configuration.
179 pub(crate) fn new(config: KBucketConfig) -> Self {
180 KBucket {
181 nodes: Vec::with_capacity(config.bucket_size),
182 capacity: config.bucket_size,
183 first_connected_pos: None,
184 pending: None,
185 pending_timeout: config.pending_timeout,
186 }
187 }
188
189 /// Returns a reference to the pending node of the bucket, if there is any.
190 pub(crate) fn pending(&self) -> Option<&PendingNode<TKey, TVal>> {
191 self.pending.as_ref()
192 }
193
194 /// Returns a mutable reference to the pending node of the bucket, if there is any.
195 pub(crate) fn pending_mut(&mut self) -> Option<&mut PendingNode<TKey, TVal>> {
196 self.pending.as_mut()
197 }
198
199 /// Returns a reference to the pending node of the bucket, if there is any
200 /// with a matching key.
201 pub(crate) fn as_pending(&self, key: &TKey) -> Option<&PendingNode<TKey, TVal>> {
202 self.pending()
203 .filter(|p| p.node.key.as_ref() == key.as_ref())
204 }
205
206 /// Returns an iterator over the nodes in the bucket, together with their status.
207 pub(crate) fn iter(&self) -> impl Iterator<Item = (&Node<TKey, TVal>, NodeStatus)> {
208 self.nodes
209 .iter()
210 .enumerate()
211 .map(move |(p, n)| (n, self.status(Position(p))))
212 }
213
214 /// Inserts the pending node into the bucket, if its timeout has elapsed,
215 /// replacing the least-recently connected node.
216 ///
217 /// If a pending node has been inserted, its key is returned together with
218 /// the node that was replaced. `None` indicates that the nodes in the
219 /// bucket remained unchanged.
220 pub(crate) fn apply_pending(&mut self) -> Option<AppliedPending<TKey, TVal>> {
221 if let Some(pending) = self.pending.take() {
222 if pending.replace <= Instant::now() {
223 if self.nodes.len() >= self.capacity {
224 if self.status(Position(0)) == NodeStatus::Connected {
225 // The bucket is full with connected nodes. Drop the pending node.
226 return None;
227 }
228 debug_assert!(self.first_connected_pos.is_none_or(|p| p > 0)); // (*)
229 // The pending node will be inserted.
230 let inserted = pending.node.clone();
231 // A connected pending node goes at the end of the list for
232 // the connected peers, removing the least-recently connected.
233 if pending.status == NodeStatus::Connected {
234 let evicted = Some(self.nodes.remove(0));
235 self.first_connected_pos = self
236 .first_connected_pos
237 .map_or_else(|| Some(self.nodes.len()), |p| p.checked_sub(1));
238 self.nodes.push(pending.node);
239 return Some(AppliedPending { inserted, evicted });
240 }
241 // A disconnected pending node goes at the end of the list
242 // for the disconnected peers.
243 else if let Some(p) = self.first_connected_pos {
244 let insert_pos = p.checked_sub(1).expect("by (*)");
245 let evicted = Some(self.nodes.remove(0));
246 self.nodes.insert(insert_pos, pending.node);
247 return Some(AppliedPending { inserted, evicted });
248 } else {
249 // All nodes are disconnected. Insert the new node as the most
250 // recently disconnected, removing the least-recently disconnected.
251 let evicted = Some(self.nodes.remove(0));
252 self.nodes.push(pending.node);
253 return Some(AppliedPending { inserted, evicted });
254 }
255 } else {
256 // There is room in the bucket, so just insert the pending node.
257 let inserted = pending.node.clone();
258 match self.insert(pending.node, pending.status) {
259 InsertResult::Inserted => {
260 return Some(AppliedPending {
261 inserted,
262 evicted: None,
263 })
264 }
265 _ => unreachable!("Bucket is not full."),
266 }
267 }
268 } else {
269 self.pending = Some(pending);
270 }
271 }
272
273 None
274 }
275
276 /// Updates the status of the pending node, if any.
277 pub(crate) fn update_pending(&mut self, status: NodeStatus) {
278 if let Some(pending) = &mut self.pending {
279 pending.status = status
280 }
281 }
282
283 /// Removes the pending node from the bucket, if any.
284 pub(crate) fn remove_pending(&mut self) -> Option<PendingNode<TKey, TVal>> {
285 self.pending.take()
286 }
287
288 /// Updates the status of the node referred to by the given key, if it is
289 /// in the bucket.
290 pub(crate) fn update(&mut self, key: &TKey, status: NodeStatus) {
291 // Remove the node from its current position and then reinsert it
292 // with the desired status, which puts it at the end of either the
293 // prefix list of disconnected nodes or the suffix list of connected
294 // nodes (i.e. most-recently disconnected or most-recently connected,
295 // respectively).
296 if let Some((node, _status, pos)) = self.remove(key) {
297 // If the least-recently connected node re-establishes its
298 // connected status, drop the pending node.
299 if pos == Position(0) && status == NodeStatus::Connected {
300 self.pending = None
301 }
302 // Reinsert the node with the desired status.
303 match self.insert(node, status) {
304 InsertResult::Inserted => {}
305 _ => unreachable!("The node is removed before being (re)inserted."),
306 }
307 }
308 }
309
310 /// Inserts a new node into the bucket with the given status.
311 ///
312 /// The status of the node to insert determines the result as follows:
313 ///
314 /// * `NodeStatus::Connected`: If the bucket is full and either all nodes are connected or
315 /// there is already a pending node, insertion fails with `InsertResult::Full`. If the
316 /// bucket is full but at least one node is disconnected and there is no pending node, the
317 /// new node is inserted as pending, yielding `InsertResult::Pending`. Otherwise the bucket
318 /// has free slots and the new node is added to the end of the bucket as the most-recently
319 /// connected node.
320 ///
321 /// * `NodeStatus::Disconnected`: If the bucket is full, insertion fails with
322 /// `InsertResult::Full`. Otherwise the bucket has free slots and the new node is inserted
323 /// at the position preceding the first connected node, i.e. as the most-recently
324 /// disconnected node. If there are no connected nodes, the new node is added as the last
325 /// element of the bucket.
326 pub(crate) fn insert(
327 &mut self,
328 node: Node<TKey, TVal>,
329 status: NodeStatus,
330 ) -> InsertResult<TKey> {
331 match status {
332 NodeStatus::Connected => {
333 if self.nodes.len() >= self.capacity {
334 if self.first_connected_pos == Some(0) || self.pending.is_some() {
335 return InsertResult::Full;
336 } else {
337 self.pending = Some(PendingNode {
338 node,
339 status: NodeStatus::Connected,
340 replace: Instant::now() + self.pending_timeout,
341 });
342 return InsertResult::Pending {
343 disconnected: self.nodes[0].key.clone(),
344 };
345 }
346 }
347 let pos = self.nodes.len();
348 self.first_connected_pos = self.first_connected_pos.or(Some(pos));
349 self.nodes.push(node);
350 InsertResult::Inserted
351 }
352 NodeStatus::Disconnected => {
353 if self.nodes.len() >= self.capacity {
354 return InsertResult::Full;
355 }
356 if let Some(ref mut p) = self.first_connected_pos {
357 self.nodes.insert(*p, node);
358 *p += 1;
359 } else {
360 self.nodes.push(node);
361 }
362 InsertResult::Inserted
363 }
364 }
365 }
366
367 /// Removes the node with the given key from the bucket, if it exists.
368 pub(crate) fn remove(
369 &mut self,
370 key: &TKey,
371 ) -> Option<(Node<TKey, TVal>, NodeStatus, Position)> {
372 if let Some(pos) = self.position(key) {
373 // Remove the node from its current position.
374 let status = self.status(pos);
375 let node = self.nodes.remove(pos.0);
376 // Adjust `first_connected_pos` accordingly.
377 match status {
378 NodeStatus::Connected => {
379 if self.first_connected_pos.is_some_and(|p| p == pos.0)
380 && pos.0 == self.nodes.len()
381 {
382 // It was the last connected node.
383 self.first_connected_pos = None
384 }
385 }
386 NodeStatus::Disconnected => {
387 if let Some(ref mut p) = self.first_connected_pos {
388 *p -= 1;
389 }
390 }
391 }
392 Some((node, status, pos))
393 } else {
394 None
395 }
396 }
397
398 /// Returns the status of the node at the given position.
399 pub(crate) fn status(&self, pos: Position) -> NodeStatus {
400 if self.first_connected_pos.is_some_and(|i| pos.0 >= i) {
401 NodeStatus::Connected
402 } else {
403 NodeStatus::Disconnected
404 }
405 }
406
407 /// Gets the number of entries currently in the bucket.
408 pub(crate) fn num_entries(&self) -> usize {
409 self.nodes.len()
410 }
411
412 /// Gets the number of entries in the bucket that are considered connected.
413 #[cfg(test)]
414 pub(crate) fn num_connected(&self) -> usize {
415 self.first_connected_pos.map_or(0, |i| self.nodes.len() - i)
416 }
417
418 /// Gets the number of entries in the bucket that are considered disconnected.
419 #[cfg(test)]
420 pub(crate) fn num_disconnected(&self) -> usize {
421 self.nodes.len() - self.num_connected()
422 }
423
424 /// Gets the position of an node in the bucket.
425 pub(crate) fn position(&self, key: &TKey) -> Option<Position> {
426 self.nodes
427 .iter()
428 .position(|p| p.key.as_ref() == key.as_ref())
429 .map(Position)
430 }
431
432 /// Gets a mutable reference to the node identified by the given key.
433 ///
434 /// Returns `None` if the given key does not refer to a node in the
435 /// bucket.
436 pub(crate) fn get_mut(&mut self, key: &TKey) -> Option<&mut Node<TKey, TVal>> {
437 self.nodes
438 .iter_mut()
439 .find(move |p| p.key.as_ref() == key.as_ref())
440 }
441}
442
443#[cfg(test)]
444mod tests {
445 use libp2p_identity::PeerId;
446 use quickcheck::*;
447
448 use super::*;
449
450 impl Arbitrary for KBucket<Key<PeerId>, ()> {
451 fn arbitrary(g: &mut Gen) -> KBucket<Key<PeerId>, ()> {
452 let timeout = Duration::from_secs(g.gen_range(1..g.size()) as u64);
453 let mut config = KBucketConfig::default();
454 config.set_pending_timeout(timeout);
455 let mut bucket = KBucket::<Key<PeerId>, ()>::new(config);
456 let num_nodes = g.gen_range(1..bucket.capacity + 1);
457 for _ in 0..num_nodes {
458 let key = Key::from(PeerId::random());
459 let node = Node { key, value: () };
460 let status = NodeStatus::arbitrary(g);
461 match bucket.insert(node, status) {
462 InsertResult::Inserted => {}
463 _ => panic!(),
464 }
465 }
466 bucket
467 }
468 }
469
470 impl Arbitrary for NodeStatus {
471 fn arbitrary(g: &mut Gen) -> NodeStatus {
472 if bool::arbitrary(g) {
473 NodeStatus::Connected
474 } else {
475 NodeStatus::Disconnected
476 }
477 }
478 }
479
480 impl Arbitrary for Position {
481 fn arbitrary(g: &mut Gen) -> Position {
482 Position(g.gen_range(0..K_VALUE.get()))
483 }
484 }
485
486 // Fill a bucket with random nodes with the given status.
487 fn fill_bucket(bucket: &mut KBucket<Key<PeerId>, ()>, status: NodeStatus) {
488 let num_entries_start = bucket.num_entries();
489 for i in 0..bucket.capacity - num_entries_start {
490 let key = Key::from(PeerId::random());
491 let node = Node { key, value: () };
492 assert_eq!(InsertResult::Inserted, bucket.insert(node, status));
493 assert_eq!(bucket.num_entries(), num_entries_start + i + 1);
494 }
495 }
496
497 #[test]
498 fn ordering() {
499 fn prop(status: Vec<NodeStatus>) -> bool {
500 let mut bucket = KBucket::<Key<PeerId>, ()>::default();
501
502 // The expected lists of connected and disconnected nodes.
503 let mut connected = VecDeque::new();
504 let mut disconnected = VecDeque::new();
505
506 // Fill the bucket, thereby populating the expected lists in insertion order.
507 for status in status {
508 let key = Key::from(PeerId::random());
509 let node = Node { key, value: () };
510 let full = bucket.num_entries() == bucket.capacity;
511 if let InsertResult::Inserted = bucket.insert(node, status) {
512 let vec = match status {
513 NodeStatus::Connected => &mut connected,
514 NodeStatus::Disconnected => &mut disconnected,
515 };
516 if full {
517 vec.pop_front();
518 }
519 vec.push_back((status, key));
520 }
521 }
522
523 // Get all nodes from the bucket, together with their status.
524 let mut nodes = bucket.iter().map(|(n, s)| (s, n.key)).collect::<Vec<_>>();
525
526 // Split the list of nodes at the first connected node.
527 let first_connected_pos = nodes.iter().position(|(s, _)| *s == NodeStatus::Connected);
528 assert_eq!(bucket.first_connected_pos, first_connected_pos);
529 let tail = first_connected_pos.map_or(Vec::new(), |p| nodes.split_off(p));
530
531 // All nodes before the first connected node must be disconnected and
532 // in insertion order. Similarly, all remaining nodes must be connected
533 // and in insertion order.
534 disconnected == nodes && connected == tail
535 }
536
537 quickcheck(prop as fn(_) -> _);
538 }
539
540 #[test]
541 fn full_bucket() {
542 let mut bucket = KBucket::<Key<PeerId>, ()>::default();
543
544 // Fill the bucket with disconnected nodes.
545 fill_bucket(&mut bucket, NodeStatus::Disconnected);
546
547 // Trying to insert another disconnected node fails.
548 let key = Key::from(PeerId::random());
549 let node = Node { key, value: () };
550 match bucket.insert(node, NodeStatus::Disconnected) {
551 InsertResult::Full => {}
552 x => panic!("{x:?}"),
553 }
554
555 // One-by-one fill the bucket with connected nodes, replacing the disconnected ones.
556 for i in 0..bucket.capacity {
557 let (first, first_status) = bucket.iter().next().unwrap();
558 let first_disconnected = first.clone();
559 assert_eq!(first_status, NodeStatus::Disconnected);
560
561 // Add a connected node, which is expected to be pending, scheduled to
562 // replace the first (i.e. least-recently connected) node.
563 let key = Key::from(PeerId::random());
564 let node = Node { key, value: () };
565 match bucket.insert(node.clone(), NodeStatus::Connected) {
566 InsertResult::Pending { disconnected } => {
567 assert_eq!(disconnected, first_disconnected.key)
568 }
569 x => panic!("{x:?}"),
570 }
571
572 // Trying to insert another connected node fails.
573 match bucket.insert(node.clone(), NodeStatus::Connected) {
574 InsertResult::Full => {}
575 x => panic!("{x:?}"),
576 }
577
578 assert!(bucket.pending().is_some());
579
580 // Apply the pending node.
581 let pending = bucket.pending_mut().expect("No pending node.");
582 pending.set_ready_at(Instant::now().checked_sub(Duration::from_secs(1)).unwrap());
583 let result = bucket.apply_pending();
584 assert_eq!(
585 result,
586 Some(AppliedPending {
587 inserted: node.clone(),
588 evicted: Some(first_disconnected)
589 })
590 );
591 assert_eq!(Some((&node, NodeStatus::Connected)), bucket.iter().last());
592 assert!(bucket.pending().is_none());
593 assert_eq!(Some(bucket.capacity - (i + 1)), bucket.first_connected_pos);
594 }
595
596 assert!(bucket.pending().is_none());
597 assert_eq!(bucket.capacity, bucket.num_entries());
598
599 // Trying to insert another connected node fails.
600 let key = Key::from(PeerId::random());
601 let node = Node { key, value: () };
602 match bucket.insert(node, NodeStatus::Connected) {
603 InsertResult::Full => {}
604 x => panic!("{x:?}"),
605 }
606 }
607
608 #[test]
609 fn full_bucket_discard_pending() {
610 let mut bucket = KBucket::<Key<PeerId>, ()>::default();
611 fill_bucket(&mut bucket, NodeStatus::Disconnected);
612 let (first, _) = bucket.iter().next().unwrap();
613 let first_disconnected = first.clone();
614
615 // Add a connected pending node.
616 let key = Key::from(PeerId::random());
617 let node = Node { key, value: () };
618 if let InsertResult::Pending { disconnected } = bucket.insert(node, NodeStatus::Connected) {
619 assert_eq!(&disconnected, &first_disconnected.key);
620 } else {
621 panic!()
622 }
623 assert!(bucket.pending().is_some());
624
625 // Update the status of the first disconnected node to be connected.
626 bucket.update(&first_disconnected.key, NodeStatus::Connected);
627
628 // The pending node has been discarded.
629 assert!(bucket.pending().is_none());
630 assert!(bucket.iter().all(|(n, _)| n.key != key));
631
632 // The initially disconnected node is now the most-recently connected.
633 assert_eq!(
634 Some((&first_disconnected, NodeStatus::Connected)),
635 bucket.iter().last()
636 );
637 assert_eq!(
638 bucket.position(&first_disconnected.key).map(|p| p.0),
639 bucket.first_connected_pos
640 );
641 assert_eq!(1, bucket.num_connected());
642 assert_eq!(bucket.capacity - 1, bucket.num_disconnected());
643 }
644
645 #[test]
646 fn bucket_update() {
647 fn prop(mut bucket: KBucket<Key<PeerId>, ()>, pos: Position, status: NodeStatus) -> bool {
648 let num_nodes = bucket.num_entries();
649
650 // Capture position and key of the random node to update.
651 let pos = pos.0 % num_nodes;
652 let key = bucket.nodes[pos].key;
653
654 // Record the (ordered) list of status of all nodes in the bucket.
655 let mut expected = bucket.iter().map(|(n, s)| (n.key, s)).collect::<Vec<_>>();
656
657 // Update the node in the bucket.
658 bucket.update(&key, status);
659
660 // Check that the bucket now contains the node with the new status,
661 // preserving the status and relative order of all other nodes.
662 let expected_pos = match status {
663 NodeStatus::Connected => num_nodes - 1,
664 NodeStatus::Disconnected => bucket.first_connected_pos.unwrap_or(num_nodes) - 1,
665 };
666 expected.remove(pos);
667 expected.insert(expected_pos, (key, status));
668 let actual = bucket.iter().map(|(n, s)| (n.key, s)).collect::<Vec<_>>();
669 expected == actual
670 }
671
672 quickcheck(prop as fn(_, _, _) -> _);
673 }
674
675 #[test]
676 fn test_custom_bucket_size() {
677 let bucket_sizes: [NonZeroUsize; 4] = [
678 NonZeroUsize::new(2).unwrap(),
679 NonZeroUsize::new(20).unwrap(),
680 NonZeroUsize::new(200).unwrap(),
681 NonZeroUsize::new(2000).unwrap(),
682 ];
683 for &size in &bucket_sizes {
684 let mut config = KBucketConfig::default();
685 config.set_bucket_size(size);
686 let mut bucket = KBucket::<Key<PeerId>, ()>::new(config);
687 fill_bucket(&mut bucket, NodeStatus::Disconnected);
688 assert_eq!(size.get(), bucket.num_entries());
689 }
690 }
691}