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

            
3
use crate::TimeoutStateHandle;
4
use crate::timeouts::{
5
    Action, TimeoutEstimator,
6
    pareto::{ParetoTimeoutEstimator, ParetoTimeoutState},
7
    readonly::ReadonlyTimeoutEstimator,
8
};
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
423
    pub(crate) fn from_storage(storage: &TimeoutStateHandle) -> Self {
34
423
        let (_, est) = estimator_from_storage(storage);
35
423
        Self {
36
423
            inner: Mutex::new(est),
37
423
        }
38
423
    }
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

            
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

            
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
411
    pub(crate) fn save_state(&self, storage: &TimeoutStateHandle) -> crate::Result<()> {
119
411
        let state = {
120
411
            let mut inner = self.inner.lock().expect("Timeout estimator lock poisoned.");
121
411
            inner.build_state()
122
        };
123
411
        if let Some(state) = state {
124
411
            storage.store(&state)?;
125
        }
126
411
        Ok(())
127
411
    }
128
}
129

            
130
impl tor_proto::client::circuit::TimeoutEstimator for Estimator {
131
    fn circuit_build_timeout(&self, length: usize) -> Duration {
132
        let (timeout, _abandon) = self
133
            .inner
134
            .lock()
135
            .expect("poisoned lock")
136
            .timeouts(&Action::BuildCircuit { length });
137

            
138
        timeout
139
    }
140
}
141

            
142
/// Try to construct a new boxed TimeoutEstimator based on the contents of
143
/// storage, and whether it is read-only.
144
///
145
/// Returns true on a read-only state.
146
425
fn estimator_from_storage(
147
425
    storage: &TimeoutStateHandle,
148
425
) -> (bool, Box<dyn TimeoutEstimator + Send + 'static>) {
149
425
    let state = match storage.load() {
150
4
        Ok(Some(v)) => v,
151
421
        Ok(None) => ParetoTimeoutState::default(),
152
        Err(e) => {
153
            warn_report!(e, "Unable to load timeout state");
154
            return (true, Box::new(ReadonlyTimeoutEstimator::new()));
155
        }
156
    };
157

            
158
425
    if storage.can_store() {
159
        // We own the lock, so we're going to use a full estimator.
160
263
        (false, Box::new(ParetoTimeoutEstimator::from_state(state)))
161
    } else {
162
162
        (true, Box::new(ReadonlyTimeoutEstimator::from_state(&state)))
163
    }
164
425
}
165

            
166
#[cfg(test)]
167
mod test {
168
    // @@ begin test lint list maintained by maint/add_warning @@
169
    #![allow(clippy::bool_assert_comparison)]
170
    #![allow(clippy::clone_on_copy)]
171
    #![allow(clippy::dbg_macro)]
172
    #![allow(clippy::mixed_attributes_style)]
173
    #![allow(clippy::print_stderr)]
174
    #![allow(clippy::print_stdout)]
175
    #![allow(clippy::single_char_pattern)]
176
    #![allow(clippy::unwrap_used)]
177
    #![allow(clippy::unchecked_duration_subtraction)]
178
    #![allow(clippy::useless_vec)]
179
    #![allow(clippy::needless_pass_by_value)]
180
    //! <!-- @@ end test lint list maintained by maint/add_warning @@ -->
181
    use super::*;
182
    use tor_persist::StateMgr;
183

            
184
    #[test]
185
    fn load_estimator() {
186
        let params = NetParameters::default();
187

            
188
        // Construct an estimator with write access to a state manager.
189
        let storage = tor_persist::TestingStateMgr::new();
190
        assert!(storage.try_lock().unwrap().held());
191
        let handle = storage.clone().create_handle("paretorama");
192

            
193
        let est = Estimator::from_storage(&handle);
194
        assert!(est.learning_timeouts());
195
        est.save_state(&handle).unwrap();
196

            
197
        // Construct another estimator that is looking at the same data,
198
        // but which only gets read-only access
199
        let storage2 = storage.new_manager();
200
        assert!(!storage2.try_lock().unwrap().held());
201
        let handle2 = storage2.clone().create_handle("paretorama");
202

            
203
        let est2 = Estimator::from_storage(&handle2);
204
        assert!(!est2.learning_timeouts());
205

            
206
        est.update_params(&params);
207
        est2.update_params(&params);
208

            
209
        // Initial timeouts, since no data is present yet.
210
        let act = Action::BuildCircuit { length: 3 };
211
        assert_eq!(
212
            est.timeouts(&act),
213
            (Duration::from_secs(60), Duration::from_secs(60))
214
        );
215
        assert_eq!(
216
            est2.timeouts(&act),
217
            (Duration::from_secs(60), Duration::from_secs(60))
218
        );
219

            
220
        // Pretend both estimators have gotten a bunch of observations...
221
        for _ in 0..500 {
222
            est.note_hop_completed(2, Duration::from_secs(7), true);
223
            est.note_hop_completed(2, Duration::from_secs(2), true);
224
            est2.note_hop_completed(2, Duration::from_secs(4), true);
225
        }
226
        assert!(!est.learning_timeouts());
227

            
228
        // Have est save and est2 load.
229
        est.save_state(&handle).unwrap();
230
        let to_1 = est.timeouts(&act);
231
        assert_ne!(
232
            est.timeouts(&act),
233
            (Duration::from_secs(60), Duration::from_secs(60))
234
        );
235
        assert_eq!(
236
            est2.timeouts(&act),
237
            (Duration::from_secs(60), Duration::from_secs(60))
238
        );
239
        est2.reload_readonly_from_storage(&handle2);
240
        let to_1_secs = to_1.0.as_secs_f64();
241
        let timeouts = est2.timeouts(&act);
242
        assert!((timeouts.0.as_secs_f64() - to_1_secs).abs() < 0.001);
243
        assert!((timeouts.1.as_secs_f64() - to_1_secs).abs() < 0.001);
244

            
245
        drop(est);
246
        drop(handle);
247
        drop(storage);
248

            
249
        // Now storage2 can upgrade...
250
        assert!(storage2.try_lock().unwrap().held());
251
        est2.upgrade_to_owning_storage(&handle2);
252
        let to_2 = est2.timeouts(&act);
253
        // This will be similar but not the same.
254
        assert!(to_2.0 > to_1.0 - Duration::from_secs(1));
255
        assert!(to_2.0 < to_1.0 + Duration::from_secs(1));
256
        // Make sure est2 is now mutable...
257
        for _ in 0..200 {
258
            est2.note_hop_completed(2, Duration::from_secs(1), true);
259
        }
260
        let to_3 = est2.timeouts(&act);
261
        assert!(to_3.0 < to_2.0);
262
    }
263
}