tor_circmgr/
preemptive.rs
1use crate::{PathConfig, PreemptiveCircuitConfig, TargetCircUsage, TargetPort};
4use std::collections::HashMap;
5use std::sync::Arc;
6use std::time::Instant;
7use tracing::warn;
8
9pub(crate) struct PreemptiveCircuitPredictor {
12 usages: HashMap<Option<TargetPort>, Instant>,
17
18 config: tor_config::MutCfg<PreemptiveCircuitConfig>,
20}
21
22impl PreemptiveCircuitPredictor {
23 pub(crate) fn new(config: PreemptiveCircuitConfig) -> Self {
25 let mut usages = HashMap::new();
26 for port in &config.initial_predicted_ports {
27 usages.insert(Some(TargetPort::ipv4(*port)), Instant::now());
29 }
30
31 usages.insert(None, Instant::now());
33
34 Self {
35 usages,
36 config: config.into(),
37 }
38 }
39
40 pub(crate) fn config(&self) -> Arc<PreemptiveCircuitConfig> {
42 self.config.get()
43 }
44
45 pub(crate) fn set_config(&self, mut new_config: PreemptiveCircuitConfig) {
48 self.config.map_and_replace(|cfg| {
49 new_config
51 .initial_predicted_ports
52 .clone_from(&cfg.initial_predicted_ports);
53 new_config
54 });
55 }
56
57 pub(crate) fn predict(&self, path_config: &PathConfig) -> Vec<TargetCircUsage> {
59 let config = self.config();
60 let now = Instant::now();
61 let circs = config.min_exit_circs_for_port;
62 self.usages
63 .iter()
64 .filter(|(_, &time)| {
65 time.checked_add(config.prediction_lifetime)
66 .map(|t| t > now)
67 .unwrap_or_else(|| {
68 warn!("failed to represent preemptive circuit prediction lifetime as an Instant");
71 false
72 })
73 })
74 .map(|(&port, _)| {
75 let require_stability = port.is_some_and(|p| path_config.long_lived_ports.contains(&p.port));
76 TargetCircUsage::Preemptive {
77 port, circs, require_stability,
78 }
79 })
80 .collect()
81 }
82
83 pub(crate) fn note_usage(&mut self, port: Option<TargetPort>, time: Instant) {
90 self.usages.insert(port, time);
91 }
92}
93
94#[cfg(test)]
95mod test {
96 #![allow(clippy::bool_assert_comparison)]
98 #![allow(clippy::clone_on_copy)]
99 #![allow(clippy::dbg_macro)]
100 #![allow(clippy::mixed_attributes_style)]
101 #![allow(clippy::print_stderr)]
102 #![allow(clippy::print_stdout)]
103 #![allow(clippy::single_char_pattern)]
104 #![allow(clippy::unwrap_used)]
105 #![allow(clippy::unchecked_duration_subtraction)]
106 #![allow(clippy::useless_vec)]
107 #![allow(clippy::needless_pass_by_value)]
108 use crate::{
110 PathConfig, PreemptiveCircuitConfig, PreemptiveCircuitPredictor, TargetCircUsage,
111 TargetPort,
112 };
113 use std::time::{Duration, Instant};
114
115 use crate::isolation::test::{assert_isoleq, IsolationTokenEq};
116
117 #[test]
118 fn predicts_starting_ports() {
119 let path_config = PathConfig::default();
120 let mut cfg = PreemptiveCircuitConfig::builder();
121 cfg.set_initial_predicted_ports(vec![]);
122 cfg.prediction_lifetime(Duration::from_secs(2));
123 let predictor = PreemptiveCircuitPredictor::new(cfg.build().unwrap());
124
125 assert_isoleq!(
126 predictor.predict(&path_config),
127 vec![TargetCircUsage::Preemptive {
128 port: None,
129 circs: 2,
130 require_stability: false,
131 }]
132 );
133
134 let mut cfg = PreemptiveCircuitConfig::builder();
135 cfg.set_initial_predicted_ports(vec![80]);
136 cfg.prediction_lifetime(Duration::from_secs(2));
137 let predictor = PreemptiveCircuitPredictor::new(cfg.build().unwrap());
138
139 let results = predictor.predict(&path_config);
140 assert_eq!(results.len(), 2);
141 assert!(results
142 .iter()
143 .any(|r| r.isol_eq(&TargetCircUsage::Preemptive {
144 port: None,
145 circs: 2,
146 require_stability: false,
147 })));
148 assert!(results
149 .iter()
150 .any(|r| r.isol_eq(&TargetCircUsage::Preemptive {
151 port: Some(TargetPort::ipv4(80)),
152 circs: 2,
153 require_stability: false,
154 })));
155 }
156
157 #[test]
158 fn predicts_used_ports() {
159 let path_config = PathConfig::default();
160 let mut cfg = PreemptiveCircuitConfig::builder();
161 cfg.set_initial_predicted_ports(vec![]);
162 cfg.prediction_lifetime(Duration::from_secs(2));
163 let mut predictor = PreemptiveCircuitPredictor::new(cfg.build().unwrap());
164
165 assert_isoleq!(
166 predictor.predict(&path_config),
167 vec![TargetCircUsage::Preemptive {
168 port: None,
169 circs: 2,
170 require_stability: false,
171 }]
172 );
173
174 predictor.note_usage(Some(TargetPort::ipv4(1234)), Instant::now());
175
176 let results = predictor.predict(&path_config);
177 assert_eq!(results.len(), 2);
178 assert!(results
179 .iter()
180 .any(|r| r.isol_eq(&TargetCircUsage::Preemptive {
181 port: None,
182 circs: 2,
183 require_stability: false,
184 })));
185 assert!(results
186 .iter()
187 .any(|r| r.isol_eq(&TargetCircUsage::Preemptive {
188 port: Some(TargetPort::ipv4(1234)),
189 circs: 2,
190 require_stability: false,
191 })));
192 }
193
194 #[test]
195 fn does_not_predict_old_ports() {
196 let path_config = PathConfig::default();
197 let mut cfg = PreemptiveCircuitConfig::builder();
198 cfg.set_initial_predicted_ports(vec![]);
199 cfg.prediction_lifetime(Duration::from_secs(2));
200 let mut predictor = PreemptiveCircuitPredictor::new(cfg.build().unwrap());
201 let now = Instant::now();
202 let three_seconds_ago = now - Duration::from_secs(2 + 1);
203
204 predictor.note_usage(Some(TargetPort::ipv4(2345)), three_seconds_ago);
205
206 assert_isoleq!(
207 predictor.predict(&path_config),
208 vec![TargetCircUsage::Preemptive {
209 port: None,
210 circs: 2,
211 require_stability: false,
212 }]
213 );
214 }
215}