fosr_lib/
stage1.rs

1use crate::models;
2use crate::stats::Stats;
3use crate::structs::{SeededData, TimePoint};
4
5use chrono::FixedOffset;
6use chrono::{DateTime, Timelike};
7use crossbeam_channel::Sender;
8use rand_core::{Rng, SeedableRng};
9use rand_distr::Distribution;
10use rand_distr::Poisson;
11use rand_distr::Uniform;
12use rand_pcg::Pcg32;
13use serde::Deserialize;
14use std::sync::Arc;
15use std::thread;
16use std::time::{Duration, SystemTime, UNIX_EPOCH};
17
18const WINDOW_WIDTH_IN_SECS: u64 = 5;
19
20/// Stage 1: generate timestamps.
21/// Generate a throughput according to some distribution and never stops. It always prepares the next windows (i.e., not
22/// the one being sent)
23pub trait Stage1:
24    Iterator<Item = SeededData<TimePoint>> + Clone + std::marker::Send + 'static
25{
26    /// Used to compute the duration of generated data
27    fn get_initial_ts(&self) -> Duration;
28}
29
30/// The generator generates timestamp from bins
31/// In `net_injection` mode, it trickles the generation
32#[derive(Debug, Clone)]
33pub struct BinBasedGenerator {
34    next_ts: Duration, // the start of the *next* generation window
35    initial_ts: Duration,
36    remaining_flows: u64,
37    lambdas: Vec<f64>,
38    current_distrib: Option<Poisson<f64>>,
39    total_flow_count: u64,
40    time_distrib: Uniform<u64>,
41    window_rng: Pcg32,
42    flow_rng: Pcg32,
43    net_injection: bool,
44    remaining_windows: Option<u64>,
45    dest_tz_offset: FixedOffset,
46}
47
48impl Stage1 for BinBasedGenerator {
49    fn get_initial_ts(&self) -> Duration {
50        self.initial_ts
51    }
52}
53
54impl Iterator for BinBasedGenerator {
55    type Item = SeededData<TimePoint>;
56
57    fn next(&mut self) -> Option<Self::Item> {
58        while self.remaining_flows == 0 {
59            if self.net_injection {
60                // wait until the end of the current window
61                if SystemTime::now().duration_since(UNIX_EPOCH).unwrap() > self.next_ts {
62                    log::warn!("Generation is too slow");
63                } else {
64                    thread::sleep(
65                        self.next_ts
66                            .saturating_sub(SystemTime::now().duration_since(UNIX_EPOCH).unwrap()),
67                    );
68                }
69            }
70
71            let ok = self.start_new_window();
72            if !ok {
73                // log::info!("{} generated flows", self.total_flow_count);
74                return None;
75            }
76        }
77        self.remaining_flows -= 1;
78        self.total_flow_count += 1;
79        let unix_time = Duration::from_millis(self.time_distrib.sample(&mut self.flow_rng.clone()));
80        let date_time = DateTime::from_timestamp_secs(unix_time.as_secs() as i64)
81            .unwrap()
82            .with_timezone(&self.dest_tz_offset);
83        Some(SeededData {
84            seed: self.flow_rng.next_u64(),
85            data: TimePoint {
86                unix_time,
87                date_time,
88            },
89        })
90    }
91}
92
93fn get_poisson(lambdas: &[f64], dest_tz_offset: FixedOffset, ts: Duration) -> Option<Poisson<f64>> {
94    let secs = DateTime::from_timestamp_secs(ts.as_secs() as i64)
95        .unwrap()
96        .with_timezone(&dest_tz_offset)
97        .time()
98        .num_seconds_from_midnight();
99    // log::info!("Hours since midnight: {}", secs/3600);
100    let secs_per_bin = (60 * 60 * 24 / lambdas.len()) as u32;
101    let index = ((secs / secs_per_bin) as usize).min(lambdas.len() - 1);
102    // return None if lambda is 0
103    Poisson::new(lambdas[index]).ok()
104}
105
106/// Compute the parameters of the Poisson distribution from the bins
107/// If "`flow_per_day`" is None, then simply reuse the values of the bins
108/// Otherwise, normalize the bins so their sum is "`flow_per_day`"
109fn get_lambdas(flow_per_day: Option<u64>, bins: &[u64]) -> Vec<f64> {
110    let bin_count: f64 = bins.len() as f64;
111    let window_per_bin: f64 = 60. * 60. * 24. / (WINDOW_WIDTH_IN_SECS as f64) / bin_count;
112    match flow_per_day {
113        Some(flow_per_day) => {
114            let sum: f64 = bins.iter().sum::<u64>() as f64;
115            // Lambda is equal to the expected value of the Poisson distribution
116            // First, we normalize the bin so the sum of all bins is flow_per_day
117            // Then, we divide by the number of windows in one bin
118            bins.iter()
119                .map(|val| (*val as f64) / sum * (flow_per_day as f64) / window_per_bin)
120                .collect()
121        }
122        None => bins
123            .iter()
124            .map(|val| (*val as f64) / window_per_bin)
125            .collect(),
126    }
127}
128
129impl BinBasedGenerator {
130    pub fn new(
131        seed: Option<u64>,
132        net_injection: bool,
133        flow_per_day: Option<u64>,
134        profile: Arc<TimeModel>,
135        initial_ts: Duration,
136        total_duration: Option<Duration>,
137        dest_tz_offset: FixedOffset,
138    ) -> Self {
139        let next_ts = initial_ts;
140        let time_distrib = Uniform::new(
141            next_ts.as_millis() as u64,
142            next_ts.as_millis() as u64 + 1000 * WINDOW_WIDTH_IN_SECS,
143        )
144        .unwrap();
145        let window_rng = match seed {
146            Some(s) => Pcg32::seed_from_u64(s),
147            None => Pcg32::from_rng(&mut rand::rng()),
148        };
149        let remaining_windows = total_duration.map(|d| {
150            d.div_duration_f32(Duration::from_secs(WINDOW_WIDTH_IN_SECS))
151                .ceil() as u64
152        });
153
154        let lambdas = get_lambdas(flow_per_day, &profile.bins);
155
156        let mut generator = BinBasedGenerator {
157            net_injection,
158            initial_ts,
159            next_ts,
160            lambdas,
161            current_distrib: None, // it will be overwritten
162            remaining_flows: 0,
163            total_flow_count: 0,
164            window_rng,
165            flow_rng: Pcg32::seed_from_u64(0), // it will be overwritten
166            time_distrib,
167            remaining_windows,
168            dest_tz_offset,
169        };
170        generator.start_new_window();
171        generator
172    }
173
174    #[cfg(feature = "net_injection")]
175    pub fn new_for_injection(
176        seed: Option<u64>,
177        total_duration: Option<Duration>,
178        flow_per_day: Option<u64>,
179        profile: TimeModel,
180        deterministic: bool,
181    ) -> Self {
182        let current_date = SystemTime::now().duration_since(UNIX_EPOCH).unwrap();
183        let window_count_since_unix_epoch =
184            ((current_date + Duration::from_secs(WINDOW_WIDTH_IN_SECS)).as_secs_f64()
185                / (WINDOW_WIDTH_IN_SECS as f64))
186                .ceil() as u64;
187        let mut window_rng = match seed {
188            Some(s) => Pcg32::seed_from_u64(s),
189            None => Pcg32::new(0xcafef00dd15ea5e5, 0xa02bdbf7bb3c0a7), // default values from the doc
190        };
191        if !deterministic {
192            // each process should use the same rng at the same time !
193            // 2 for a next_u64 call
194            window_rng.advance(2 * window_count_since_unix_epoch);
195        }
196        let next_ts = Duration::from_secs(window_count_since_unix_epoch * WINDOW_WIDTH_IN_SECS);
197        let time_distrib = Uniform::new(
198            next_ts.as_millis() as u64,
199            next_ts.as_millis() as u64 + 1000 * WINDOW_WIDTH_IN_SECS,
200        )
201        .unwrap();
202        let remaining_windows = total_duration.map(|d| {
203            d.div_duration_f32(Duration::from_secs(WINDOW_WIDTH_IN_SECS))
204                .ceil() as u64
205        });
206        let lambdas = get_lambdas(flow_per_day, profile.bins);
207
208        let mut generator = BinBasedGenerator {
209            net_injection: true,
210            initial_ts: next_ts,
211            next_ts,
212            lambdas,
213            current_distrib: None, // it will be overwritten
214            remaining_flows: 0,
215            total_flow_count: 0,
216            window_rng,
217            flow_rng: Pcg32::seed_from_u64(0), // it will be overwritten
218            time_distrib,
219            remaining_windows,
220            dest_tz_offset: *chrono::Local::now().fixed_offset().offset(), // use local timezone with
221                                                                           // current offset
222        };
223        generator.start_new_window();
224        generator
225    }
226
227    /// Start a new window
228    /// Returns "false" it’s not possible
229    fn start_new_window(&mut self) -> bool {
230        self.flow_rng = Pcg32::seed_from_u64(self.window_rng.next_u64());
231        if let Some(ref mut r) = self.remaining_windows {
232            if *r == 0 {
233                return false;
234            }
235            *r -= 1;
236        }
237
238        self.current_distrib = get_poisson(&self.lambdas, self.dest_tz_offset, self.next_ts);
239
240        if self.current_distrib.is_none() {
241            log_once::warn_once!("Time window without any flow according to the time model");
242        }
243
244        self.remaining_flows = self
245            .current_distrib
246            .map_or(0, |s| s.sample(&mut self.flow_rng.clone()) as u64);
247        self.time_distrib = Uniform::new(
248            self.next_ts.as_millis() as u64,
249            self.next_ts.as_millis() as u64 + 1000 * WINDOW_WIDTH_IN_SECS,
250        )
251        .unwrap();
252        self.next_ts += Duration::from_secs(WINDOW_WIDTH_IN_SECS);
253        true
254    }
255}
256
257#[derive(Deserialize, Debug, Clone)]
258#[serde(from = "Config")]
259pub struct TimeModel {
260    bins: Vec<u64>, // we assume the bins start at midnight (in the local dataset timezone)
261    metadata: Metadata,
262}
263
264impl TimeModel {
265    pub fn from_source(m: &models::ModelsSource) -> Result<Self, String> {
266        let mut models: Vec<TimeModel> = vec![];
267        for s in m.get_time_profiles()? {
268            models.push(
269                serde_json::from_str(&s)
270                    .map_err(|e| format!("Cannot parse the time model: {e}"))?,
271            );
272        }
273        assert!(!models.is_empty());
274
275        let profile = Self::from_list(models);
276        log::info!(
277            "Time model learned on {} has been loaded",
278            profile.metadata.input_file
279        );
280        Ok(profile)
281    }
282
283    fn from_list(mut models: Vec<TimeModel>) -> Self {
284        let metadata = Metadata {
285            creation_time: "On-the-fly".to_string(),
286            input_file: models
287                .iter()
288                .map(|m| m.metadata.input_file.clone())
289                .collect::<Vec<String>>()
290                .join(", "),
291        };
292        // the overall number of packets of the new time model is the sum of the packets of each
293        // dataset, but each dataset has an equal weight
294        let overall_sum: u64 = models
295            .iter()
296            .map(|m: &TimeModel| -> u64 { m.bins.iter().sum() })
297            .sum();
298        let nb_models = models.len() as u64;
299        for m in &mut models {
300            let sum: u64 = m.bins.iter().sum();
301            m.bins = m
302                .bins
303                .clone()
304                .into_iter()
305                .map(|v| -> u64 { v * overall_sum / (sum * nb_models) })
306                .collect();
307        }
308        let mut bins: Vec<u64> = vec![];
309        assert!(!models.is_empty());
310        for i in 0..models[0].bins.len() {
311            bins.push(models.iter().map(|m| m.bins[i]).sum());
312        }
313        TimeModel { bins, metadata }
314    }
315}
316
317impl From<Config> for TimeModel {
318    fn from(c: Config) -> TimeModel {
319        TimeModel {
320            bins: c.histogram,
321            metadata: c.metadata,
322        }
323    }
324}
325
326#[derive(Deserialize, Debug, Clone)]
327struct Config {
328    histogram: Vec<u64>,
329    metadata: Metadata,
330}
331
332#[derive(Deserialize, Debug, Clone)]
333#[allow(unused)]
334struct Metadata {
335    creation_time: String,
336    input_file: String,
337}
338
339/// Generate data and send them progressively to a channel
340pub fn run_channel(
341    generator: impl Stage1,
342    tx_s1: Sender<SeededData<TimePoint>>,
343    stats: Arc<Stats>,
344) -> Result<(), Box<dyn std::error::Error>> {
345    log::trace!("Start S1 (channel)");
346    let initial_ts = generator.get_initial_ts();
347    for ts in generator {
348        if stats.should_stop() {
349            break;
350        }
351        stats.set_current_duration(
352            ts.data.unix_time.as_secs() - initial_ts.as_secs() + 2 * WINDOW_WIDTH_IN_SECS,
353        ); // this hack (adding WINDOW_WIDTH_IN_SECS) is just a way to be sure to reach the duration target for the progress bar
354        // log::trace!("S0 generates {ts:?}");
355        tx_s1.send(ts)?;
356    }
357    log::trace!("S1 stops");
358    Ok(())
359}
360
361/// Generate data into a vector
362pub fn run_vec(generator: impl Stage1) -> Vec<SeededData<TimePoint>> {
363    log::trace!("Start S1 (vec)");
364    let mut vector = vec![];
365    for ts in generator {
366        vector.push(ts);
367    }
368    log::trace!("S1 stops");
369    vector
370}
fosr_lib/stage1.rs

fosr_lib/
stage1.rs
