1use chrono::{DateTime, Duration, Utc};
2use std::time;
3
4const COALESCE_TIMEOUT: time::Duration = time::Duration::from_secs(20);
5const SIMULATED_DURATION_FOR_SINGLE_EVENT: time::Duration = time::Duration::from_millis(1);
6
7pub struct EventCoalescer {
8 environment: Option<&'static str>,
9 period_start: Option<DateTime<Utc>>,
10 period_end: Option<DateTime<Utc>>,
11}
12
13impl EventCoalescer {
14 pub fn new() -> Self {
15 Self {
16 environment: None,
17 period_start: None,
18 period_end: None,
19 }
20 }
21
22 pub fn log_event(
23 &mut self,
24 environment: &'static str,
25 ) -> Option<(DateTime<Utc>, DateTime<Utc>)> {
26 self.log_event_with_time(Utc::now(), environment)
27 }
28
29 // pub fn close_current_period(&mut self) -> Option<(DateTime<Utc>, DateTime<Utc>)> {
30 // self.environment.map(|env| self.log_event(env)).flatten()
31 // }
32
33 fn log_event_with_time(
34 &mut self,
35 log_time: DateTime<Utc>,
36 environment: &'static str,
37 ) -> Option<(DateTime<Utc>, DateTime<Utc>)> {
38 let coalesce_timeout = Duration::from_std(COALESCE_TIMEOUT).unwrap();
39
40 let Some(period_start) = self.period_start else {
41 self.period_start = Some(log_time);
42 self.environment = Some(environment);
43 return None;
44 };
45
46 let period_end = self
47 .period_end
48 .unwrap_or(period_start + SIMULATED_DURATION_FOR_SINGLE_EVENT);
49 let within_timeout = log_time - period_end < coalesce_timeout;
50 let environment_is_same = self.environment == Some(environment);
51 let should_coaelesce = !within_timeout || !environment_is_same;
52
53 if should_coaelesce {
54 self.period_start = Some(log_time);
55 self.period_end = None;
56 self.environment = Some(environment);
57 return Some((
58 period_start,
59 if within_timeout { log_time } else { period_end },
60 ));
61 }
62
63 self.period_end = Some(log_time);
64
65 None
66 }
67}
68
69#[cfg(test)]
70mod tests {
71 use chrono::TimeZone;
72
73 use super::*;
74
75 #[test]
76 fn test_same_context_exceeding_timeout() {
77 let environment_1 = "environment_1";
78 let mut event_coalescer = EventCoalescer::new();
79
80 assert_eq!(event_coalescer.period_start, None);
81 assert_eq!(event_coalescer.period_end, None);
82 assert_eq!(event_coalescer.environment, None);
83
84 let period_start = Utc.with_ymd_and_hms(1990, 4, 12, 0, 0, 0).unwrap();
85 let coalesced_duration = event_coalescer.log_event_with_time(period_start, environment_1);
86
87 assert_eq!(coalesced_duration, None);
88 assert_eq!(event_coalescer.period_start, Some(period_start));
89 assert_eq!(event_coalescer.period_end, None);
90 assert_eq!(event_coalescer.environment, Some(environment_1));
91
92 let within_timeout_adjustment = Duration::from_std(COALESCE_TIMEOUT / 2).unwrap();
93 let mut period_end = period_start;
94
95 // Ensure that many calls within the timeout don't start a new period
96 for _ in 0..100 {
97 period_end += within_timeout_adjustment;
98 let coalesced_duration = event_coalescer.log_event_with_time(period_end, environment_1);
99
100 assert_eq!(coalesced_duration, None);
101 assert_eq!(event_coalescer.period_start, Some(period_start));
102 assert_eq!(event_coalescer.period_end, Some(period_end));
103 assert_eq!(event_coalescer.environment, Some(environment_1));
104 }
105
106 let exceed_timeout_adjustment = Duration::from_std(COALESCE_TIMEOUT * 2).unwrap();
107 // Logging an event exceeding the timeout should start a new period
108 let new_period_start = period_end + exceed_timeout_adjustment;
109 let coalesced_duration =
110 event_coalescer.log_event_with_time(new_period_start, environment_1);
111
112 assert_eq!(coalesced_duration, Some((period_start, period_end)));
113 assert_eq!(event_coalescer.period_start, Some(new_period_start));
114 assert_eq!(event_coalescer.period_end, None);
115 assert_eq!(event_coalescer.environment, Some(environment_1));
116 }
117
118 #[test]
119 fn test_different_environment_under_timeout() {
120 let environment_1 = "environment_1";
121 let mut event_coalescer = EventCoalescer::new();
122
123 assert_eq!(event_coalescer.period_start, None);
124 assert_eq!(event_coalescer.period_end, None);
125 assert_eq!(event_coalescer.environment, None);
126
127 let period_start = Utc.with_ymd_and_hms(1990, 4, 12, 0, 0, 0).unwrap();
128 let coalesced_duration = event_coalescer.log_event_with_time(period_start, environment_1);
129
130 assert_eq!(coalesced_duration, None);
131 assert_eq!(event_coalescer.period_start, Some(period_start));
132 assert_eq!(event_coalescer.period_end, None);
133 assert_eq!(event_coalescer.environment, Some(environment_1));
134
135 let within_timeout_adjustment = Duration::from_std(COALESCE_TIMEOUT / 2).unwrap();
136 let period_end = period_start + within_timeout_adjustment;
137 let coalesced_duration = event_coalescer.log_event_with_time(period_end, environment_1);
138
139 assert_eq!(coalesced_duration, None);
140 assert_eq!(event_coalescer.period_start, Some(period_start));
141 assert_eq!(event_coalescer.period_end, Some(period_end));
142 assert_eq!(event_coalescer.environment, Some(environment_1));
143
144 // Logging an event within the timeout but with a different environment should start a new period
145 let period_end = period_end + within_timeout_adjustment;
146 let environment_2 = "environment_2";
147 let coalesced_duration = event_coalescer.log_event_with_time(period_end, environment_2);
148
149 assert_eq!(coalesced_duration, Some((period_start, period_end)));
150 assert_eq!(event_coalescer.period_start, Some(period_end));
151 assert_eq!(event_coalescer.period_end, None);
152 assert_eq!(event_coalescer.environment, Some(environment_2));
153 }
154
155 #[test]
156 fn test_switching_environment_while_within_timeout() {
157 let environment_1 = "environment_1";
158 let mut event_coalescer = EventCoalescer::new();
159
160 assert_eq!(event_coalescer.period_start, None);
161 assert_eq!(event_coalescer.period_end, None);
162 assert_eq!(event_coalescer.environment, None);
163
164 let period_start = Utc.with_ymd_and_hms(1990, 4, 12, 0, 0, 0).unwrap();
165 let coalesced_duration = event_coalescer.log_event_with_time(period_start, environment_1);
166
167 assert_eq!(coalesced_duration, None);
168 assert_eq!(event_coalescer.period_start, Some(period_start));
169 assert_eq!(event_coalescer.period_end, None);
170 assert_eq!(event_coalescer.environment, Some(environment_1));
171
172 let within_timeout_adjustment = Duration::from_std(COALESCE_TIMEOUT / 2).unwrap();
173 let period_end = period_start + within_timeout_adjustment;
174 let environment_2 = "environment_2";
175 let coalesced_duration = event_coalescer.log_event_with_time(period_end, environment_2);
176
177 assert_eq!(coalesced_duration, Some((period_start, period_end)));
178 assert_eq!(event_coalescer.period_start, Some(period_end));
179 assert_eq!(event_coalescer.period_end, None);
180 assert_eq!(event_coalescer.environment, Some(environment_2));
181 }
182 // 0 20 40 60
183 // |-------------------|-------------------|-------------------|-------------------
184 // |--------|----------env change
185 // |-------------------
186 // |period_start |period_end
187 // |new_period_start
188
189 #[test]
190 fn test_switching_environment_while_exceeding_timeout() {
191 let environment_1 = "environment_1";
192 let mut event_coalescer = EventCoalescer::new();
193
194 assert_eq!(event_coalescer.period_start, None);
195 assert_eq!(event_coalescer.period_end, None);
196 assert_eq!(event_coalescer.environment, None);
197
198 let period_start = Utc.with_ymd_and_hms(1990, 4, 12, 0, 0, 0).unwrap();
199 let coalesced_duration = event_coalescer.log_event_with_time(period_start, environment_1);
200
201 assert_eq!(coalesced_duration, None);
202 assert_eq!(event_coalescer.period_start, Some(period_start));
203 assert_eq!(event_coalescer.period_end, None);
204 assert_eq!(event_coalescer.environment, Some(environment_1));
205
206 let exceed_timeout_adjustment = Duration::from_std(COALESCE_TIMEOUT * 2).unwrap();
207 let period_end = period_start + exceed_timeout_adjustment;
208 let environment_2 = "environment_2";
209 let coalesced_duration = event_coalescer.log_event_with_time(period_end, environment_2);
210
211 assert_eq!(
212 coalesced_duration,
213 Some((
214 period_start,
215 period_start + SIMULATED_DURATION_FOR_SINGLE_EVENT
216 ))
217 );
218 assert_eq!(event_coalescer.period_start, Some(period_end));
219 assert_eq!(event_coalescer.period_end, None);
220 assert_eq!(event_coalescer.environment, Some(environment_2));
221 }
222 // 0 20 40 60
223 // |-------------------|-------------------|-------------------|-------------------
224 // |--------|----------------------------------------env change
225 // |-------------------|
226 // |period_start |period_end
227 // |new_period_start
228}