statsig_rust/event_logging/event_queue/
queue.rs1use super::{batch::EventBatch, queued_event::QueuedEvent};
2use crate::{
3 event_logging::statsig_event_internal::StatsigEventInternal, log_d, read_lock_or_return,
4 write_lock_or_return,
5};
6use parking_lot::RwLock;
7use std::collections::VecDeque;
8
9const TAG: &str = stringify!(EventQueue);
10
11pub enum QueueAddResult {
12 Noop,
13 NeedsFlush,
14 NeedsFlushAndDropped(u64),
15}
16
17pub enum QueueReconcileResult {
18 Success,
19 DroppedEvents(u64),
20 LockFailure,
21}
22
23pub struct EventQueue {
24 pub batch_size: usize,
25 pub max_pending_batches: usize,
26
27 pending_events: RwLock<VecDeque<QueuedEvent>>,
28 batches: RwLock<VecDeque<EventBatch>>,
29 max_pending_events: usize,
30}
31
32impl EventQueue {
33 pub fn new(batch_size: u32, max_queue_size: u32) -> Self {
34 let batch_size = batch_size as usize;
35 let max_queue_size = max_queue_size as usize;
36
37 Self {
38 pending_events: RwLock::new(VecDeque::new()),
39 batches: RwLock::new(VecDeque::new()),
40 batch_size,
41 max_pending_batches: max_queue_size,
42 max_pending_events: batch_size * max_queue_size,
43 }
44 }
45
46 pub fn approximate_pending_events_count(&self) -> usize {
47 let pending_len = read_lock_or_return!(TAG, self.pending_events, 0).len();
48 let batches_len = read_lock_or_return!(TAG, self.batches, 0).len();
49 pending_len + (batches_len * self.batch_size)
50 }
51
52 pub fn add(&self, pending_event: QueuedEvent) -> QueueAddResult {
53 let mut pending_events =
54 write_lock_or_return!(TAG, self.pending_events, QueueAddResult::Noop);
55 pending_events.push_back(pending_event);
56
57 let mut dropped_events = 0;
58 while pending_events.len() > self.max_pending_events {
59 pending_events.pop_front();
60 dropped_events += 1;
61 }
62
63 if dropped_events > 0 {
64 return QueueAddResult::NeedsFlushAndDropped(dropped_events);
65 }
66
67 if pending_events.len() % self.batch_size == 0 {
68 return QueueAddResult::NeedsFlush;
69 }
70
71 QueueAddResult::Noop
72 }
73
74 pub fn requeue_batch(&self, batch: EventBatch) -> QueueReconcileResult {
75 let len = batch.events.len() as u64;
76 let mut batches =
77 write_lock_or_return!(TAG, self.batches, QueueReconcileResult::DroppedEvents(len));
78
79 if batches.len() > self.max_pending_batches {
80 return QueueReconcileResult::DroppedEvents(len);
81 }
82
83 log_d!(
84 TAG,
85 "Requeueing batch with {} events and {} attempts to flush",
86 batch.events.len(),
87 batch.attempts
88 );
89
90 batches.push_back(batch);
91 QueueReconcileResult::Success
92 }
93
94 pub fn contains_at_least_one_full_batch(&self) -> bool {
95 let pending_events_count = self
96 .pending_events
97 .try_read_for(std::time::Duration::from_secs(5))
98 .map(|e| e.len())
99 .unwrap_or(0);
100 if pending_events_count >= self.batch_size {
101 return true;
102 }
103
104 let batches = read_lock_or_return!(TAG, self.batches, false);
105 for batch in batches.iter() {
106 if batch.events.len() >= self.batch_size {
107 return true;
108 }
109 }
110
111 false
112 }
113
114 pub fn take_all_batches(&self) -> VecDeque<EventBatch> {
115 let mut batches = write_lock_or_return!(TAG, self.batches, VecDeque::new());
116 std::mem::take(&mut *batches)
117 }
118
119 pub fn take_next_batch(&self) -> Option<EventBatch> {
120 let mut batches = write_lock_or_return!(TAG, self.batches, None);
121 batches.pop_front()
122 }
123
124 pub fn reconcile_batching(&self) -> QueueReconcileResult {
125 let pending_events: VecDeque<StatsigEventInternal> = self
126 .take_all_pending_events()
127 .into_iter()
128 .map(|evt| evt.into_statsig_event_internal())
129 .collect();
130
131 if pending_events.is_empty() {
132 return QueueReconcileResult::Success;
133 }
134
135 let mut batches =
136 write_lock_or_return!(TAG, self.batches, QueueReconcileResult::LockFailure);
137 let old_batches = std::mem::take(&mut *batches);
138
139 let (full_batches, partial_batches): (VecDeque<_>, VecDeque<_>) = old_batches
140 .into_iter()
141 .partition(|batch| batch.events.len() >= self.batch_size);
142
143 let mut events_to_batch = VecDeque::new();
144 for batch in partial_batches {
145 events_to_batch.extend(batch.events);
146 }
147 events_to_batch.extend(pending_events);
148
149 let new_batches = self.create_batches(events_to_batch);
150
151 batches.extend(full_batches);
152 batches.extend(new_batches);
153
154 let dropped_events_count = self.clamp_batches(&mut batches);
155 if dropped_events_count > 0 {
156 return QueueReconcileResult::DroppedEvents(dropped_events_count);
157 }
158
159 QueueReconcileResult::Success
160 }
161
162 fn take_all_pending_events(&self) -> VecDeque<QueuedEvent> {
163 let mut pending_events = write_lock_or_return!(TAG, self.pending_events, VecDeque::new());
164 std::mem::take(&mut *pending_events)
165 }
166
167 fn create_batches(
168 &self,
169 mut pending_events: VecDeque<StatsigEventInternal>,
170 ) -> Vec<EventBatch> {
171 let mut batches = Vec::new();
172 while !pending_events.is_empty() {
173 let drain_count = self.batch_size.min(pending_events.len());
174 let chunk = pending_events.drain(..drain_count).collect::<Vec<_>>();
175 batches.push(EventBatch::new(chunk));
176 }
177
178 batches
179 }
180
181 fn clamp_batches(&self, batches: &mut VecDeque<EventBatch>) -> u64 {
182 if batches.len() <= self.max_pending_batches {
183 return 0;
184 }
185
186 let mut dropped_events_count = 0;
187 while batches.len() > self.max_pending_batches {
188 if let Some(batch) = batches.pop_front() {
189 dropped_events_count += batch.events.len() as u64;
190 }
191 }
192
193 dropped_events_count
194 }
195}
196
197#[cfg(test)]
198mod tests {
199 use std::borrow::Cow;
200
201 use super::*;
202 use crate::event_logging::event_queue::queued_event::EnqueueOperation;
203 use crate::event_logging::event_queue::queued_gate_expo::EnqueueGateExpoOp;
204 use crate::event_logging::exposure_sampling::EvtSamplingDecision::ForceSampled;
205 use crate::{
206 event_logging::event_logger::ExposureTrigger, statsig_types::FeatureGate,
207 user::StatsigUserInternal, EvaluationDetails, StatsigUser,
208 };
209
210 #[test]
211 fn test_adding_single_to_queue() {
212 let (queue, user, gate) = setup(10, 10);
213 let user_internal = StatsigUserInternal::new(&user, None);
214
215 let enqueue_op = EnqueueGateExpoOp {
216 exposure_time: 1,
217 user: &user_internal,
218 queried_gate_name: &gate.name,
219 evaluation: gate.__evaluation.as_ref().map(Cow::Borrowed),
220 details: EvaluationDetails::unrecognized_no_data(),
221 trigger: ExposureTrigger::Auto,
222 };
223
224 let queued_event = enqueue_op.into_queued_event(ForceSampled);
225
226 let result = queue.add(queued_event);
227
228 assert!(matches!(result, QueueAddResult::Noop));
229 assert_eq!(
230 queue
231 .pending_events
232 .try_read_for(std::time::Duration::from_secs(5))
233 .unwrap()
234 .len(),
235 1
236 );
237 }
238
239 #[test]
240 fn test_adding_multiple_to_queue() {
241 let (queue, user, gate) = setup(1000, 20);
242 let user_internal = StatsigUserInternal::new(&user, None);
243
244 let mut triggered_count = 0;
245 for _ in 0..4567 {
246 let enqueue_op = EnqueueGateExpoOp {
247 exposure_time: 1,
248 user: &user_internal,
249 queried_gate_name: &gate.name,
250 evaluation: gate.__evaluation.as_ref().map(Cow::Borrowed),
251 details: EvaluationDetails::unrecognized_no_data(),
252 trigger: ExposureTrigger::Auto,
253 };
254
255 let result = queue.add(enqueue_op.into_queued_event(ForceSampled));
256
257 if let QueueAddResult::NeedsFlush = result {
258 triggered_count += 1;
259 }
260 }
261
262 assert_eq!(
263 queue
264 .pending_events
265 .try_read_for(std::time::Duration::from_secs(5))
266 .unwrap()
267 .len(),
268 4567
269 );
270 assert_eq!(triggered_count, (4567 / 1000) as usize);
271 }
272
273 #[test]
274 fn test_take_all_batches() {
275 let batch_size = 200;
276 let max_pending_batches = 40;
277
278 let (queue, user, gate) = setup(batch_size, max_pending_batches);
279 let user_internal = StatsigUserInternal::new(&user, None);
280
281 for _ in 0..4567 {
282 let enqueue_op = EnqueueGateExpoOp {
283 exposure_time: 1,
284 user: &user_internal,
285 queried_gate_name: &gate.name,
286 evaluation: gate.__evaluation.as_ref().map(Cow::Borrowed),
287 details: EvaluationDetails::unrecognized_no_data(),
288 trigger: ExposureTrigger::Auto,
289 };
290 queue.add(enqueue_op.into_queued_event(ForceSampled));
291 }
292
293 queue.reconcile_batching();
294 let batches = queue.take_all_batches();
295 assert_eq!(batches.len(), (4567.0 / batch_size as f64).ceil() as usize,);
296 }
297
298 #[test]
299 fn test_take_next_batch() {
300 let batch_size = 200;
301 let max_pending_batches = 20;
302
303 let (queue, user, gate) = setup(batch_size, max_pending_batches);
304 let user_internal = StatsigUserInternal::new(&user, None);
305
306 for _ in 0..4567 {
307 let enqueue_op = EnqueueGateExpoOp {
308 exposure_time: 1,
309 user: &user_internal,
310 queried_gate_name: &gate.name,
311 evaluation: gate.__evaluation.as_ref().map(Cow::Borrowed),
312 details: EvaluationDetails::unrecognized_no_data(),
313 trigger: ExposureTrigger::Auto,
314 };
315 queue.add(enqueue_op.into_queued_event(ForceSampled));
316 }
317
318 queue.reconcile_batching();
319 let batch = queue.take_next_batch();
320 assert_eq!(batch.unwrap().events.len(), batch_size as usize);
321
322 assert_eq!(
323 queue
324 .batches
325 .try_read_for(std::time::Duration::from_secs(5))
326 .unwrap()
327 .len(),
328 (max_pending_batches - 1) as usize
329 ); }
331
332 fn setup(batch_size: u32, max_queue_size: u32) -> (EventQueue, StatsigUser, FeatureGate) {
333 let queue = EventQueue::new(batch_size, max_queue_size);
334 let user = StatsigUser::with_user_id("user-id");
335 let gate = FeatureGate {
336 name: "gate-name".into(),
337 value: true,
338 rule_id: "rule-id".into(),
339 id_type: "user-id".into(),
340 details: EvaluationDetails::unrecognized_no_data(),
341 __evaluation: None,
342 };
343 (queue, user, gate)
344 }
345}