1use std::sync::Mutex;
45use std::sync::atomic::{AtomicU64, Ordering};
46use std::thread::ThreadId;
47
48#[derive(Debug, Clone, Default)]
50pub struct CombinerStats {
51 pub combine_passes: u64,
53 pub total_ops: u64,
55 pub max_batch_size: usize,
57 pub contention_events: u64,
59}
60
61impl CombinerStats {
62 pub fn avg_batch_size(&self) -> f64 {
64 if self.combine_passes == 0 {
65 0.0
66 } else {
67 self.total_ops as f64 / self.combine_passes as f64
68 }
69 }
70}
71
72pub struct FlatCombiner<S> {
82 state: Mutex<S>,
84 queue: Mutex<Vec<BoxedOp<S>>>,
86 generation: AtomicU64,
88 stats: Mutex<CombinerStats>,
90 combine_owner: Mutex<Option<ThreadId>>,
93}
94
95type BoxedOp<S> = Box<dyn FnOnce(&mut S) + Send>;
96
97struct CombineOwnerGuard<'a> {
98 owner: &'a Mutex<Option<ThreadId>>,
99}
100
101impl Drop for CombineOwnerGuard<'_> {
102 fn drop(&mut self) {
103 let mut owner = self.owner.lock().unwrap_or_else(|e| e.into_inner());
104 *owner = None;
105 }
106}
107
108impl<'a> CombineOwnerGuard<'a> {
109 fn new(owner: &'a Mutex<Option<ThreadId>>) -> Self {
110 let current = std::thread::current().id();
111 let mut owner_guard = owner.lock().unwrap_or_else(|e| e.into_inner());
112 *owner_guard = Some(current);
113 drop(owner_guard);
114 Self { owner }
115 }
116}
117
118impl<S> FlatCombiner<S> {
119 pub fn new(state: S) -> Self {
121 Self {
122 state: Mutex::new(state),
123 queue: Mutex::new(Vec::new()),
124 generation: AtomicU64::new(0),
125 stats: Mutex::new(CombinerStats::default()),
126 combine_owner: Mutex::new(None),
127 }
128 }
129
130 fn assert_not_reentrant(&self, operation: &str) {
131 let current = std::thread::current().id();
132 let owner = self.combine_owner.lock().unwrap_or_else(|e| e.into_inner());
133 if owner
134 .as_ref()
135 .is_some_and(|thread_id| *thread_id == current)
136 {
137 panic!("FlatCombiner::{operation} cannot be called reentrantly from a combine pass");
138 }
139 }
140
141 fn lock_queue(&self) -> std::sync::MutexGuard<'_, Vec<BoxedOp<S>>> {
144 match self.queue.try_lock() {
145 Ok(guard) => guard,
146 Err(std::sync::TryLockError::Poisoned(e)) => e.into_inner(),
147 Err(std::sync::TryLockError::WouldBlock) => {
148 if let Ok(mut stats) = self.stats.lock() {
149 stats.contention_events += 1;
150 }
151 self.queue.lock().unwrap_or_else(|e| e.into_inner())
152 }
153 }
154 }
155
156 pub fn execute<R>(&self, op: impl FnOnce(&mut S) -> R) -> R {
161 self.assert_not_reentrant("execute");
162 let mut state = self.state.lock().unwrap_or_else(|e| e.into_inner());
163 op(&mut state)
164 }
165
166 pub fn with_state<R>(&self, f: impl FnOnce(&S) -> R) -> R {
168 self.assert_not_reentrant("with_state");
169 let state = self.state.lock().unwrap_or_else(|e| e.into_inner());
170 f(&state)
171 }
172
173 pub fn submit(&self, op: impl FnOnce(&mut S) + Send + 'static) {
180 let mut queue = self.lock_queue();
181 queue.push(Box::new(op));
182 }
183
184 pub fn submit_batch(&self, ops: impl IntoIterator<Item = BoxedOp<S>>) {
186 let mut queue = self.lock_queue();
187 queue.extend(ops);
188 }
189
190 pub fn combine(&self) -> usize {
205 self.assert_not_reentrant("combine");
206 let ops: Vec<BoxedOp<S>> = {
208 let mut queue = self.queue.lock().unwrap_or_else(|e| e.into_inner());
209 std::mem::take(&mut *queue)
210 };
211
212 if ops.is_empty() {
213 return 0;
214 }
215
216 let count = ops.len();
217
218 {
220 let mut state = self.state.lock().unwrap_or_else(|e| e.into_inner());
221 let _owner_guard = CombineOwnerGuard::new(&self.combine_owner);
222 for op in ops {
223 op(&mut state);
224 }
225 }
226
227 self.generation.fetch_add(1, Ordering::Release);
229 if let Ok(mut stats) = self.stats.lock() {
230 stats.combine_passes += 1;
231 stats.total_ops += count as u64;
232 stats.max_batch_size = stats.max_batch_size.max(count);
233 }
234
235 count
236 }
237
238 pub fn combine_with<R>(&self, around: impl FnOnce(&mut S, &dyn Fn(&mut S)) -> R) -> (usize, R) {
255 self.assert_not_reentrant("combine_with");
256 let ops: Vec<BoxedOp<S>> = {
257 let mut queue = self.queue.lock().unwrap_or_else(|e| e.into_inner());
258 std::mem::take(&mut *queue)
259 };
260
261 let count = ops.len();
262 let mut state = self.state.lock().unwrap_or_else(|e| e.into_inner());
263 let _owner_guard = CombineOwnerGuard::new(&self.combine_owner);
264
265 let ops_cell = std::cell::RefCell::new(Some(ops));
268 let apply = |s: &mut S| {
269 if let Some(ops) = ops_cell.borrow_mut().take() {
270 for op in ops {
271 op(s);
272 }
273 }
274 };
275
276 let result = around(&mut state, &apply);
277 apply(&mut state);
281
282 if count > 0 {
283 self.generation.fetch_add(1, Ordering::Release);
284 if let Ok(mut stats) = self.stats.lock() {
285 stats.combine_passes += 1;
286 stats.total_ops += count as u64;
287 stats.max_batch_size = stats.max_batch_size.max(count);
288 }
289 }
290
291 (count, result)
292 }
293
294 pub fn pending_count(&self) -> usize {
296 self.queue.lock().unwrap_or_else(|e| e.into_inner()).len()
297 }
298
299 pub fn generation(&self) -> u64 {
301 self.generation.load(Ordering::Acquire)
302 }
303
304 pub fn stats(&self) -> CombinerStats {
306 self.stats.lock().unwrap_or_else(|e| e.into_inner()).clone()
307 }
308
309 pub fn reset_stats(&self) {
311 if let Ok(mut stats) = self.stats.lock() {
312 *stats = CombinerStats::default();
313 }
314 }
315}
316
317impl<S: std::fmt::Debug> std::fmt::Debug for FlatCombiner<S> {
321 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
322 let pending = self.pending_count();
323 let current_gen = self.generation();
324 f.debug_struct("FlatCombiner")
325 .field("pending", &pending)
326 .field("generation", ¤t_gen)
327 .finish_non_exhaustive()
328 }
329}
330
331#[cfg(test)]
332mod tests {
333 use super::*;
334 use std::sync::Arc;
335
336 #[test]
337 fn new_creates_empty_combiner() {
338 let fc = FlatCombiner::new(0u64);
339 assert_eq!(fc.pending_count(), 0);
340 assert_eq!(fc.generation(), 0);
341 assert_eq!(fc.stats().combine_passes, 0);
342 }
343
344 #[test]
345 fn execute_applies_directly() {
346 let fc = FlatCombiner::new(10u64);
347 let result = fc.execute(|s| {
348 *s += 5;
349 *s
350 });
351 assert_eq!(result, 15);
352 }
353
354 #[test]
355 fn with_state_reads_without_mutation() {
356 let fc = FlatCombiner::new(vec![1, 2, 3]);
357 let len = fc.with_state(|s| s.len());
358 assert_eq!(len, 3);
359 }
360
361 #[test]
362 fn submit_queues_operations() {
363 let fc = FlatCombiner::new(0u64);
364 fc.submit(|s| *s += 1);
365 fc.submit(|s| *s += 2);
366 assert_eq!(fc.pending_count(), 2);
367
368 let val = fc.with_state(|s| *s);
370 assert_eq!(val, 0);
371 }
372
373 #[test]
374 fn combine_drains_and_applies() {
375 let fc = FlatCombiner::new(0u64);
376 fc.submit(|s| *s += 10);
377 fc.submit(|s| *s += 20);
378 fc.submit(|s| *s += 30);
379
380 let count = fc.combine();
381 assert_eq!(count, 3);
382 assert_eq!(fc.pending_count(), 0);
383
384 let val = fc.with_state(|s| *s);
385 assert_eq!(val, 60);
386 }
387
388 #[test]
389 fn combine_empty_returns_zero() {
390 let fc = FlatCombiner::new(0u64);
391 assert_eq!(fc.combine(), 0);
392 assert_eq!(fc.generation(), 0);
393 }
394
395 #[test]
396 fn combine_increments_generation() {
397 let fc = FlatCombiner::new(0u64);
398 assert_eq!(fc.generation(), 0);
399
400 fc.submit(|s| *s += 1);
401 fc.combine();
402 assert_eq!(fc.generation(), 1);
403
404 fc.submit(|s| *s += 1);
405 fc.combine();
406 assert_eq!(fc.generation(), 2);
407 }
408
409 #[test]
410 fn stats_track_batches() {
411 let fc = FlatCombiner::new(0u64);
412
413 fc.submit(|s| *s += 1);
415 fc.submit(|s| *s += 1);
416 fc.submit(|s| *s += 1);
417 fc.combine();
418
419 fc.submit(|s| *s += 1);
421 fc.combine();
422
423 let stats = fc.stats();
424 assert_eq!(stats.combine_passes, 2);
425 assert_eq!(stats.total_ops, 4);
426 assert_eq!(stats.max_batch_size, 3);
427 assert!((stats.avg_batch_size() - 2.0).abs() < f64::EPSILON);
428 }
429
430 #[test]
431 fn reset_stats_clears_counters() {
432 let fc = FlatCombiner::new(0u64);
433 fc.submit(|s| *s += 1);
434 fc.combine();
435 assert_eq!(fc.stats().combine_passes, 1);
436
437 fc.reset_stats();
438 let stats = fc.stats();
439 assert_eq!(stats.combine_passes, 0);
440 assert_eq!(stats.total_ops, 0);
441 }
442
443 #[test]
444 fn operations_execute_in_order() {
445 let fc = FlatCombiner::new(Vec::<u32>::new());
446 fc.submit(|s| s.push(1));
447 fc.submit(|s| s.push(2));
448 fc.submit(|s| s.push(3));
449 fc.combine();
450
451 let values = fc.with_state(|s| s.clone());
452 assert_eq!(values, vec![1, 2, 3]);
453 }
454
455 #[test]
456 fn submit_batch_adds_multiple() {
457 let fc = FlatCombiner::new(0u64);
458 let ops: Vec<BoxedOp<u64>> = vec![
459 Box::new(|s: &mut u64| *s += 10),
460 Box::new(|s: &mut u64| *s += 20),
461 ];
462 fc.submit_batch(ops);
463 assert_eq!(fc.pending_count(), 2);
464 fc.combine();
465 assert_eq!(fc.with_state(|s| *s), 30);
466 }
467
468 #[test]
469 fn combine_with_wraps_batch() {
470 let fc = FlatCombiner::new(Vec::<String>::new());
471 fc.submit(|s| s.push("a".into()));
472 fc.submit(|s| s.push("b".into()));
473
474 let (count, len_before) = fc.combine_with(|state, apply| {
475 let before = state.len();
476 apply(state);
477 before
478 });
479
480 assert_eq!(count, 2);
481 assert_eq!(len_before, 0);
482 assert_eq!(fc.with_state(|s| s.len()), 2);
483 }
484
485 #[test]
486 fn multiple_combine_passes() {
487 let fc = FlatCombiner::new(0u64);
488
489 for i in 0..10 {
490 fc.submit(move |s| *s += i);
491 }
492 fc.combine();
493 assert_eq!(fc.with_state(|s| *s), 45); for i in 0..5 {
496 fc.submit(move |s| *s += i);
497 }
498 fc.combine();
499 assert_eq!(fc.with_state(|s| *s), 55); }
501
502 #[test]
503 fn debug_impl() {
504 let fc = FlatCombiner::new(42u64);
505 let debug = format!("{fc:?}");
506 assert!(debug.contains("FlatCombiner"));
507 assert!(debug.contains("pending"));
508 assert!(debug.contains("generation"));
509 }
510
511 #[test]
512 fn concurrent_submit_and_combine() {
513 let fc = Arc::new(FlatCombiner::new(0u64));
514
515 let handles: Vec<_> = (0..8)
517 .map(|_| {
518 let fc = Arc::clone(&fc);
519 std::thread::spawn(move || {
520 for _ in 0..100 {
521 fc.submit(|s| *s += 1);
522 }
523 })
524 })
525 .collect();
526
527 for h in handles {
529 h.join().unwrap();
530 }
531
532 let mut total = 0;
534 loop {
535 let count = fc.combine();
536 if count == 0 {
537 break;
538 }
539 total += count;
540 }
541
542 assert_eq!(total, 800);
543 assert_eq!(fc.with_state(|s| *s), 800);
544 }
545
546 #[test]
547 fn concurrent_submit_and_combine_interleaved() {
548 let fc = Arc::new(FlatCombiner::new(0u64));
549
550 let submit_handles: Vec<_> = (0..4)
552 .map(|_| {
553 let fc = Arc::clone(&fc);
554 std::thread::spawn(move || {
555 for _ in 0..100 {
556 fc.submit(|s| *s += 1);
557 std::thread::yield_now();
558 }
559 })
560 })
561 .collect();
562
563 let fc_c = Arc::clone(&fc);
565 let combiner = std::thread::spawn(move || {
566 let mut total = 0;
567 for _ in 0..500 {
568 total += fc_c.combine();
569 std::thread::yield_now();
570 }
571 total
572 });
573
574 for h in submit_handles {
575 h.join().unwrap();
576 }
577
578 let combined_during = combiner.join().unwrap();
580 let remaining = fc.combine();
581 let final_val = fc.with_state(|s| *s);
582
583 assert_eq!(
584 final_val,
585 (combined_during + remaining) as u64,
586 "total combined ({} + {}) should match state ({})",
587 combined_during,
588 remaining,
589 final_val
590 );
591 assert_eq!(final_val, 400);
592 }
593
594 #[test]
595 fn poison_recovery() {
596 let fc = FlatCombiner::new(0u64);
600 fc.submit(|_| panic!("op panics"));
601 let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| fc.combine()));
602 assert!(result.is_err(), "panic must propagate out of combine");
603 assert_eq!(fc.generation(), 0, "aborted pass must not bump generation");
604
605 fc.execute(|s| *s += 1);
606 fc.submit(|s| *s += 1);
607 assert_eq!(fc.combine(), 1);
608 assert_eq!(fc.with_state(|s| *s), 2);
609 assert_eq!(fc.generation(), 1);
610 }
611
612 #[test]
613 fn combine_panics_on_reentrant_call_from_op_instead_of_deadlocking() {
614 let fc = Arc::new(FlatCombiner::new(0u64));
615 let fc2 = Arc::clone(&fc);
616 fc.submit(move |_| {
617 let _ = fc2.with_state(|s| *s);
618 });
619
620 let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| fc.combine()));
621 assert!(result.is_err(), "reentrant with_state must panic, not hang");
622 }
623
624 #[test]
625 fn submit_from_inside_combine_op_lands_in_next_batch() {
626 let fc = Arc::new(FlatCombiner::new(0u64));
627 let fc2 = Arc::clone(&fc);
628 fc.submit(move |s| {
629 *s += 1;
630 fc2.submit(|s| *s += 10);
631 });
632
633 assert_eq!(fc.combine(), 1);
634 assert_eq!(fc.with_state(|s| *s), 1);
635 assert_eq!(fc.pending_count(), 1);
636 assert_eq!(fc.combine(), 1);
637 assert_eq!(fc.with_state(|s| *s), 11);
638 }
639
640 #[test]
641 fn combine_with_executes_batch_even_if_apply_not_called() {
642 let fc = FlatCombiner::new(0u64);
643 fc.submit(|s| *s += 5);
644
645 let (count, ()) = fc.combine_with(|_, _apply| ());
646
647 assert_eq!(count, 1);
648 assert_eq!(
649 fc.with_state(|s| *s),
650 5,
651 "batch must be applied even when `around` skips `apply`"
652 );
653 assert_eq!(fc.generation(), 1);
654 }
655
656 #[test]
657 fn contention_events_recorded_when_queue_is_held() {
658 let fc = Arc::new(FlatCombiner::new(0u64));
659
660 let queue_guard = fc.queue.lock().unwrap();
663 let fc2 = Arc::clone(&fc);
664 let submitter = std::thread::spawn(move || fc2.submit(|s| *s += 1));
665
666 while fc.stats().contention_events == 0 {
669 std::thread::yield_now();
670 }
671 drop(queue_guard);
672 submitter.join().unwrap();
673
674 assert!(fc.stats().contention_events >= 1);
675 assert_eq!(fc.combine(), 1);
676 assert_eq!(fc.with_state(|s| *s), 1);
677 }
678
679 #[test]
680 fn avg_batch_size_zero_when_no_combines() {
681 let stats = CombinerStats::default();
682 assert_eq!(stats.avg_batch_size(), 0.0);
683 }
684
685 #[test]
686 fn combine_with_panics_on_reentrant_execute_instead_of_deadlocking() {
687 let fc = FlatCombiner::new(0u64);
688
689 let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
690 let _ = fc.combine_with(|_, _| fc.execute(|state| *state));
691 }));
692
693 assert!(result.is_err());
694 }
695
696 #[test]
697 fn combine_with_panics_on_reentrant_with_state_instead_of_deadlocking() {
698 let fc = FlatCombiner::new(7u64);
699
700 let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
701 let _ = fc.combine_with(|_, _| fc.with_state(|state| *state));
702 }));
703
704 assert!(result.is_err());
705 }
706}