tor_circmgr/hspool/
pool.rs

1//! An internal pool object that we use to implement HsCircPool.
2
3use std::time::{Duration, Instant};
4
5use crate::{
6    hspool::{HsCircStem, HsCircStemKind},
7    AbstractCirc,
8};
9use rand::Rng;
10use tor_basic_utils::RngExt as _;
11
12/// A collection of circuits used to fulfil onion-service-related requests.
13pub(super) struct Pool<C: AbstractCirc> {
14    /// The collection of circuits themselves, in no particular order.
15    circuits: Vec<HsCircStem<C>>,
16
17    /// The number of NAIVE elements that we would like to have in our pool.
18    stem_target: usize,
19
20    /// The number of GUARDED elements that we would like to have in our pool.
21    guarded_stem_target: usize,
22
23    /// True if we have exhausted our pool since the last time we decided
24    /// whether to change our target level.
25    have_been_exhausted: bool,
26
27    /// True if we have been under 4/5 of our target since the last time we
28    /// decided whether to change it.
29    have_been_under_highwater: bool,
30
31    /// Last time when we changed our target size.
32    last_changed_target: Option<Instant>,
33}
34
35/// Our default (and minimum) target NAIVE pool size.
36const DEFAULT_NAIVE_STEM_TARGET: usize = 3;
37
38/// Our default (and minimum) target GUARDED pool size.
39const DEFAULT_GUARDED_STEM_TARGET: usize = 1;
40
41/// Our maximum target NAIVE pool size.  We will never let our NAIVE target grow above this
42/// value.
43const MAX_NAIVE_STEM_TARGET: usize = 384;
44
45/// Our maximum target GUARDED pool size.  We will never let our GUARDED target grow above this
46/// value.
47const MAX_GUARDED_STEM_TARGET: usize = 128;
48
49/// A type of circuit we would like to launch.
50///
51/// [`ForLaunch::note_circ_launched`] should be called whenever a circuit
52/// of this [`HsCircStemKind`] is launched, to decrement the internal target `count`.
53pub(super) struct ForLaunch<'a> {
54    /// The kind of circuit we want to launch.
55    kind: HsCircStemKind,
56    /// How many circuits of this kind do we need?
57    ///
58    /// This is a mutable reference to one of the target values from [`CircsToLaunch`];
59    /// we decrement it when we have launched a circuit of this type.
60    count: &'a mut usize,
61}
62
63impl<'a> ForLaunch<'a> {
64    /// A circuit was launched, decrement the current target for its kind.
65    pub(super) fn note_circ_launched(self) {
66        *self.count -= 1;
67    }
68
69    /// The kind of circuit we want to launch.
70    pub(super) fn kind(&self) -> HsCircStemKind {
71        self.kind
72    }
73}
74
75/// The circuits we need to launch.
76pub(super) struct CircsToLaunch {
77    /// The number of NAIVE circuits we want to launch.
78    stem_target: usize,
79    /// The number of GUARDED circuits we want to launch.
80    guarded_stem_target: usize,
81}
82
83impl CircsToLaunch {
84    /// Return a [`ForLaunch`] representing a circuit we would like to launch.
85    pub(super) fn for_launch(&mut self) -> ForLaunch {
86        // We start by launching NAIVE circuits.
87        if self.stem_target > 0 {
88            ForLaunch {
89                kind: HsCircStemKind::Naive,
90                count: &mut self.stem_target,
91            }
92        } else {
93            // If we have enough NAIVE circuits, we can start launching GUARDED ones too.
94            ForLaunch {
95                kind: HsCircStemKind::Guarded,
96                count: &mut self.guarded_stem_target,
97            }
98        }
99    }
100
101    /// Return the number of NAIVE circuits we would like to launch.
102    pub(super) fn stem(&self) -> usize {
103        self.stem_target
104    }
105
106    /// Return the number of GUARDED circuits we would like to launch.
107    pub(super) fn guarded_stem(&self) -> usize {
108        self.guarded_stem_target
109    }
110
111    /// Return the total number of circuits we would currently like to launch.
112    pub(super) fn n_to_launch(&self) -> usize {
113        self.stem_target + self.guarded_stem_target
114    }
115}
116
117impl<C: AbstractCirc> Default for Pool<C> {
118    fn default() -> Self {
119        Self {
120            circuits: Vec::new(),
121            stem_target: DEFAULT_NAIVE_STEM_TARGET,
122            guarded_stem_target: DEFAULT_GUARDED_STEM_TARGET,
123            have_been_exhausted: false,
124            have_been_under_highwater: false,
125            last_changed_target: None,
126        }
127    }
128}
129
130impl<C: AbstractCirc> Pool<C> {
131    /// Add `circ` to this pool
132    pub(super) fn insert(&mut self, circ: HsCircStem<C>) {
133        self.circuits.push(circ);
134    }
135
136    /// Remove every circuit from this pool for which `f` returns false.
137    pub(super) fn retain<F>(&mut self, f: F)
138    where
139        F: FnMut(&HsCircStem<C>) -> bool,
140    {
141        self.circuits.retain(f);
142    }
143
144    /// Return true if we are very low on circuits and should build more immediately.
145    pub(super) fn very_low(&self) -> bool {
146        self.circuits.len() <= self.target() / 3
147    }
148
149    /// Return a [`CircsToLaunch`] describing the circuits we would currently like to launch.
150    pub(super) fn circs_to_launch(&self) -> CircsToLaunch {
151        CircsToLaunch {
152            stem_target: self.stems_to_launch(),
153            guarded_stem_target: self.guarded_stems_to_launch(),
154        }
155    }
156
157    /// Return the number of NAIVE circuits we would currently like to launch.
158    fn stems_to_launch(&self) -> usize {
159        let circ_count = self
160            .circuits
161            .iter()
162            .filter(|c| c.kind == HsCircStemKind::Naive)
163            .count();
164
165        self.stem_target.saturating_sub(circ_count)
166    }
167
168    /// Return the number of GUARDED circuits we would currently like to launch.
169    fn guarded_stems_to_launch(&self) -> usize {
170        let circ_count = self
171            .circuits
172            .iter()
173            .filter(|c| c.kind == HsCircStemKind::Guarded)
174            .count();
175
176        self.guarded_stem_target.saturating_sub(circ_count)
177    }
178
179    /// Return the total number of circuits we would like to launch.
180    ///
181    /// We do not discard when we are _above_ this threshold, but we do
182    /// try to build when we are low.
183    fn target(&self) -> usize {
184        self.stem_target + self.guarded_stem_target
185    }
186
187    /// If there is any circuit in this pool for which `f`  returns true and that satisfies
188    /// all of the specified [`HsCircPrefs`], return one such circuit at random, and remove
189    /// it from the pool.
190    ///
191    /// If none of the circuits satisfy `prefs`, return a randomly selected circuit for which `f`
192    /// returns true, and remove it from the pool.
193    pub(super) fn take_one_where<R, F>(
194        &mut self,
195        rng: &mut R,
196        f: F,
197        prefs: &HsCircPrefs,
198    ) -> Option<HsCircStem<C>>
199    where
200        R: Rng,
201        F: Fn(&HsCircStem<C>) -> bool,
202    {
203        let rv = match random_idx_where(rng, &mut self.circuits[..], |circ_stem| {
204            // First, check if any circuit matches _all_ the prefs
205            circ_stem.satisfies_prefs(prefs) && f(circ_stem)
206        })
207        .or_else(|| {
208            // Select a circuit satisfying `f` at random.
209            random_idx_where(rng, &mut self.circuits[..], f)
210        }) {
211            Some(idx) => Some(self.circuits.swap_remove(idx)),
212            None => None,
213        };
214
215        if self.circuits.is_empty() {
216            self.have_been_exhausted = true;
217            self.have_been_under_highwater = true;
218        } else if self.circuits.len() < self.target() * 4 / 5 {
219            self.have_been_under_highwater = true;
220        }
221
222        rv
223    }
224
225    /// Update the target sizes for our pool.
226    ///
227    /// This updates our target numbers of NAIVE and GUARDED circuits.
228    pub(super) fn update_target_size(&mut self, now: Instant) {
229        /// Minimum amount of time that must elapse between a change and a
230        /// decision to grow our pool.  We use this to control the rate of
231        /// growth and make sure that we are allowing enough time for circuits
232        /// to complete.
233        const MIN_TIME_TO_GROW: Duration = Duration::from_secs(120);
234        /// Minimum amount of time that must elapse between a target change and
235        /// a decisions to shrink our target.  We use this to make sure that we
236        /// aren't shrinking too rapidly, and that we are allowing enough time
237        /// for the pool to actually get used.
238        const MIN_TIME_TO_SHRINK: Duration = Duration::from_secs(600);
239
240        let last_changed = self.last_changed_target.get_or_insert(now);
241        let time_since_last_change = now.saturating_duration_since(*last_changed);
242
243        // TODO: we may want to have separate have_been_exhausted/have_been_under_highwater
244        // flags for NAIVE and GUARDED circuits.
245        //
246        // TODO: stem_target and guarded_stem_target currently grow/shrink at the same rate,
247        // which is not ideal.
248        //
249        // Instead, we should switch to an adaptive strategy, where the two targets are updated
250        // based on how many NAIVE/GUARDED circuit requests we got.
251        if self.have_been_exhausted {
252            if time_since_last_change < MIN_TIME_TO_GROW {
253                return;
254            }
255            self.stem_target *= 2;
256            self.guarded_stem_target *= 2;
257        } else if !self.have_been_under_highwater {
258            if time_since_last_change < MIN_TIME_TO_SHRINK {
259                return;
260            }
261
262            self.stem_target /= 2;
263            self.guarded_stem_target /= 2;
264        }
265        self.last_changed_target = Some(now);
266        self.stem_target = self
267            .stem_target
268            .clamp(DEFAULT_NAIVE_STEM_TARGET, MAX_NAIVE_STEM_TARGET);
269        self.guarded_stem_target = self
270            .guarded_stem_target
271            .clamp(DEFAULT_GUARDED_STEM_TARGET, MAX_GUARDED_STEM_TARGET);
272        self.have_been_exhausted = false;
273        self.have_been_under_highwater = false;
274    }
275
276    /// Purge all the circuits from the pool.
277    #[allow(clippy::unnecessary_wraps)] // for consistency and future-proofing
278    pub(super) fn retire_all_circuits(&mut self) -> Result<(), tor_config::ReconfigureError> {
279        self.have_been_exhausted = true;
280
281        // Purge all circuits from this pool
282        self.circuits.clear();
283
284        Ok(())
285    }
286}
287
288/// Preferences for what kind of circuit to select from the pool.
289#[derive(Default, Debug, Clone)]
290pub(super) struct HsCircPrefs {
291    /// If `Some`, specifies the [`HsCircStemKind`] we would like.
292    pub(super) kind_prefs: Option<HsCircStemKind>,
293}
294
295impl HsCircPrefs {
296    /// Set the preferred [`HsCircStemKind`].
297    pub(super) fn preferred_stem_kind(&mut self, kind: HsCircStemKind) -> &mut Self {
298        self.kind_prefs = Some(kind);
299        self
300    }
301}
302
303/// Helper: find a random item `elt` in `slice` such that `predicate(elt)` is
304/// true. Return the index of that item.
305///
306/// Can arbitrarily reorder `slice`. This allows us to visit the indices in uniform-at-random
307/// order, without having to do any O(N) operations or allocations.
308fn random_idx_where<R, T, P>(rng: &mut R, mut slice: &mut [T], predicate: P) -> Option<usize>
309where
310    R: Rng,
311    P: Fn(&T) -> bool,
312{
313    while !slice.is_empty() {
314        let idx = rng
315            .gen_range_checked(0..slice.len())
316            .expect("slice was not empty but is now empty");
317        if predicate(&slice[idx]) {
318            return Some(idx);
319        }
320        let last_idx = slice.len() - 1;
321        // Move the one we just tried to the end,
322        // and eliminate it from consideration.
323        slice.swap(idx, last_idx);
324        slice = &mut slice[..last_idx];
325    }
326    // We didn't find any.
327    None
328}
329
330#[cfg(test)]
331mod test {
332    // @@ begin test lint list maintained by maint/add_warning @@
333    #![allow(clippy::bool_assert_comparison)]
334    #![allow(clippy::clone_on_copy)]
335    #![allow(clippy::dbg_macro)]
336    #![allow(clippy::mixed_attributes_style)]
337    #![allow(clippy::print_stderr)]
338    #![allow(clippy::print_stdout)]
339    #![allow(clippy::single_char_pattern)]
340    #![allow(clippy::unwrap_used)]
341    #![allow(clippy::unchecked_duration_subtraction)]
342    #![allow(clippy::useless_vec)]
343    #![allow(clippy::needless_pass_by_value)]
344    //! <!-- @@ end test lint list maintained by maint/add_warning @@ -->
345
346    use super::*;
347    use tor_basic_utils::test_rng::testing_rng;
348
349    #[test]
350    fn random_idx() {
351        let mut rng = testing_rng();
352        let mut orig_numbers: Vec<i32> = vec![1, 3, 4, 8, 11, 19, 12, 6, 27];
353        let mut numbers = orig_numbers.clone();
354
355        let mut found: std::collections::HashMap<i32, bool> =
356            numbers.iter().map(|n| (*n, false)).collect();
357
358        for _ in 0..1000 {
359            let idx = random_idx_where(&mut rng, &mut numbers[..], |n| n & 1 == 1).unwrap();
360            assert!(numbers[idx] & 1 == 1);
361            found.insert(numbers[idx], true);
362        }
363
364        for num in numbers.iter() {
365            assert!(found[num] == (num & 1 == 1));
366        }
367
368        // Number may be reordered, but should still have the same elements.
369        numbers.sort();
370        orig_numbers.sort();
371        assert_eq!(numbers, orig_numbers);
372    }
373
374    #[test]
375    fn random_idx_empty() {
376        let mut rng = testing_rng();
377        let idx = random_idx_where(&mut rng, &mut [], |_: &i32| panic!());
378        assert_eq!(idx, None);
379    }
380
381    #[test]
382    fn random_idx_none() {
383        let mut rng = testing_rng();
384        let mut numbers: Vec<i32> = vec![1, 3, 4, 8, 11, 19, 12, 6, 27];
385        assert_eq!(
386            random_idx_where(&mut rng, &mut numbers[..], |_: &i32| false),
387            None
388        );
389    }
390}
tor_circmgr/hspool/pool.rs

tor_circmgr/hspool/
pool.rs
