Commit 3e48aa73 authored by Jim Newsome's avatar Jim Newsome
Browse files

NetDir::pick_n_relays: move implementation to a testable helper

parent 5e8bf486
Loading
Loading
Loading
Loading
+60 −35
Original line number Diff line number Diff line
@@ -1686,11 +1686,11 @@ impl NetDir {
        }
    }

    /// Choose `n` relay at random.
    /// Choose `n` items (relays) at random, using the provided weights.
    ///
    /// Each relay is chosen with probability proportional to its weight
    /// in the role `role`, and is only selected if the predicate `usable`
    /// returns true for it.
    /// This is intended as an internal, easier-to-test, implementation of
    /// `pick_n_relays`. `T` is generic for testing, but intended to be `Relay`;
    /// we call the items relays below.
    ///
    /// Relays are chosen without replacement: no relay will be
    /// returned twice. Therefore, the resulting vector may be smaller
@@ -1700,30 +1700,17 @@ impl NetDir {
    /// are no relays with nonzero weight where `usable` returned
    /// true.
    #[allow(clippy::cognitive_complexity)] // all due to tracing crate.
    pub fn pick_n_relays<'a, R, P>(
        &'a self,
    fn pick_n_filtered_weighted<R, T>(
        rng: &mut R,
        n: usize,
        role: WeightRole,
        mut usable: P,
    ) -> Vec<Relay<'a>>
        filtered_weighted_items: &[(T, u64)],
    ) -> Vec<T>
    where
        R: rand::Rng,
        P: FnMut(&Relay<'a>) -> bool,
        T: Clone,
    {
        let filtered_weighted_relays: Vec<_> = self
            .relays()
            .filter_map(|r| {
                if usable(&r) {
                    Some((r.clone(), self.weights.weight_rs_for_role(r.rs, role)))
                } else {
                    None
                }
            })
            .collect();

        // NOTE: See discussion in pick_relay().
        let mut sampled_relays = match filtered_weighted_relays[..].sample_weighted(
        let mut sampled_relays = match filtered_weighted_items[..].sample_weighted(
            rng,
            n,
            |(_r, w)| *w as f64,
@@ -1732,16 +1719,17 @@ impl NetDir {
                // Too few relays had nonzero weights: return all of those that are okay.
                // (This is behavior used to come up with rand 0.9; it no longer does.
                // We still detect it.)
                let nonzero_weight_relays: Vec<_> = filtered_weighted_relays
                let nonzero_weight_relays: Vec<_> = filtered_weighted_items
                    .iter()
                    .filter_map(|(r, w)| if *w > 0 { Some(r) } else { None })
                    .cloned()
                    .collect();
                if nonzero_weight_relays.is_empty() {
                    warn!(?self.weights, ?role,
                    warn!(
                        "After filtering, all {} relays had zero weight! Picking some at random. See bug #1907.",
                          filtered_weighted_relays.len());
                    let filtered_relays: Vec<_> = filtered_weighted_relays
                        filtered_weighted_items.len()
                    );
                    let filtered_relays: Vec<_> = filtered_weighted_items
                        .iter()
                        .map(|(r, _w)| r.clone())
                        .collect();
@@ -1751,9 +1739,11 @@ impl NetDir {
                        filtered_relays
                    }
                } else {
                    warn!(?self.weights, ?role,
                    warn!(
                        "After filtering, only had {}/{} relays with nonzero weight. Returning them all. See bug #1907.",
                           nonzero_weight_relays.len(), filtered_weighted_relays.len());
                        nonzero_weight_relays.len(),
                        filtered_weighted_items.len()
                    );
                    nonzero_weight_relays
                }
            }
@@ -1762,12 +1752,14 @@ impl NetDir {
                Vec::new()
            }
            Ok(iter) => {
                let selection: Vec<_> = iter.map(|(r, _w)| Relay::clone(r)).collect();
                if selection.len() < n && selection.len() < filtered_weighted_relays.len() {
                    warn!(?self.weights, ?role,
                let selection: Vec<_> = iter.map(|(r, _w)| r.clone()).collect();
                if selection.len() < n && selection.len() < filtered_weighted_items.len() {
                    warn!(
                        "sample_weighted returned only {returned}, despite requesting {n}, \
                          and having {filtered_len} available after filtering. See bug #1907.",
                          returned=selection.len(), filtered_len=filtered_weighted_relays.len());
                        returned = selection.len(),
                        filtered_len = filtered_weighted_items.len()
                    );
                }
                selection
            }
@@ -1776,6 +1768,39 @@ impl NetDir {
        sampled_relays
    }

    /// Choose `n` relay at random.
    ///
    /// Each relay is chosen with probability proportional to its weight
    /// in the role `role`, and is only selected if the predicate `usable`
    /// returns true for it.
    ///
    /// Relays are chosen without replacement: no relay will be
    /// returned twice. Therefore, the resulting vector may be smaller
    /// than `n` if we happen to have fewer than `n` appropriate relays.
    ///
    /// This function returns an empty vector if (and only if) there
    /// are no relays with nonzero weight where `usable` returned
    /// true.
    #[allow(clippy::cognitive_complexity)] // all due to tracing crate.
    pub fn pick_n_relays<'a, R, P>(
        &'a self,
        rng: &mut R,
        n: usize,
        role: WeightRole,
        usable: P,
    ) -> Vec<Relay<'a>>
    where
        R: rand::Rng,
        P: FnMut(&Relay<'a>) -> bool,
    {
        let filtered_weighted_relays: Vec<(Relay<'a>, u64)> = self
            .relays()
            .filter(usable)
            .map(|r| (r.clone(), self.weights.weight_rs_for_role(r.rs, role)))
            .collect();
        return NetDir::pick_n_filtered_weighted(rng, n, filtered_weighted_relays.as_slice());
    }

    /// Compute the weight with which `relay` will be selected for a given
    /// `role`.
    pub fn relay_weight<'a>(&'a self, relay: &Relay<'a>, role: WeightRole) -> RelayWeight {