1use crate::budget::Budget;
4use std::collections::VecDeque;
5use std::sync::Arc;
6use std::sync::atomic::{AtomicBool, Ordering};
7use yo_common::Result;
8use yo_shard::Epochs;
9use yo_shard::spsc::Receiver;
10
11pub const BATCH_MAX: usize = 64;
18
19#[derive(Debug, Clone, Copy, PartialEq, Eq)]
21pub enum Flow {
22 Next,
24 Break,
32}
33
34pub trait Engine {
42 type Work;
44
45 fn key_hash(&self, work: &Self::Work) -> Option<u64>;
51
52 fn prefetch(&self, work: &Self::Work, hash: u64);
63
64 fn run(&mut self, work: Self::Work, hash: Option<u64>) -> Flow;
69
70 fn flush(&mut self);
76
77 fn submit_io(&mut self) -> Result<()> {
87 Ok(())
88 }
89
90 fn drain_io(&mut self) -> Result<()> {
100 Ok(())
101 }
102
103 fn maintain(&mut self, budget: &mut Budget) {
109 let _ = budget;
110 }
111}
112
113#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
115pub struct Turn {
116 pub commands: usize,
118 pub broke: bool,
120 pub carried: usize,
122 pub maintained: u32,
124}
125
126impl Turn {
127 #[must_use]
133 pub const fn is_idle(&self) -> bool {
134 self.commands == 0 && self.carried == 0
135 }
136}
137
138pub struct Reactor<E: Engine> {
143 engine: E,
144 id: usize,
145 epochs: Arc<Epochs>,
146 lanes: Vec<Receiver<E::Work>>,
147 pending: VecDeque<E::Work>,
150 hashes: Vec<Option<u64>>,
152 lane: usize,
155 budget: u32,
156 turns: u64,
157 commands: u64,
158 batches: u64,
159 full: u64,
160 breaks: u64,
161 idle: u64,
162}
163
164impl<E: Engine> Reactor<E> {
165 pub fn new(engine: E, id: usize, epochs: Arc<Epochs>, lanes: Vec<Receiver<E::Work>>) -> Self {
176 assert!(id < epochs.len(), "shard {id} has no epoch slot");
177 Reactor {
178 engine,
179 id,
180 epochs,
181 lanes,
182 pending: VecDeque::with_capacity(BATCH_MAX),
183 hashes: Vec::with_capacity(BATCH_MAX),
184 lane: 0,
185 budget: crate::MAINTENANCE_UNITS,
186 turns: 0,
187 commands: 0,
188 batches: 0,
189 full: 0,
190 breaks: 0,
191 idle: 0,
192 }
193 }
194
195 pub fn inline(engine: E) -> Self {
203 Reactor::new(engine, 0, Epochs::new(1), Vec::new())
204 }
205
206 #[must_use]
211 pub fn with_maintenance(mut self, units: u32) -> Self {
212 self.budget = units;
213 self
214 }
215
216 pub const fn engine(&self) -> &E {
218 &self.engine
219 }
220
221 pub const fn engine_mut(&mut self) -> &mut E {
223 &mut self.engine
224 }
225
226 #[must_use]
228 pub const fn id(&self) -> usize {
229 self.id
230 }
231
232 #[must_use]
234 pub const fn turns(&self) -> u64 {
235 self.turns
236 }
237
238 #[must_use]
240 pub const fn commands(&self) -> u64 {
241 self.commands
242 }
243
244 #[must_use]
246 pub const fn batches(&self) -> u64 {
247 self.batches
248 }
249
250 #[must_use]
257 pub const fn full_batches(&self) -> u64 {
258 self.full
259 }
260
261 #[must_use]
264 pub const fn breaks(&self) -> u64 {
265 self.breaks
266 }
267
268 #[must_use]
270 pub const fn idle_turns(&self) -> u64 {
271 self.idle
272 }
273
274 #[must_use]
277 pub fn carried(&self) -> usize {
278 self.pending.len()
279 }
280
281 pub fn tick(&mut self) -> Result<Turn> {
289 self.turns += 1;
290
291 self.engine.submit_io()?;
293
294 if self.pending.is_empty() {
298 self.fill();
299 if !self.pending.is_empty() {
300 self.batches += 1;
301 if self.pending.len() == BATCH_MAX {
302 self.full += 1;
303 }
304 }
305 }
306
307 let mut turn = Turn::default();
308 if self.pending.is_empty() {
309 self.idle += 1;
315 } else {
316 self.epochs.enter(self.id);
318
319 self.hashes.clear();
321 for w in &self.pending {
322 let h = self.engine.key_hash(w);
323 if let Some(h) = h {
324 self.engine.prefetch(w, h);
325 }
326 self.hashes.push(h);
327 }
328
329 let mut n = 0;
331 while let Some(w) = self.pending.pop_front() {
332 let h = self.hashes[n];
333 n += 1;
334 if self.engine.run(w, h) == Flow::Break {
335 turn.broke = true;
336 self.breaks += 1;
337 break;
338 }
339 }
340
341 self.epochs.leave(self.id);
343 self.engine.flush();
344
345 self.commands += n as u64;
346 turn.commands = n;
347 turn.carried = self.pending.len();
348 }
349
350 self.engine.drain_io()?;
352 if self.budget > 0 {
353 let mut budget = Budget::new(self.budget);
354 self.engine.maintain(&mut budget);
355 turn.maintained = budget.spent();
356 }
357 Ok(turn)
358 }
359
360 pub fn run_until(&mut self, stop: &AtomicBool) -> Result<()> {
371 const SPINS: u32 = 128;
372 let mut idle = 0u32;
373 loop {
374 if !self.tick()?.is_idle() {
375 idle = 0;
376 continue;
377 }
378 if stop.load(Ordering::Acquire) {
379 if self.tick()?.is_idle() {
381 return Ok(());
382 }
383 continue;
384 }
385 idle += 1;
386 if idle < SPINS {
387 std::hint::spin_loop();
388 } else {
389 std::thread::yield_now();
390 idle = 0;
391 }
392 }
393 }
394
395 pub fn execute(&mut self, work: E::Work) -> Flow {
406 self.turns += 1;
407 self.commands += 1;
408 self.epochs.enter(self.id);
409 let hash = self.engine.key_hash(&work);
410 if let Some(h) = hash {
411 self.engine.prefetch(&work, h);
412 }
413 let flow = self.engine.run(work, hash);
414 self.epochs.leave(self.id);
415 flow
416 }
417
418 pub fn execute_all<I>(&mut self, work: I) -> usize
428 where
429 I: IntoIterator<Item = E::Work>,
430 {
431 self.pending.extend(work);
432 if self.pending.is_empty() {
433 return 0;
434 }
435 self.turns += 1;
436 self.epochs.enter(self.id);
437
438 self.hashes.clear();
439 for w in &self.pending {
440 let h = self.engine.key_hash(w);
441 if let Some(h) = h {
442 self.engine.prefetch(w, h);
443 }
444 self.hashes.push(h);
445 }
446
447 let mut n = 0;
448 while let Some(w) = self.pending.pop_front() {
449 let h = self.hashes[n];
450 n += 1;
451 if self.engine.run(w, h) == Flow::Break {
452 self.breaks += 1;
453 break;
454 }
455 }
456 self.pending.clear();
457
458 self.epochs.leave(self.id);
459 self.commands += n as u64;
460 n
461 }
462
463 fn fill(&mut self) {
466 if self.lanes.is_empty() {
467 return;
468 }
469 let n = self.lanes.len();
470 let mut room = BATCH_MAX;
471 let mut at = self.lane;
472 loop {
473 let mut took = 0;
474 for _ in 0..n {
475 if room == 0 {
476 self.lane = at;
477 return;
478 }
479 if let Some(w) = self.lanes[at].pop() {
480 self.pending.push_back(w);
481 room -= 1;
482 took += 1;
483 }
484 at = (at + 1) % n;
485 }
486 if took == 0 {
487 self.lane = at;
488 return;
489 }
490 }
491 }
492}
493
494impl<E: Engine> std::fmt::Debug for Reactor<E> {
495 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
496 f.debug_struct("Reactor")
497 .field("id", &self.id)
498 .field("lanes", &self.lanes.len())
499 .field("turns", &self.turns)
500 .field("commands", &self.commands)
501 .field("batches", &self.batches)
502 .field("full_batches", &self.full)
503 .field("breaks", &self.breaks)
504 .field("idle_turns", &self.idle)
505 .field("carried", &self.pending.len())
506 .finish()
507 }
508}
509
510#[cfg(test)]
511mod tests {
512 use super::*;
513 use std::cell::RefCell;
514 use yo_common::{Code, Error};
515 use yo_shard::spsc::{Sender, lane};
516
517 #[derive(Debug, Clone, PartialEq, Eq)]
519 enum Step {
520 Submit,
521 Prefetch(u64),
522 Run(u64),
523 Flush,
524 Drain,
525 Maintain,
526 }
527
528 const KEYLESS: u64 = u64::MAX;
531
532 struct Recorder {
539 steps: RefCell<Vec<Step>>,
540 break_on: Option<u64>,
541 fail_submit: bool,
542 fail_drain: bool,
543 maintenance_item: u32,
544 }
545
546 impl Recorder {
547 fn new() -> Recorder {
548 Recorder {
549 steps: RefCell::new(Vec::new()),
550 break_on: None,
551 fail_submit: false,
552 fail_drain: false,
553 maintenance_item: 0,
554 }
555 }
556
557 fn push(&self, step: Step) {
558 self.steps.borrow_mut().push(step);
559 }
560
561 fn steps(&self) -> Vec<Step> {
562 self.steps.borrow().clone()
563 }
564
565 fn runs(&self) -> Vec<u64> {
566 self.steps
567 .borrow()
568 .iter()
569 .filter_map(|s| match s {
570 Step::Run(v) => Some(*v),
571 _ => None,
572 })
573 .collect()
574 }
575
576 fn count(&self, want: &Step) -> usize {
577 self.steps.borrow().iter().filter(|s| *s == want).count()
578 }
579 }
580
581 impl Engine for Recorder {
582 type Work = u64;
583
584 fn key_hash(&self, work: &u64) -> Option<u64> {
585 if *work == KEYLESS { None } else { Some(*work) }
586 }
587
588 fn prefetch(&self, _work: &u64, hash: u64) {
589 self.push(Step::Prefetch(hash));
590 }
591
592 fn run(&mut self, work: u64, hash: Option<u64>) -> Flow {
593 assert_eq!(
594 hash,
595 if work == KEYLESS { None } else { Some(work) },
596 "the second walk gets the hash the first walk took"
597 );
598 self.push(Step::Run(work));
599 if self.break_on == Some(work) {
600 return Flow::Break;
601 }
602 Flow::Next
603 }
604
605 fn flush(&mut self) {
606 self.push(Step::Flush);
607 }
608
609 fn submit_io(&mut self) -> Result<()> {
610 self.push(Step::Submit);
611 if self.fail_submit {
612 return Err(Error::new(Code::Io, "submit said no"));
613 }
614 Ok(())
615 }
616
617 fn drain_io(&mut self) -> Result<()> {
618 self.push(Step::Drain);
619 if self.fail_drain {
620 return Err(Error::new(Code::Io, "drain said no"));
621 }
622 Ok(())
623 }
624
625 fn maintain(&mut self, budget: &mut Budget) {
626 self.push(Step::Maintain);
627 if self.maintenance_item == 0 {
628 return;
629 }
630 while budget.spend(self.maintenance_item) {}
631 }
632 }
633
634 fn wired(lanes: usize) -> (Reactor<Recorder>, Vec<Sender<u64>>, Arc<Epochs>) {
637 let mut rxs = Vec::new();
638 let mut txs = Vec::new();
639 for _ in 0..lanes {
640 let (tx, rx) = lane(1024);
641 txs.push(tx);
642 rxs.push(rx);
643 }
644 let epochs = Epochs::new(1);
645 let r = Reactor::new(Recorder::new(), 0, Arc::clone(&epochs), rxs);
646 (r, txs, epochs)
647 }
648
649 fn walk_boundary(steps: &[Step]) -> (usize, usize) {
652 let last_prefetch = steps
653 .iter()
654 .rposition(|s| matches!(s, Step::Prefetch(_)))
655 .expect("nothing was prefetched");
656 let first_run = steps
657 .iter()
658 .position(|s| matches!(s, Step::Run(_)))
659 .expect("nothing ran");
660 (last_prefetch, first_run)
661 }
662
663 #[test]
664 fn the_stages_run_in_the_order_the_spec_lists_them() {
665 let (mut r, tx, _e) = wired(1);
666 tx[0].push(7).unwrap();
667 let turn = r.tick().unwrap();
668 assert_eq!(turn.commands, 1);
669 assert_eq!(
670 r.engine().steps(),
671 vec![
672 Step::Submit,
673 Step::Prefetch(7),
674 Step::Run(7),
675 Step::Flush,
676 Step::Drain,
677 Step::Maintain,
678 ]
679 );
680 }
681
682 #[test]
683 fn every_command_is_prefetched_before_any_of_them_runs() {
684 let (mut r, tx, _e) = wired(1);
685 for i in 0..8 {
686 tx[0].push(i).unwrap();
687 }
688 r.tick().unwrap();
689 let (last_prefetch, first_run) = walk_boundary(&r.engine().steps());
690 assert!(
691 last_prefetch < first_run,
692 "the two walks overlapped, which makes the prefetch distance one"
693 );
694 assert_eq!(r.engine().runs(), (0..8).collect::<Vec<_>>());
695 }
696
697 #[test]
698 fn a_batch_stops_at_sixty_four() {
699 let (mut r, tx, _e) = wired(1);
700 for i in 0..200 {
701 tx[0].push(i).unwrap();
702 }
703 assert_eq!(r.tick().unwrap().commands, BATCH_MAX);
704 assert_eq!(r.full_batches(), 1);
705 assert_eq!(r.tick().unwrap().commands, BATCH_MAX);
706 assert_eq!(r.tick().unwrap().commands, BATCH_MAX);
707 assert_eq!(r.tick().unwrap().commands, 200 - 3 * BATCH_MAX);
708 assert_eq!(r.batches(), 4);
709 assert_eq!(r.full_batches(), 3);
710 assert_eq!(r.commands(), 200);
711 assert_eq!(r.engine().runs(), (0..200).collect::<Vec<_>>());
712 }
713
714 #[test]
715 fn a_break_leaves_the_rest_of_the_batch_for_the_next_turn() {
716 let (mut r, tx, _e) = wired(1);
717 for i in 0..10 {
718 tx[0].push(i).unwrap();
719 }
720 r.engine_mut().break_on = Some(3);
721 let turn = r.tick().unwrap();
722 assert!(turn.broke);
723 assert_eq!(turn.commands, 4, "the command that broke it still ran");
724 assert_eq!(turn.carried, 6);
725 assert_eq!(r.engine().count(&Step::Flush), 1, "the replies still went");
726
727 r.engine_mut().break_on = None;
730 let turn = r.tick().unwrap();
731 assert_eq!(turn.commands, 6);
732 assert_eq!(turn.carried, 0);
733 assert_eq!(r.engine().runs(), (0..10).collect::<Vec<_>>());
734 assert_eq!(r.breaks(), 1);
735 assert_eq!(r.batches(), 1, "a broken batch is one batch, not two");
736 }
737
738 #[test]
739 fn the_epoch_moves_once_per_batch_and_not_once_per_command() {
740 let (mut r, tx, epochs) = wired(1);
741 for i in 0..10 {
742 tx[0].push(i).unwrap();
743 }
744 let before = epochs.get(0);
745 r.tick().unwrap();
746 assert_eq!(
747 epochs.get(0),
748 before + 2,
749 "one enter and one leave for ten commands"
750 );
751 assert_eq!(r.commands(), 10);
752 }
753
754 #[test]
755 fn an_idle_turn_does_not_touch_the_epoch_or_the_replies() {
756 let (mut r, _tx, epochs) = wired(1);
757 let before = epochs.get(0);
758 let turn = r.tick().unwrap();
759 assert!(turn.is_idle());
760 assert_eq!(epochs.get(0), before, "an idle shard holds nothing");
761 assert_eq!(r.engine().count(&Step::Flush), 0);
762 assert_eq!(
763 r.engine().count(&Step::Drain),
764 1,
765 "completions still get picked up"
766 );
767 assert_eq!(
768 r.engine().count(&Step::Maintain),
769 1,
770 "and background work still runs"
771 );
772 assert_eq!(r.idle_turns(), 1);
773 assert_eq!(r.batches(), 0);
774 }
775
776 #[test]
777 fn work_comes_off_every_lane_rather_than_the_first_one() {
778 let (mut r, tx, _e) = wired(4);
779 for (i, t) in tx.iter().enumerate() {
780 for j in 0..4u64 {
781 t.push(i as u64 * 10 + j).unwrap();
782 }
783 }
784 let turn = r.tick().unwrap();
785 assert_eq!(turn.commands, 16);
786 let runs = r.engine().runs();
787 for lane in 0..4u64 {
788 let from_lane = runs.iter().filter(|v| **v / 10 == lane).count();
789 assert_eq!(from_lane, 4, "lane {lane} was skipped or drained twice");
790 }
791 assert_eq!(
792 runs[..4].iter().map(|v| v / 10).collect::<Vec<_>>(),
793 vec![0, 1, 2, 3],
794 "round robin, so the first four are one from each lane"
795 );
796 }
797
798 #[test]
799 fn a_lane_that_never_stops_cannot_starve_the_others() {
800 let (mut r, tx, _e) = wired(2);
801 for i in 0..200 {
804 tx[0].push(i).unwrap();
805 }
806 tx[1].push(9_999).unwrap();
807 r.tick().unwrap();
808 assert!(
809 r.engine().runs().contains(&9_999),
810 "the quiet lane waited behind a whole batch of the busy one"
811 );
812 }
813
814 #[test]
815 fn maintenance_gets_a_budget_and_stops_when_it_is_spent() {
816 let (mut r, _tx, _e) = wired(1);
817 r.engine_mut().maintenance_item = 100;
818 let turn = r.tick().unwrap();
819 assert!(
822 (crate::MAINTENANCE_UNITS..crate::MAINTENANCE_UNITS + 100).contains(&turn.maintained),
823 "spent {}",
824 turn.maintained
825 );
826
827 let mut r = r.with_maintenance(0);
828 let before = r.engine().count(&Step::Maintain);
829 let turn = r.tick().unwrap();
830 assert_eq!(turn.maintained, 0);
831 assert_eq!(
832 r.engine().count(&Step::Maintain),
833 before,
834 "a zero budget is no slice at all rather than an empty one"
835 );
836 }
837
838 #[test]
839 fn a_failed_submit_stops_the_turn_before_the_batch() {
840 let (mut r, tx, _e) = wired(1);
841 tx[0].push(1).unwrap();
842 r.engine_mut().fail_submit = true;
843 assert_eq!(r.tick().unwrap_err().code(), Code::Io);
844 assert_eq!(r.engine().runs(), Vec::<u64>::new());
845 assert_eq!(r.carried(), 0, "nothing was drained, so nothing is held");
846
847 r.engine_mut().fail_submit = false;
849 assert_eq!(r.tick().unwrap().commands, 1);
850 }
851
852 #[test]
853 fn a_failed_drain_still_ran_the_batch_and_flushed_it() {
854 let (mut r, tx, _e) = wired(1);
855 tx[0].push(1).unwrap();
856 r.engine_mut().fail_drain = true;
857 assert_eq!(r.tick().unwrap_err().code(), Code::Io);
858 assert_eq!(r.engine().runs(), vec![1]);
859 assert_eq!(r.engine().count(&Step::Flush), 1);
860 }
861
862 #[test]
863 fn a_command_with_no_key_is_run_without_a_hash() {
864 let (mut r, tx, _e) = wired(1);
865 tx[0].push(KEYLESS).unwrap();
866 let turn = r.tick().unwrap();
867 assert_eq!(turn.commands, 1);
868 assert_eq!(
869 r.engine().count(&Step::Prefetch(KEYLESS)),
870 0,
871 "there is nothing to warm for a command with no key"
872 );
873 }
874
875 #[test]
876 fn inline_execution_takes_the_same_walk_as_the_loop() {
877 let mut r = Reactor::inline(Recorder::new());
878 assert_eq!(r.execute(9), Flow::Next);
879 assert_eq!(
880 r.engine().steps(),
881 vec![Step::Prefetch(9), Step::Run(9)],
882 "no submit, no flush and no maintenance, and the same two calls"
883 );
884 assert_eq!(r.commands(), 1);
885 }
886
887 #[test]
888 fn inline_batches_pay_for_the_epoch_once() {
889 let mut r = Reactor::inline(Recorder::new());
890 assert_eq!(r.execute_all(0..20), 20);
891 assert_eq!(r.engine().runs(), (0..20).collect::<Vec<_>>());
892 let (last_prefetch, first_run) = walk_boundary(&r.engine().steps());
893 assert!(last_prefetch < first_run, "inline gets the two walks too");
894 }
895
896 #[test]
897 fn an_inline_batch_stops_at_a_break() {
898 let mut r = Reactor::inline(Recorder::new());
899 r.engine_mut().break_on = Some(2);
900 assert_eq!(r.execute_all(0..10), 3);
901 assert_eq!(r.engine().runs(), vec![0, 1, 2]);
902 assert_eq!(r.carried(), 0, "inline holds nothing over to next time");
903
904 r.engine_mut().break_on = None;
906 assert_eq!(r.execute_all(100..103), 3);
907 assert_eq!(r.engine().runs(), vec![0, 1, 2, 100, 101, 102]);
908 }
909
910 #[test]
911 fn an_empty_inline_batch_is_not_a_turn() {
912 let mut r = Reactor::inline(Recorder::new());
913 assert_eq!(r.execute_all(Vec::new()), 0);
914 assert_eq!(r.turns(), 0);
915 assert!(r.engine().steps().is_empty());
916 }
917
918 #[test]
919 fn run_until_returns_when_the_flag_is_set_and_the_lanes_are_dry() {
920 let (mut r, tx, _e) = wired(1);
921 for i in 0..300 {
922 tx[0].push(i).unwrap();
923 }
924 let stop = AtomicBool::new(true);
925 r.run_until(&stop).unwrap();
926 assert_eq!(r.commands(), 300, "the flag does not drop queued work");
927 assert!(r.turns() >= 5);
928 }
929
930 #[test]
931 fn a_reactor_says_what_it_has_been_doing() {
932 let (mut r, tx, _e) = wired(1);
933 tx[0].push(1).unwrap();
934 r.tick().unwrap();
935 let said = format!("{r:?}");
936 assert!(said.contains("commands: 1"), "{said}");
937 assert!(said.contains("lanes: 1"), "{said}");
938 }
939}