1
//! Declarations for a [`TimeoutEstimator`] type that can change implementation.
2

            
3
use crate::timeouts::{
4
    pareto::{ParetoTimeoutEstimator, ParetoTimeoutState},
5
    readonly::ReadonlyTimeoutEstimator,
6
    Action, TimeoutEstimator,
7
};
8
use crate::TimeoutStateHandle;
9
use std::sync::Mutex;
10
use std::time::Duration;
11
use tor_error::warn_report;
12
use tor_netdir::params::NetParameters;
13
use tracing::{debug, warn};
14

            
15
/// A timeout estimator that can change its inner implementation and share its
16
/// implementation among multiple threads.
17
pub(crate) struct Estimator {
18
    /// The estimator we're currently using.
19
    inner: Mutex<Box<dyn TimeoutEstimator + Send + 'static>>,
20
}
21

            
22
impl Estimator {
23
    /// Construct a new estimator from some variant.
24
    #[cfg(test)]
25
16
    pub(crate) fn new(est: impl TimeoutEstimator + Send + 'static) -> Self {
26
16
        Self {
27
16
            inner: Mutex::new(Box::new(est)),
28
16
        }
29
16
    }
30

            
31
    /// Create this estimator based on the values stored in `storage`, and whether
32
    /// this storage is read-only.
33
233
    pub(crate) fn from_storage(storage: &TimeoutStateHandle) -> Self {
34
233
        let (_, est) = estimator_from_storage(storage);
35
233
        Self {
36
233
            inner: Mutex::new(est),
37
233
        }
38
233
    }
39

            
40
    /// Assuming that we can read and write to `storage`, replace our state with
41
    /// a new state that estimates timeouts.
42
2
    pub(crate) fn upgrade_to_owning_storage(&self, storage: &TimeoutStateHandle) {
43
2
        let (readonly, est) = estimator_from_storage(storage);
44
2
        if readonly {
45
            warn!("Unable to upgrade to owned persistent storage.");
46
            return;
47
2
        }
48
2
        *self.inner.lock().expect("Timeout estimator lock poisoned") = est;
49
2
    }
50

            
51
    /// Replace the contents of this estimator with a read-only state estimator
52
    /// based on the contents of `storage`.
53
2
    pub(crate) fn reload_readonly_from_storage(&self, storage: &TimeoutStateHandle) {
54
2
        if let Ok(Some(v)) = storage.load() {
55
2
            let est = ReadonlyTimeoutEstimator::from_state(&v);
56
2
            *self.inner.lock().expect("Timeout estimator lock poisoned") = Box::new(est);
57
2
        } else {
58
            debug!("Unable to reload timeout state.");
59
        }
60
2
    }
61

            
62
    /// Record that a given circuit hop has completed.
63
    ///
64
    /// The `hop` number is a zero-indexed value for which hop just completed.
65
    ///
66
    /// The `delay` value is the amount of time after we first launched the
67
    /// circuit.
68
    ///
69
    /// If this is the last hop of the circuit, then `is_last` is true.
70
3436
    pub(crate) fn note_hop_completed(&self, hop: u8, delay: Duration, is_last: bool) {
71
3436
        let mut inner = self.inner.lock().expect("Timeout estimator lock poisoned.");
72
3436

            
73
3436
        inner.note_hop_completed(hop, delay, is_last);
74
3436
    }
75

            
76
    /// Record that a circuit failed to complete because it took too long.
77
    ///
78
    /// The `hop` number is a the number of hops that were successfully
79
    /// completed.
80
    ///
81
    /// The `delay` number is the amount of time after we first launched the
82
    /// circuit.
83
8
    pub(crate) fn note_circ_timeout(&self, hop: u8, delay: Duration) {
84
8
        let mut inner = self.inner.lock().expect("Timeout estimator lock poisoned.");
85
8
        inner.note_circ_timeout(hop, delay);
86
8
    }
87

            
88
    /// Return the current estimation for how long we should wait for a given
89
    /// [`Action`] to complete.
90
    ///
91
    /// This function should return a 2-tuple of `(timeout, abandon)`
92
    /// durations.  After `timeout` has elapsed since circuit launch,
93
    /// the circuit should no longer be used, but we should still keep
94
    /// building it in order see how long it takes.  After `abandon`
95
    /// has elapsed since circuit launch, the circuit should be
96
    /// abandoned completely.
97
32
    pub(crate) fn timeouts(&self, action: &Action) -> (Duration, Duration) {
98
32
        let mut inner = self.inner.lock().expect("Timeout estimator lock poisoned.");
99
32

            
100
32
        inner.timeouts(action)
101
32
    }
102

            
103
    /// Return true if we're currently trying to learn more timeouts
104
    /// by launching testing circuits.
105
6
    pub(crate) fn learning_timeouts(&self) -> bool {
106
6
        let inner = self.inner.lock().expect("Timeout estimator lock poisoned.");
107
6
        inner.learning_timeouts()
108
6
    }
109

            
110
    /// Replace the network parameters used by this estimator (if any)
111
    /// with ones derived from `params`.
112
4
    pub(crate) fn update_params(&self, params: &NetParameters) {
113
4
        let mut inner = self.inner.lock().expect("Timeout estimator lock poisoned.");
114
4
        inner.update_params(params);
115
4
    }
116

            
117
    /// Store any state associated with this timeout estimator into `storage`.
118
97
    pub(crate) fn save_state(&self, storage: &TimeoutStateHandle) -> crate::Result<()> {
119
97
        let state = {
120
97
            let mut inner = self.inner.lock().expect("Timeout estimator lock poisoned.");
121
97
            inner.build_state()
122
        };
123
97
        if let Some(state) = state {
124
97
            storage.store(&state)?;
125
        }
126
97
        Ok(())
127
97
    }
128
}
129

            
130
/// Try to construct a new boxed TimeoutEstimator based on the contents of
131
/// storage, and whether it is read-only.
132
///
133
/// Returns true on a read-only state.
134
235
fn estimator_from_storage(
135
235
    storage: &TimeoutStateHandle,
136
235
) -> (bool, Box<dyn TimeoutEstimator + Send + 'static>) {
137
235
    let state = match storage.load() {
138
4
        Ok(Some(v)) => v,
139
231
        Ok(None) => ParetoTimeoutState::default(),
140
        Err(e) => {
141
            warn_report!(e, "Unable to load timeout state");
142
            return (true, Box::new(ReadonlyTimeoutEstimator::new()));
143
        }
144
    };
145

            
146
235
    if storage.can_store() {
147
        // We own the lock, so we're going to use a full estimator.
148
97
        (false, Box::new(ParetoTimeoutEstimator::from_state(state)))
149
    } else {
150
138
        (true, Box::new(ReadonlyTimeoutEstimator::from_state(&state)))
151
    }
152
235
}
153

            
154
#[cfg(test)]
155
mod test {
156
    // @@ begin test lint list maintained by maint/add_warning @@
157
    #![allow(clippy::bool_assert_comparison)]
158
    #![allow(clippy::clone_on_copy)]
159
    #![allow(clippy::dbg_macro)]
160
    #![allow(clippy::mixed_attributes_style)]
161
    #![allow(clippy::print_stderr)]
162
    #![allow(clippy::print_stdout)]
163
    #![allow(clippy::single_char_pattern)]
164
    #![allow(clippy::unwrap_used)]
165
    #![allow(clippy::unchecked_duration_subtraction)]
166
    #![allow(clippy::useless_vec)]
167
    #![allow(clippy::needless_pass_by_value)]
168
    //! <!-- @@ end test lint list maintained by maint/add_warning @@ -->
169
    use super::*;
170
    use tor_persist::StateMgr;
171

            
172
    #[test]
173
    fn load_estimator() {
174
        let params = NetParameters::default();
175

            
176
        // Construct an estimator with write access to a state manager.
177
        let storage = tor_persist::TestingStateMgr::new();
178
        assert!(storage.try_lock().unwrap().held());
179
        let handle = storage.clone().create_handle("paretorama");
180

            
181
        let est = Estimator::from_storage(&handle);
182
        assert!(est.learning_timeouts());
183
        est.save_state(&handle).unwrap();
184

            
185
        // Construct another estimator that is looking at the same data,
186
        // but which only gets read-only access
187
        let storage2 = storage.new_manager();
188
        assert!(!storage2.try_lock().unwrap().held());
189
        let handle2 = storage2.clone().create_handle("paretorama");
190

            
191
        let est2 = Estimator::from_storage(&handle2);
192
        assert!(!est2.learning_timeouts());
193

            
194
        est.update_params(&params);
195
        est2.update_params(&params);
196

            
197
        // Initial timeouts, since no data is present yet.
198
        let act = Action::BuildCircuit { length: 3 };
199
        assert_eq!(
200
            est.timeouts(&act),
201
            (Duration::from_secs(60), Duration::from_secs(60))
202
        );
203
        assert_eq!(
204
            est2.timeouts(&act),
205
            (Duration::from_secs(60), Duration::from_secs(60))
206
        );
207

            
208
        // Pretend both estimators have gotten a bunch of observations...
209
        for _ in 0..500 {
210
            est.note_hop_completed(2, Duration::from_secs(7), true);
211
            est.note_hop_completed(2, Duration::from_secs(2), true);
212
            est2.note_hop_completed(2, Duration::from_secs(4), true);
213
        }
214
        assert!(!est.learning_timeouts());
215

            
216
        // Have est save and est2 load.
217
        est.save_state(&handle).unwrap();
218
        let to_1 = est.timeouts(&act);
219
        assert_ne!(
220
            est.timeouts(&act),
221
            (Duration::from_secs(60), Duration::from_secs(60))
222
        );
223
        assert_eq!(
224
            est2.timeouts(&act),
225
            (Duration::from_secs(60), Duration::from_secs(60))
226
        );
227
        est2.reload_readonly_from_storage(&handle2);
228
        let to_1_secs = to_1.0.as_secs_f64();
229
        let timeouts = est2.timeouts(&act);
230
        assert!((timeouts.0.as_secs_f64() - to_1_secs).abs() < 0.001);
231
        assert!((timeouts.1.as_secs_f64() - to_1_secs).abs() < 0.001);
232

            
233
        drop(est);
234
        drop(handle);
235
        drop(storage);
236

            
237
        // Now storage2 can upgrade...
238
        assert!(storage2.try_lock().unwrap().held());
239
        est2.upgrade_to_owning_storage(&handle2);
240
        let to_2 = est2.timeouts(&act);
241
        // This will be similar but not the same.
242
        assert!(to_2.0 > to_1.0 - Duration::from_secs(1));
243
        assert!(to_2.0 < to_1.0 + Duration::from_secs(1));
244
        // Make sure est2 is now mutable...
245
        for _ in 0..200 {
246
            est2.note_hop_completed(2, Duration::from_secs(1), true);
247
        }
248
        let to_3 = est2.timeouts(&act);
249
        assert!(to_3.0 < to_2.0);
250
    }
251
}