libp2p_kad/
kbucket.rs

1// Copyright 2018 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//! Implementation of a Kademlia routing table as used by a single peer
22//! participating in a Kademlia DHT.
23//!
24//! The entry point for the API of this module is a [`KBucketsTable`].
25//!
26//! ## Pending Insertions
27//!
28//! When the bucket associated with the `Key` of an inserted entry is full
29//! but contains disconnected nodes, it accepts a [`PendingEntry`].
30//! Pending entries are inserted lazily when their timeout is found to be expired
31//! upon querying the `KBucketsTable`. When that happens, the `KBucketsTable` records
32//! an [`AppliedPending`] result which must be consumed by calling [`take_applied_pending`]
33//! regularly and / or after performing lookup operations like [`entry`] and [`closest`].
34//!
35//! [`entry`]: KBucketsTable::entry
36//! [`closest`]: KBucketsTable::closest
37//! [`AppliedPending`]: bucket::AppliedPending
38//! [`take_applied_pending`]: KBucketsTable::take_applied_pending
39//! [`PendingEntry`]: entry::PendingEntry
40
41// [Implementation Notes]
42//
43// 1. Routing Table Layout
44//
45// The routing table is currently implemented as a fixed-size "array" of
46// buckets, ordered by increasing distance relative to a local key
47// that identifies the local peer. This is an often-used, simplified
48// implementation that approximates the properties of the b-tree (or prefix tree)
49// implementation described in the full paper [0], whereby buckets are split on-demand.
50// This should be treated as an implementation detail, however, so that the
51// implementation may change in the future without breaking the API.
52//
53// 2. Replacement Cache
54//
55// In this implementation, the "replacement cache" for unresponsive peers
56// consists of a single entry per bucket. Furthermore, this implementation is
57// currently tailored to connection-oriented transports, meaning that the
58// "LRU"-based ordering of entries in a bucket is actually based on the last reported
59// connection status of the corresponding peers, from least-recently (dis)connected to
60// most-recently (dis)connected, and controlled through the `Entry` API. As a result,
61// the nodes in the buckets are not reordered as a result of RPC activity, but only as a
62// result of nodes being marked as connected or disconnected. In particular,
63// if a bucket is full and contains only entries for peers that are considered
64// connected, no pending entry is accepted. See the `bucket` submodule for
65// further details.
66//
67// [0]: https://pdos.csail.mit.edu/~petar/papers/maymounkov-kademlia-lncs.pdf
68
69mod bucket;
70mod entry;
71#[allow(clippy::ptr_offset_with_cast)]
72#[allow(clippy::assign_op_pattern)]
73mod key;
74
75pub use bucket::NodeStatus;
76pub use entry::*;
77
78use arrayvec::{self, ArrayVec};
79use bucket::KBucket;
80use std::collections::VecDeque;
81use std::time::{Duration, Instant};
82
83/// Maximum number of k-buckets.
84const NUM_BUCKETS: usize = 256;
85
86/// A `KBucketsTable` represents a Kademlia routing table.
87#[derive(Debug, Clone)]
88pub(crate) struct KBucketsTable<TKey, TVal> {
89    /// The key identifying the local peer that owns the routing table.
90    local_key: TKey,
91    /// The buckets comprising the routing table.
92    buckets: Vec<KBucket<TKey, TVal>>,
93    /// The list of evicted entries that have been replaced with pending
94    /// entries since the last call to [`KBucketsTable::take_applied_pending`].
95    applied_pending: VecDeque<AppliedPending<TKey, TVal>>,
96}
97
98/// A (type-safe) index into a `KBucketsTable`, i.e. a non-negative integer in the
99/// interval `[0, NUM_BUCKETS)`.
100#[derive(Debug, Copy, Clone, PartialEq, Eq)]
101struct BucketIndex(usize);
102
103impl BucketIndex {
104    /// Creates a new `BucketIndex` for a `Distance`.
105    ///
106    /// The given distance is interpreted as the distance from a `local_key` of
107    /// a `KBucketsTable`. If the distance is zero, `None` is returned, in
108    /// recognition of the fact that the only key with distance `0` to a
109    /// `local_key` is the `local_key` itself, which does not belong in any
110    /// bucket.
111    fn new(d: &Distance) -> Option<BucketIndex> {
112        d.ilog2().map(|i| BucketIndex(i as usize))
113    }
114
115    /// Gets the index value as an unsigned integer.
116    fn get(&self) -> usize {
117        self.0
118    }
119
120    /// Returns the minimum inclusive and maximum inclusive [`Distance`]
121    /// included in the bucket for this index.
122    fn range(&self) -> (Distance, Distance) {
123        let min = Distance(U256::pow(U256::from(2), U256::from(self.0)));
124        if self.0 == usize::from(u8::MAX) {
125            (min, Distance(U256::MAX))
126        } else {
127            let max = Distance(U256::pow(U256::from(2), U256::from(self.0 + 1)) - 1);
128            (min, max)
129        }
130    }
131
132    /// Generates a random distance that falls into the bucket for this index.
133    fn rand_distance(&self, rng: &mut impl rand::Rng) -> Distance {
134        let mut bytes = [0u8; 32];
135        let quot = self.0 / 8;
136        for i in 0..quot {
137            bytes[31 - i] = rng.gen();
138        }
139        let rem = (self.0 % 8) as u32;
140        let lower = usize::pow(2, rem);
141        let upper = usize::pow(2, rem + 1);
142        bytes[31 - quot] = rng.gen_range(lower..upper) as u8;
143        Distance(U256::from(bytes))
144    }
145}
146
147impl<TKey, TVal> KBucketsTable<TKey, TVal>
148where
149    TKey: Clone + AsRef<KeyBytes>,
150    TVal: Clone,
151{
152    /// Creates a new, empty Kademlia routing table with entries partitioned
153    /// into buckets as per the Kademlia protocol.
154    ///
155    /// The given `pending_timeout` specifies the duration after creation of
156    /// a [`PendingEntry`] after which it becomes eligible for insertion into
157    /// a full bucket, replacing the least-recently (dis)connected node.
158    pub(crate) fn new(local_key: TKey, pending_timeout: Duration) -> Self {
159        KBucketsTable {
160            local_key,
161            buckets: (0..NUM_BUCKETS)
162                .map(|_| KBucket::new(pending_timeout))
163                .collect(),
164            applied_pending: VecDeque::new(),
165        }
166    }
167
168    /// Returns the local key.
169    pub(crate) fn local_key(&self) -> &TKey {
170        &self.local_key
171    }
172
173    /// Returns an `Entry` for the given key, representing the state of the entry
174    /// in the routing table.
175    pub(crate) fn entry<'a>(&'a mut self, key: &'a TKey) -> Entry<'a, TKey, TVal> {
176        let index = BucketIndex::new(&self.local_key.as_ref().distance(key));
177        if let Some(i) = index {
178            let bucket = &mut self.buckets[i.get()];
179            if let Some(applied) = bucket.apply_pending() {
180                self.applied_pending.push_back(applied)
181            }
182            Entry::new(bucket, key)
183        } else {
184            Entry::SelfEntry
185        }
186    }
187
188    /// Returns an iterator over all buckets.
189    ///
190    /// The buckets are ordered by proximity to the `local_key`, i.e. the first
191    /// bucket is the closest bucket (containing at most one key).
192    pub(crate) fn iter(&mut self) -> impl Iterator<Item = KBucketRef<'_, TKey, TVal>> + '_ {
193        let applied_pending = &mut self.applied_pending;
194        self.buckets.iter_mut().enumerate().map(move |(i, b)| {
195            if let Some(applied) = b.apply_pending() {
196                applied_pending.push_back(applied)
197            }
198            KBucketRef {
199                index: BucketIndex(i),
200                bucket: b,
201            }
202        })
203    }
204
205    /// Returns the bucket for the distance to the given key.
206    ///
207    /// Returns `None` if the given key refers to the local key.
208    pub(crate) fn bucket<K>(&mut self, key: &K) -> Option<KBucketRef<'_, TKey, TVal>>
209    where
210        K: AsRef<KeyBytes>,
211    {
212        let d = self.local_key.as_ref().distance(key);
213        if let Some(index) = BucketIndex::new(&d) {
214            let bucket = &mut self.buckets[index.0];
215            if let Some(applied) = bucket.apply_pending() {
216                self.applied_pending.push_back(applied)
217            }
218            Some(KBucketRef { bucket, index })
219        } else {
220            None
221        }
222    }
223
224    /// Consumes the next applied pending entry, if any.
225    ///
226    /// When an entry is attempted to be inserted and the respective bucket is full,
227    /// it may be recorded as pending insertion after a timeout, see [`InsertResult::Pending`].
228    ///
229    /// If the oldest currently disconnected entry in the respective bucket does not change
230    /// its status until the timeout of pending entry expires, it is evicted and
231    /// the pending entry inserted instead. These insertions of pending entries
232    /// happens lazily, whenever the `KBucketsTable` is accessed, and the corresponding
233    /// buckets are updated accordingly. The fact that a pending entry was applied is
234    /// recorded in the `KBucketsTable` in the form of `AppliedPending` results, which must be
235    /// consumed by calling this function.
236    pub(crate) fn take_applied_pending(&mut self) -> Option<AppliedPending<TKey, TVal>> {
237        self.applied_pending.pop_front()
238    }
239
240    /// Returns an iterator over the keys closest to `target`, ordered by
241    /// increasing distance.
242    pub(crate) fn closest_keys<'a, T>(
243        &'a mut self,
244        target: &'a T,
245    ) -> impl Iterator<Item = TKey> + 'a
246    where
247        T: AsRef<KeyBytes>,
248    {
249        let distance = self.local_key.as_ref().distance(target);
250        ClosestIter {
251            target,
252            iter: None,
253            table: self,
254            buckets_iter: ClosestBucketsIter::new(distance),
255            fmap: |b: &KBucket<TKey, _>| -> ArrayVec<_, { K_VALUE.get() }> {
256                b.iter().map(|(n, _)| n.key.clone()).collect()
257            },
258        }
259    }
260
261    /// Returns an iterator over the nodes closest to the `target` key, ordered by
262    /// increasing distance.
263    pub(crate) fn closest<'a, T>(
264        &'a mut self,
265        target: &'a T,
266    ) -> impl Iterator<Item = EntryView<TKey, TVal>> + 'a
267    where
268        T: Clone + AsRef<KeyBytes>,
269        TVal: Clone,
270    {
271        let distance = self.local_key.as_ref().distance(target);
272        ClosestIter {
273            target,
274            iter: None,
275            table: self,
276            buckets_iter: ClosestBucketsIter::new(distance),
277            fmap: |b: &KBucket<_, TVal>| -> ArrayVec<_, { K_VALUE.get() }> {
278                b.iter()
279                    .map(|(n, status)| EntryView {
280                        node: n.clone(),
281                        status,
282                    })
283                    .collect()
284            },
285        }
286    }
287
288    /// Counts the number of nodes between the local node and the node
289    /// closest to `target`.
290    ///
291    /// The number of nodes between the local node and the target are
292    /// calculated by backtracking from the target towards the local key.
293    pub(crate) fn count_nodes_between<T>(&mut self, target: &T) -> usize
294    where
295        T: AsRef<KeyBytes>,
296    {
297        let local_key = self.local_key.clone();
298        let distance = target.as_ref().distance(&local_key);
299        let mut iter = ClosestBucketsIter::new(distance).take_while(|i| i.get() != 0);
300        if let Some(i) = iter.next() {
301            let num_first = self.buckets[i.get()]
302                .iter()
303                .filter(|(n, _)| n.key.as_ref().distance(&local_key) <= distance)
304                .count();
305            let num_rest: usize = iter.map(|i| self.buckets[i.get()].num_entries()).sum();
306            num_first + num_rest
307        } else {
308            0
309        }
310    }
311}
312
313/// An iterator over (some projection of) the closest entries in a
314/// `KBucketsTable` w.r.t. some target `Key`.
315struct ClosestIter<'a, TTarget, TKey, TVal, TMap, TOut> {
316    /// A reference to the target key whose distance to the local key determines
317    /// the order in which the buckets are traversed. The resulting
318    /// array from projecting the entries of each bucket using `fmap` is
319    /// sorted according to the distance to the target.
320    target: &'a TTarget,
321    /// A reference to all buckets of the `KBucketsTable`.
322    table: &'a mut KBucketsTable<TKey, TVal>,
323    /// The iterator over the bucket indices in the order determined by the
324    /// distance of the local key to the target.
325    buckets_iter: ClosestBucketsIter,
326    /// The iterator over the entries in the currently traversed bucket.
327    iter: Option<arrayvec::IntoIter<TOut, { K_VALUE.get() }>>,
328    /// The projection function / mapping applied on each bucket as
329    /// it is encountered, producing the next `iter`ator.
330    fmap: TMap,
331}
332
333/// An iterator over the bucket indices, in the order determined by the `Distance` of
334/// a target from the `local_key`, such that the entries in the buckets are incrementally
335/// further away from the target, starting with the bucket covering the target.
336struct ClosestBucketsIter {
337    /// The distance to the `local_key`.
338    distance: Distance,
339    /// The current state of the iterator.
340    state: ClosestBucketsIterState,
341}
342
343/// Operating states of a `ClosestBucketsIter`.
344enum ClosestBucketsIterState {
345    /// The starting state of the iterator yields the first bucket index and
346    /// then transitions to `ZoomIn`.
347    Start(BucketIndex),
348    /// The iterator "zooms in" to to yield the next bucket cotaining nodes that
349    /// are incrementally closer to the local node but further from the `target`.
350    /// These buckets are identified by a `1` in the corresponding bit position
351    /// of the distance bit string. When bucket `0` is reached, the iterator
352    /// transitions to `ZoomOut`.
353    ZoomIn(BucketIndex),
354    /// Once bucket `0` has been reached, the iterator starts "zooming out"
355    /// to buckets containing nodes that are incrementally further away from
356    /// both the local key and the target. These are identified by a `0` in
357    /// the corresponding bit position of the distance bit string. When bucket
358    /// `255` is reached, the iterator transitions to state `Done`.
359    ZoomOut(BucketIndex),
360    /// The iterator is in this state once it has visited all buckets.
361    Done,
362}
363
364impl ClosestBucketsIter {
365    fn new(distance: Distance) -> Self {
366        let state = match BucketIndex::new(&distance) {
367            Some(i) => ClosestBucketsIterState::Start(i),
368            None => ClosestBucketsIterState::Start(BucketIndex(0)),
369        };
370        Self { distance, state }
371    }
372
373    fn next_in(&self, i: BucketIndex) -> Option<BucketIndex> {
374        (0..i.get()).rev().find_map(|i| {
375            if self.distance.0.bit(i) {
376                Some(BucketIndex(i))
377            } else {
378                None
379            }
380        })
381    }
382
383    fn next_out(&self, i: BucketIndex) -> Option<BucketIndex> {
384        (i.get() + 1..NUM_BUCKETS).find_map(|i| {
385            if !self.distance.0.bit(i) {
386                Some(BucketIndex(i))
387            } else {
388                None
389            }
390        })
391    }
392}
393
394impl Iterator for ClosestBucketsIter {
395    type Item = BucketIndex;
396
397    fn next(&mut self) -> Option<Self::Item> {
398        match self.state {
399            ClosestBucketsIterState::Start(i) => {
400                self.state = ClosestBucketsIterState::ZoomIn(i);
401                Some(i)
402            }
403            ClosestBucketsIterState::ZoomIn(i) => {
404                if let Some(i) = self.next_in(i) {
405                    self.state = ClosestBucketsIterState::ZoomIn(i);
406                    Some(i)
407                } else {
408                    let i = BucketIndex(0);
409                    self.state = ClosestBucketsIterState::ZoomOut(i);
410                    Some(i)
411                }
412            }
413            ClosestBucketsIterState::ZoomOut(i) => {
414                if let Some(i) = self.next_out(i) {
415                    self.state = ClosestBucketsIterState::ZoomOut(i);
416                    Some(i)
417                } else {
418                    self.state = ClosestBucketsIterState::Done;
419                    None
420                }
421            }
422            ClosestBucketsIterState::Done => None,
423        }
424    }
425}
426
427impl<TTarget, TKey, TVal, TMap, TOut> Iterator for ClosestIter<'_, TTarget, TKey, TVal, TMap, TOut>
428where
429    TTarget: AsRef<KeyBytes>,
430    TKey: Clone + AsRef<KeyBytes>,
431    TVal: Clone,
432    TMap: Fn(&KBucket<TKey, TVal>) -> ArrayVec<TOut, { K_VALUE.get() }>,
433    TOut: AsRef<KeyBytes>,
434{
435    type Item = TOut;
436
437    fn next(&mut self) -> Option<Self::Item> {
438        loop {
439            match &mut self.iter {
440                Some(iter) => match iter.next() {
441                    Some(k) => return Some(k),
442                    None => self.iter = None,
443                },
444                None => {
445                    if let Some(i) = self.buckets_iter.next() {
446                        let bucket = &mut self.table.buckets[i.get()];
447                        if let Some(applied) = bucket.apply_pending() {
448                            self.table.applied_pending.push_back(applied)
449                        }
450                        let mut v = (self.fmap)(bucket);
451                        v.sort_by(|a, b| {
452                            self.target
453                                .as_ref()
454                                .distance(a.as_ref())
455                                .cmp(&self.target.as_ref().distance(b.as_ref()))
456                        });
457                        self.iter = Some(v.into_iter());
458                    } else {
459                        return None;
460                    }
461                }
462            }
463        }
464    }
465}
466
467/// A reference to a bucket.
468pub struct KBucketRef<'a, TKey, TVal> {
469    index: BucketIndex,
470    bucket: &'a mut KBucket<TKey, TVal>,
471}
472
473impl<'a, TKey, TVal> KBucketRef<'a, TKey, TVal>
474where
475    TKey: Clone + AsRef<KeyBytes>,
476    TVal: Clone,
477{
478    /// Returns the minimum inclusive and maximum inclusive distance for
479    /// this bucket.
480    pub fn range(&self) -> (Distance, Distance) {
481        self.index.range()
482    }
483
484    /// Checks whether the bucket is empty.
485    pub fn is_empty(&self) -> bool {
486        self.num_entries() == 0
487    }
488
489    /// Returns the number of entries in the bucket.
490    pub fn num_entries(&self) -> usize {
491        self.bucket.num_entries()
492    }
493
494    /// Returns true if the bucket has a pending node.
495    pub fn has_pending(&self) -> bool {
496        self.bucket.pending().map_or(false, |n| !n.is_ready())
497    }
498
499    /// Tests whether the given distance falls into this bucket.
500    pub fn contains(&self, d: &Distance) -> bool {
501        BucketIndex::new(d).map_or(false, |i| i == self.index)
502    }
503
504    /// Generates a random distance that falls into this bucket.
505    ///
506    /// Together with a known key `a` (e.g. the local key), a random distance `d` for
507    /// this bucket w.r.t `k` gives rise to the corresponding (random) key `b` s.t.
508    /// the XOR distance between `a` and `b` is `d`. In other words, it gives
509    /// rise to a random key falling into this bucket. See [`key::Key::for_distance`].
510    pub fn rand_distance(&self, rng: &mut impl rand::Rng) -> Distance {
511        self.index.rand_distance(rng)
512    }
513
514    /// Returns an iterator over the entries in the bucket.
515    pub fn iter(&'a self) -> impl Iterator<Item = EntryRefView<'a, TKey, TVal>> {
516        self.bucket.iter().map(move |(n, status)| EntryRefView {
517            node: NodeRefView {
518                key: &n.key,
519                value: &n.value,
520            },
521            status,
522        })
523    }
524}
525
526#[cfg(test)]
527mod tests {
528    use super::*;
529    use libp2p_identity::PeerId;
530    use quickcheck::*;
531
532    type TestTable = KBucketsTable<KeyBytes, ()>;
533
534    impl Arbitrary for TestTable {
535        fn arbitrary(g: &mut Gen) -> TestTable {
536            let local_key = Key::from(PeerId::random());
537            let timeout = Duration::from_secs(g.gen_range(1..360));
538            let mut table = TestTable::new(local_key.clone().into(), timeout);
539            let mut num_total = g.gen_range(0..100);
540            for (i, b) in &mut table.buckets.iter_mut().enumerate().rev() {
541                let ix = BucketIndex(i);
542                let num = g.gen_range(0..usize::min(K_VALUE.get(), num_total) + 1);
543                num_total -= num;
544                for _ in 0..num {
545                    let distance = ix.rand_distance(&mut rand::thread_rng());
546                    let key = local_key.for_distance(distance);
547                    let node = Node {
548                        key: key.clone(),
549                        value: (),
550                    };
551                    let status = NodeStatus::arbitrary(g);
552                    match b.insert(node, status) {
553                        InsertResult::Inserted => {}
554                        _ => panic!(),
555                    }
556                }
557            }
558            table
559        }
560    }
561
562    #[test]
563    fn buckets_are_non_overlapping_and_exhaustive() {
564        let local_key = Key::from(PeerId::random());
565        let timeout = Duration::from_secs(0);
566        let mut table = KBucketsTable::<KeyBytes, ()>::new(local_key.into(), timeout);
567
568        let mut prev_max = U256::from(0);
569
570        for bucket in table.iter() {
571            let (min, max) = bucket.range();
572            assert_eq!(Distance(prev_max + U256::from(1)), min);
573            prev_max = max.0;
574        }
575
576        assert_eq!(U256::MAX, prev_max);
577    }
578
579    #[test]
580    fn bucket_contains_range() {
581        fn prop(ix: u8) {
582            let index = BucketIndex(ix as usize);
583            let mut bucket = KBucket::<Key<PeerId>, ()>::new(Duration::from_secs(0));
584            let bucket_ref = KBucketRef {
585                index,
586                bucket: &mut bucket,
587            };
588
589            let (min, max) = bucket_ref.range();
590
591            assert!(min <= max);
592
593            assert!(bucket_ref.contains(&min));
594            assert!(bucket_ref.contains(&max));
595
596            if min != Distance(0.into()) {
597                // ^ avoid underflow
598                assert!(!bucket_ref.contains(&Distance(min.0 - 1)));
599            }
600
601            if max != Distance(U256::max_value()) {
602                // ^ avoid overflow
603                assert!(!bucket_ref.contains(&Distance(max.0 + 1)));
604            }
605        }
606
607        quickcheck(prop as fn(_));
608    }
609
610    #[test]
611    fn rand_distance() {
612        fn prop(ix: u8) -> bool {
613            let d = BucketIndex(ix as usize).rand_distance(&mut rand::thread_rng());
614            let n = U256::from(<[u8; 32]>::from(d.0));
615            let b = U256::from(2);
616            let e = U256::from(ix);
617            let lower = b.pow(e);
618            let upper = b.checked_pow(e + U256::from(1)).unwrap_or(U256::MAX) - U256::from(1);
619            lower <= n && n <= upper
620        }
621        quickcheck(prop as fn(_) -> _);
622    }
623
624    #[test]
625    fn entry_inserted() {
626        let local_key = Key::from(PeerId::random());
627        let other_id = Key::from(PeerId::random());
628
629        let mut table = KBucketsTable::<_, ()>::new(local_key, Duration::from_secs(5));
630        if let Entry::Absent(entry) = table.entry(&other_id) {
631            match entry.insert((), NodeStatus::Connected) {
632                InsertResult::Inserted => (),
633                _ => panic!(),
634            }
635        } else {
636            panic!()
637        }
638
639        let res = table.closest_keys(&other_id).collect::<Vec<_>>();
640        assert_eq!(res.len(), 1);
641        assert_eq!(res[0], other_id);
642    }
643
644    #[test]
645    fn entry_self() {
646        let local_key = Key::from(PeerId::random());
647        let mut table = KBucketsTable::<_, ()>::new(local_key.clone(), Duration::from_secs(5));
648        match table.entry(&local_key) {
649            Entry::SelfEntry => (),
650            _ => panic!(),
651        }
652    }
653
654    #[test]
655    fn closest() {
656        let local_key = Key::from(PeerId::random());
657        let mut table = KBucketsTable::<_, ()>::new(local_key, Duration::from_secs(5));
658        let mut count = 0;
659        loop {
660            if count == 100 {
661                break;
662            }
663            let key = Key::from(PeerId::random());
664            if let Entry::Absent(e) = table.entry(&key) {
665                match e.insert((), NodeStatus::Connected) {
666                    InsertResult::Inserted => count += 1,
667                    _ => continue,
668                }
669            } else {
670                panic!("entry exists")
671            }
672        }
673
674        let mut expected_keys: Vec<_> = table
675            .buckets
676            .iter()
677            .flat_map(|t| t.iter().map(|(n, _)| n.key.clone()))
678            .collect();
679
680        for _ in 0..10 {
681            let target_key = Key::from(PeerId::random());
682            let keys = table.closest_keys(&target_key).collect::<Vec<_>>();
683            // The list of keys is expected to match the result of a full-table scan.
684            expected_keys.sort_by_key(|k| k.distance(&target_key));
685            assert_eq!(keys, expected_keys);
686        }
687    }
688
689    #[test]
690    fn applied_pending() {
691        let local_key = Key::from(PeerId::random());
692        let mut table = KBucketsTable::<_, ()>::new(local_key.clone(), Duration::from_millis(1));
693        let expected_applied;
694        let full_bucket_index;
695        loop {
696            let key = Key::from(PeerId::random());
697            if let Entry::Absent(e) = table.entry(&key) {
698                match e.insert((), NodeStatus::Disconnected) {
699                    InsertResult::Full => {
700                        if let Entry::Absent(e) = table.entry(&key) {
701                            match e.insert((), NodeStatus::Connected) {
702                                InsertResult::Pending { disconnected } => {
703                                    expected_applied = AppliedPending {
704                                        inserted: Node {
705                                            key: key.clone(),
706                                            value: (),
707                                        },
708                                        evicted: Some(Node {
709                                            key: disconnected,
710                                            value: (),
711                                        }),
712                                    };
713                                    full_bucket_index = BucketIndex::new(&key.distance(&local_key));
714                                    break;
715                                }
716                                _ => panic!(),
717                            }
718                        } else {
719                            panic!()
720                        }
721                    }
722                    _ => continue,
723                }
724            } else {
725                panic!("entry exists")
726            }
727        }
728
729        // Expire the timeout for the pending entry on the full bucket.`
730        let full_bucket = &mut table.buckets[full_bucket_index.unwrap().get()];
731        let elapsed = Instant::now().checked_sub(Duration::from_secs(1)).unwrap();
732        full_bucket.pending_mut().unwrap().set_ready_at(elapsed);
733
734        match table.entry(&expected_applied.inserted.key) {
735            Entry::Present(_, NodeStatus::Connected) => {}
736            x => panic!("Unexpected entry: {x:?}"),
737        }
738
739        match table.entry(&expected_applied.evicted.as_ref().unwrap().key) {
740            Entry::Absent(_) => {}
741            x => panic!("Unexpected entry: {x:?}"),
742        }
743
744        assert_eq!(Some(expected_applied), table.take_applied_pending());
745        assert_eq!(None, table.take_applied_pending());
746    }
747
748    #[test]
749    fn count_nodes_between() {
750        fn prop(mut table: TestTable, target: Key<PeerId>) -> bool {
751            let num_to_target = table.count_nodes_between(&target);
752            let distance = table.local_key.distance(&target);
753            let base2 = U256::from(2);
754            let mut iter = ClosestBucketsIter::new(distance);
755            iter.all(|i| {
756                // Flip the distance bit related to the bucket.
757                let d = Distance(distance.0 ^ (base2.pow(U256::from(i.get()))));
758                let k = table.local_key.for_distance(d);
759                if distance.0.bit(i.get()) {
760                    // Bit flip `1` -> `0`, the key must be closer than `target`.
761                    d < distance && table.count_nodes_between(&k) <= num_to_target
762                } else {
763                    // Bit flip `0` -> `1`, the key must be farther than `target`.
764                    d > distance && table.count_nodes_between(&k) >= num_to_target
765                }
766            })
767        }
768
769        QuickCheck::new()
770            .tests(10)
771            .quickcheck(prop as fn(_, _) -> _)
772    }
773}