1use std::{
23 any::Any,
24 cell::{RefCell, UnsafeCell},
25 collections::VecDeque,
26 fmt::{Debug, Display},
27 marker::PhantomData,
28 num::{NonZeroU64, NonZeroUsize},
29 rc::Rc,
30};
31
32use nautilus_core::{DurationNanos, UnixNanos, correctness::FAILED};
33use serde::{Deserialize, Serialize};
34use ustr::Ustr;
35
36use crate::{
37 actor::{
38 Actor,
39 registry::{register_actor, try_get_actor_unchecked, with_actor_registry},
40 },
41 clock::Clock,
42 msgbus::{self, Endpoint, Handler, MStr, ShareableMessageHandler},
43 timer::{TimeEvent, TimeEventCallback},
44};
45
46const MAX_INITIAL_TIMESTAMPS_CAPACITY: usize = 1024;
47
48pub struct Throttler<T, F> {
75 clock: Rc<RefCell<dyn Clock>>,
76 actor_id: Ustr,
77 timer_name: Ustr,
78 limit: NonZeroUsize,
79 interval_ns: NonZeroU64,
80 large_limit: bool,
81 buffer: VecDeque<T>,
82 timestamps: VecDeque<UnixNanos>,
83 is_limiting: bool,
84 recv_count: usize,
85 sent_count: usize,
86 output_send: F,
87 output_drop: Option<F>,
88}
89
90impl<T, F> Actor for Throttler<T, F>
91where
92 T: 'static + Debug,
93 F: Fn(T) + 'static,
94{
95 fn id(&self) -> Ustr {
96 self.actor_id
97 }
98
99 fn handle(&mut self, _msg: &dyn Any) {}
100
101 fn as_any(&self) -> &dyn Any {
102 self
103 }
104}
105
106impl<T, F> Debug for Throttler<T, F>
107where
108 T: Debug,
109{
110 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
111 f.debug_struct(stringify!(Throttler))
112 .field("actor_id", &self.actor_id)
113 .field("timer_name", &self.timer_name)
114 .field("limit", &self.limit())
115 .field("interval_ns", &self.interval_ns().as_u64())
116 .field("buffer", &self.buffer)
117 .field("timestamps", &self.timestamps)
118 .field("is_limiting", &self.is_limiting)
119 .field("recv_count", &self.recv_count)
120 .field("sent_count", &self.sent_count)
121 .finish()
122 }
123}
124
125impl<T, F> Throttler<T, F>
126where
127 T: Debug,
128{
129 #[inline]
134 pub fn new(
135 rate_limit: RateLimit,
136 clock: Rc<RefCell<dyn Clock>>,
137 timer_name: &str,
138 output_send: F,
139 output_drop: Option<F>,
140 actor_id: Ustr,
141 ) -> Self {
142 Self {
143 clock,
144 actor_id,
145 timer_name: Ustr::from(format!("{timer_name}-{actor_id}").as_str()),
146 limit: rate_limit.limit,
147 interval_ns: rate_limit.interval_ns,
148 large_limit: rate_limit.limit.get() > MAX_INITIAL_TIMESTAMPS_CAPACITY,
149 buffer: VecDeque::new(),
150 timestamps: VecDeque::with_capacity(
151 rate_limit.limit.get().min(MAX_INITIAL_TIMESTAMPS_CAPACITY),
152 ),
153 is_limiting: false,
154 recv_count: 0,
155 sent_count: 0,
156 output_send,
157 output_drop,
158 }
159 }
160
161 #[inline]
174 pub(crate) fn set_timer(&self, callback: Option<TimeEventCallback>) {
175 let delta = self.delta_next();
176 self.set_timer_after(delta, callback);
177 }
178
179 #[inline]
180 fn set_timer_after(&self, delta: DurationNanos, callback: Option<TimeEventCallback>) {
181 let mut clock = self.clock.borrow_mut();
182 if clock.timer_exists(&self.timer_name) {
183 clock.cancel_timer(&self.timer_name);
184 }
185 let alert_ts = clock.timestamp_ns() + delta;
186
187 clock
188 .set_time_alert_ns(self.timer_name.as_str(), alert_ts, callback, Some(true))
189 .expect(FAILED);
190 }
191
192 #[inline]
197 pub fn delta_next(&self) -> DurationNanos {
198 match self.timestamps.get(self.limit.get() - 1) {
199 Some(ts) => {
200 let diff = self
201 .clock
202 .borrow()
203 .timestamp_ns()
204 .saturating_duration_since(*ts);
205 DurationNanos::new(self.interval_ns.get()).saturating_sub(diff)
206 }
207 None => DurationNanos::default(),
208 }
209 }
210
211 #[inline]
212 fn delta_next_at(&self, now: UnixNanos) -> DurationNanos {
213 match self.timestamps.get(self.limit.get() - 1) {
214 Some(ts) => {
215 let diff = now.saturating_duration_since(*ts);
216 DurationNanos::new(self.interval_ns.get()).saturating_sub(diff)
217 }
218 None => DurationNanos::default(),
219 }
220 }
221
222 #[inline]
225 pub fn reset(&mut self) {
226 self.cancel_timer_internal();
227 self.buffer.clear();
228 self.recv_count = 0;
229 self.sent_count = 0;
230 self.is_limiting = false;
231 self.timestamps.clear();
232 }
233}
234
235impl<T, F> Throttler<T, F> {
236 fn cancel_timer_internal(&self) {
243 if let Ok(mut clock) = self.clock.try_borrow_mut() {
244 clock.cancel_timer(&self.timer_name);
245 }
246 }
247
248 fn count_in_window(&self, now: UnixNanos) -> usize {
251 let Some(interval_start) = now.checked_sub(DurationNanos::new(self.interval_ns.get()))
252 else {
253 return self.timestamps.len();
254 };
255
256 if let Some(oldest) = self.timestamps.back()
257 && *oldest > interval_start
258 {
259 return self.timestamps.len();
260 }
261
262 match self.timestamps.front() {
263 Some(newest) if *newest > interval_start => {}
264 _ => return 0,
265 }
266
267 self.timestamps
268 .iter()
269 .take_while(|&&ts| ts > interval_start)
270 .count()
271 }
272
273 #[inline]
275 pub const fn limit(&self) -> usize {
276 self.limit.get()
277 }
278
279 #[inline]
281 pub const fn interval_ns(&self) -> DurationNanos {
282 DurationNanos::new(self.interval_ns.get())
283 }
284
285 #[inline]
287 pub const fn rate_limit(&self) -> RateLimit {
288 RateLimit {
289 limit: self.limit,
290 interval_ns: self.interval_ns,
291 }
292 }
293
294 #[inline]
296 pub fn qsize(&self) -> usize {
297 self.buffer.len()
298 }
299
300 #[inline]
302 pub fn used(&self) -> f64 {
303 if self.timestamps.is_empty() {
304 return 0.0;
305 }
306 let messages_in_current_interval = self.count_in_window(self.clock.borrow().timestamp_ns());
307 (messages_in_current_interval as f64) / (self.limit.get() as f64)
308 }
309
310 #[inline]
312 pub const fn is_limiting(&self) -> bool {
313 self.is_limiting
314 }
315
316 #[inline]
318 pub const fn recv_count(&self) -> usize {
319 self.recv_count
320 }
321
322 #[inline]
324 pub const fn sent_count(&self) -> usize {
325 self.sent_count
326 }
327}
328
329impl<T, F> Throttler<T, F>
330where
331 T: 'static + Debug,
332 F: Fn(T) + 'static,
333{
334 pub fn to_actor(self) -> Rc<UnsafeCell<Self>> {
337 let process_handler = ThrottlerProcess::<T, F>::new(self.actor_id);
339 msgbus::register_any(
340 process_handler.id().into(),
341 ShareableMessageHandler::from(Rc::new(process_handler) as Rc<dyn Handler<dyn Any>>),
342 );
343
344 register_actor(self)
346 }
347
348 pub fn dispose(&mut self) {
357 self.cancel_timer_internal();
358 msgbus::deregister_any(process_endpoint(self.actor_id));
359 with_actor_registry(|registry| {
360 registry.remove(&self.actor_id);
361 });
362 }
363
364 #[inline]
365 pub(crate) fn send_msg(&mut self, msg: T) {
366 let now = self.clock.borrow().timestamp_ns();
367
368 if self.timestamps.len() >= self.limit.get() {
369 self.timestamps.pop_back();
370 }
371 self.timestamps.push_front(now);
372
373 self.sent_count += 1;
374 (self.output_send)(msg);
375 }
376
377 #[inline]
385 pub fn try_reserve(&mut self, count: usize) -> bool {
386 self.recv_count += count;
387
388 if count == 0 {
389 return true;
390 }
391
392 let now = self.clock.borrow().timestamp_ns();
393 let delta = self.delta_next_at(now);
394 if self.is_limiting && delta.is_zero() && self.buffer.is_empty() {
395 self.is_limiting = false;
396 }
397
398 if self.is_limiting {
399 return false;
400 }
401
402 let used = self.count_in_window(now);
403
404 if self.limit.get().saturating_sub(used) < count {
405 self.is_limiting = true;
406
407 if !delta.is_zero() {
408 self.set_timer_after(delta, Some(throttler_resume::<T, F>(self.actor_id)));
409 }
410 return false;
411 }
412
413 for _ in 0..count {
414 if self.timestamps.len() >= self.limit.get() {
415 self.timestamps.pop_back();
416 }
417 self.timestamps.push_front(now);
418 }
419 self.sent_count += count;
420 true
421 }
422
423 #[inline]
424 pub(crate) fn limit_msg(&mut self, msg: T) {
425 if self.output_drop.is_none() {
426 self.buffer.push_front(msg);
427 log::debug!("Buffering {}", self.buffer.len());
428
429 if !self.is_limiting {
430 log::debug!("Limiting");
431 let cb = Some(ThrottlerProcess::<T, F>::new(self.actor_id).get_timer_callback());
432 let delta = self.delta_next();
433 self.set_timer_after(delta, cb);
434 self.is_limiting = true;
435 }
436 } else {
437 log::debug!("Dropping");
438
439 if let Some(drop) = &self.output_drop {
440 drop(msg);
441 }
442
443 if !self.is_limiting {
444 log::debug!("Limiting");
445 let delta = self.delta_next();
446 self.set_timer_after(delta, Some(throttler_resume::<T, F>(self.actor_id)));
447 self.is_limiting = true;
448 }
449 }
450 }
451
452 #[inline]
453 pub fn send(&mut self, msg: T)
454 where
455 T: 'static,
456 F: Fn(T) + 'static,
457 {
458 self.recv_count += 1;
459
460 if self.large_limit && self.timestamps.len() < self.limit.get() && !self.is_limiting {
461 self.send_msg(msg);
462 return;
463 }
464
465 let delta = if self.is_limiting && !self.buffer.is_empty() {
466 DurationNanos::default()
467 } else {
468 self.delta_next()
469 };
470
471 if self.is_limiting && delta.is_zero() && self.buffer.is_empty() {
476 self.is_limiting = false;
477 }
478
479 if self.is_limiting || !delta.is_zero() {
480 self.limit_msg(msg);
481 } else {
482 self.send_msg(msg);
483 }
484 }
485}
486
487fn process_endpoint(actor_id: Ustr) -> MStr<Endpoint> {
490 MStr::endpoint(format!("{actor_id}_process")).expect(FAILED)
491}
492
493struct ThrottlerProcess<T, F> {
498 actor_id: Ustr,
499 endpoint: MStr<Endpoint>,
500 phantom_t: PhantomData<T>,
501 phantom_f: PhantomData<F>,
502}
503
504impl<T, F> ThrottlerProcess<T, F>
505where
506 T: Debug,
507{
508 pub(crate) fn new(actor_id: Ustr) -> Self {
509 Self {
510 actor_id,
511 endpoint: process_endpoint(actor_id),
512 phantom_t: PhantomData,
513 phantom_f: PhantomData,
514 }
515 }
516
517 pub(crate) fn get_timer_callback(&self) -> TimeEventCallback {
518 let endpoint = self.endpoint;
519 TimeEventCallback::from(move |event: TimeEvent| {
520 msgbus::send_any(endpoint, &(event));
521 })
522 }
523}
524
525impl<T, F> Handler<dyn Any> for ThrottlerProcess<T, F>
526where
527 T: 'static + Debug,
528 F: Fn(T) + 'static,
529{
530 fn id(&self) -> Ustr {
531 *self.endpoint
532 }
533
534 fn handle(&self, _message: &dyn Any) {
535 let Some(mut throttler) = try_get_actor_unchecked::<Throttler<T, F>>(&self.actor_id) else {
538 return;
539 };
540
541 while let Some(msg) = throttler.buffer.pop_back() {
542 throttler.send_msg(msg);
543
544 if !throttler.buffer.is_empty() && !throttler.delta_next().is_zero() {
548 throttler.is_limiting = true;
549 throttler.set_timer(Some(self.get_timer_callback()));
550 return;
551 }
552 }
553
554 throttler.is_limiting = false;
555 }
556}
557
558impl<T, F> Drop for Throttler<T, F> {
559 fn drop(&mut self) {
560 self.cancel_timer_internal();
564 }
565}
566
567pub fn throttler_resume<T, F>(actor_id: Ustr) -> TimeEventCallback
577where
578 T: 'static + Debug,
579 F: Fn(T) + 'static,
580{
581 TimeEventCallback::from(move |_event: TimeEvent| {
582 if let Some(mut throttler) = try_get_actor_unchecked::<Throttler<T, F>>(&actor_id) {
583 if throttler.buffer.is_empty() {
584 throttler.is_limiting = false;
585 } else {
586 let cb = Some(ThrottlerProcess::<T, F>::new(actor_id).get_timer_callback());
587 throttler.set_timer(cb);
588 }
589 }
590 })
591}
592
593#[derive(Debug, Clone, Copy, PartialEq, Eq, Deserialize, Serialize)]
602#[serde(deny_unknown_fields)]
603pub struct RateLimit {
604 limit: NonZeroUsize,
605 interval_ns: NonZeroU64,
606}
607
608impl RateLimit {
609 pub fn new_checked(limit: usize, interval_ns: DurationNanos) -> anyhow::Result<Self> {
615 let limit = NonZeroUsize::new(limit)
616 .ok_or_else(|| anyhow::anyhow!("Invalid limit: {limit} (must be non-zero)"))?;
617
618 let interval_ns = NonZeroU64::new(interval_ns.as_u64()).ok_or_else(|| {
619 anyhow::anyhow!("Invalid interval_ns: {interval_ns} (must be non-zero)")
620 })?;
621
622 Ok(Self { limit, interval_ns })
623 }
624
625 #[must_use]
631 pub fn new(limit: usize, interval_ns: DurationNanos) -> Self {
632 Self::new_checked(limit, interval_ns).expect(FAILED)
633 }
634
635 #[must_use]
637 pub const fn limit(&self) -> usize {
638 self.limit.get()
639 }
640
641 #[must_use]
643 pub const fn interval_ns(&self) -> DurationNanos {
644 DurationNanos::new(self.interval_ns.get())
645 }
646}
647
648impl Display for RateLimit {
649 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
650 let interval = self.interval_ns();
651 let limit = self.limit();
652 let total_secs = interval.as_secs();
653 let remainder_ns = interval.subsec_nanos();
654 let hours = total_secs / 3600;
655 let minutes = (total_secs % 3600) / 60;
656 let seconds = total_secs % 60;
657
658 if remainder_ns == 0 {
659 write!(f, "{limit}/{hours:02}:{minutes:02}:{seconds:02}")
660 } else {
661 let micros = remainder_ns / 1_000;
662 write!(
663 f,
664 "{limit}/{hours:02}:{minutes:02}:{seconds:02}.{micros:06}"
665 )
666 }
667 }
668}
669
670#[cfg(test)]
671mod tests {
672 use std::{
673 cell::{RefCell, UnsafeCell},
674 rc::Rc,
675 };
676
677 use nautilus_core::{DurationNanos, UUID4, UnixNanos};
678 use rstest::{fixture, rstest};
679 use ustr::Ustr;
680
681 use super::{MAX_INITIAL_TIMESTAMPS_CAPACITY, RateLimit, Throttler, ThrottlerProcess};
682 use crate::{
683 clock::{Clock, VirtualClock},
684 msgbus::{self, Handler},
685 };
686 type SharedThrottler = Rc<UnsafeCell<Throttler<u64, Box<dyn Fn(u64)>>>>;
687
688 #[derive(Clone)]
693 struct TestThrottler {
694 throttler: SharedThrottler,
695 clock: Rc<RefCell<VirtualClock>>,
696 interval: DurationNanos,
697 }
698
699 #[allow(unsafe_code)]
700 impl TestThrottler {
701 #[expect(clippy::mut_from_ref)]
702 pub(crate) fn get_throttler(&self) -> &mut Throttler<u64, Box<dyn Fn(u64)>> {
703 unsafe { &mut *self.throttler.get() }
704 }
705 }
706
707 fn timer_count_with_prefix(
711 throttler: &Throttler<u64, Box<dyn Fn(u64)>>,
712 prefix: &str,
713 ) -> usize {
714 throttler
715 .clock
716 .borrow()
717 .timer_names()
718 .iter()
719 .filter(|name| name.starts_with(prefix))
720 .count()
721 }
722
723 #[allow(unsafe_code)]
724 #[expect(clippy::mut_from_ref)]
725 fn access_shared(t: &SharedThrottler) -> &mut Throttler<u64, Box<dyn Fn(u64)>> {
726 unsafe { &mut *t.get() }
727 }
728
729 #[rstest]
730 #[case(0, 1_000)]
731 #[case(1_000, 0)]
732 fn test_rate_limit_new_checked_rejects_zero(#[case] limit: usize, #[case] interval_ns: u64) {
733 assert!(RateLimit::new_checked(limit, DurationNanos::new(interval_ns)).is_err());
734 }
735
736 #[rstest]
737 #[case(0, 1_000)]
738 #[case(1_000, 0)]
739 #[should_panic]
740 fn test_rate_limit_new_panics_on_zero(#[case] limit: usize, #[case] interval_ns: u64) {
741 let _ = RateLimit::new(limit, DurationNanos::new(interval_ns));
742 }
743
744 #[rstest]
745 fn test_rate_limit_new_checked_accepts_positive() {
746 let rate = RateLimit::new_checked(5, DurationNanos::new(10)).unwrap();
747
748 assert_eq!(rate.limit(), 5);
749 assert_eq!(rate.interval_ns(), DurationNanos::new(10));
750 }
751
752 #[rstest]
753 #[case::whole_seconds(100, 1_000_000_000, "100/00:00:01")]
754 #[case::hours_minutes_seconds(17, 3_723_000_000_000, "17/01:02:03")]
755 #[case::large_hours(23, 360_000_000_000_000, "23/100:00:00")]
756 #[case::fractional_seconds(31, 1_234_567_000, "31/00:00:01.234567")]
757 #[case::truncated_nanoseconds(41, 1_234_567_890, "41/00:00:01.234567")]
758 #[case::below_microsecond(53, 999, "53/00:00:00.000000")]
759 fn test_rate_limit_display(
760 #[case] limit: usize,
761 #[case] interval_ns: u64,
762 #[case] expected: &str,
763 ) {
764 let rate = RateLimit::new(limit, DurationNanos::new(interval_ns));
765
766 assert_eq!(rate.to_string(), expected);
767 }
768
769 #[fixture]
770 pub fn test_throttler_buffered() -> TestThrottler {
771 let output_send: Box<dyn Fn(u64)> = Box::new(|msg: u64| {
772 log::debug!("Sent: {msg}");
773 });
774 let clock = Rc::new(RefCell::new(VirtualClock::new()));
775 let inner_clock = Rc::clone(&clock);
776 let rate_limit = RateLimit::new(5, DurationNanos::new(10));
777 let interval = rate_limit.interval_ns();
778 let actor_id = Ustr::from(UUID4::new().as_str());
779
780 TestThrottler {
781 throttler: Throttler::new(
782 rate_limit,
783 clock,
784 "buffer_timer",
785 output_send,
786 None,
787 actor_id,
788 )
789 .to_actor(),
790 clock: inner_clock,
791 interval,
792 }
793 }
794
795 #[fixture]
796 pub fn test_throttler_unbuffered() -> TestThrottler {
797 let output_send: Box<dyn Fn(u64)> = Box::new(|msg: u64| {
798 log::debug!("Sent: {msg}");
799 });
800 let output_drop: Box<dyn Fn(u64)> = Box::new(|msg: u64| {
801 log::debug!("Dropped: {msg}");
802 });
803 let clock = Rc::new(RefCell::new(VirtualClock::new()));
804 let inner_clock = Rc::clone(&clock);
805 let rate_limit = RateLimit::new(5, DurationNanos::new(10));
806 let interval = rate_limit.interval_ns();
807 let actor_id = Ustr::from(UUID4::new().as_str());
808
809 TestThrottler {
810 throttler: Throttler::new(
811 rate_limit,
812 clock,
813 "dropper_timer",
814 output_send,
815 Some(output_drop),
816 actor_id,
817 )
818 .to_actor(),
819 clock: inner_clock,
820 interval,
821 }
822 }
823
824 #[rstest]
825 fn test_buffering_send_to_limit_becomes_throttled(test_throttler_buffered: TestThrottler) {
826 let throttler = test_throttler_buffered.get_throttler();
827 for _ in 0..6 {
828 throttler.send(42);
829 }
830 assert_eq!(throttler.qsize(), 1);
831
832 assert!(throttler.is_limiting);
833 assert_eq!(throttler.recv_count, 6);
834 assert_eq!(throttler.sent_count, 5);
835 assert_eq!(timer_count_with_prefix(throttler, "buffer_timer"), 1);
836 }
837
838 #[rstest]
839 fn test_buffering_used_when_sent_to_limit_returns_one(test_throttler_buffered: TestThrottler) {
840 let throttler = test_throttler_buffered.get_throttler();
841
842 for _ in 0..5 {
843 throttler.send(42);
844 }
845
846 assert_eq!(throttler.used(), 1.0);
847 assert_eq!(throttler.recv_count, 5);
848 assert_eq!(throttler.sent_count, 5);
849 }
850
851 #[rstest]
852 fn test_buffering_used_when_half_interval_from_limit_returns_one(
853 test_throttler_buffered: TestThrottler,
854 ) {
855 let throttler = test_throttler_buffered.get_throttler();
856
857 for _ in 0..5 {
858 throttler.send(42);
859 }
860
861 let half_interval = test_throttler_buffered.interval / 2;
862 {
864 let mut clock = test_throttler_buffered.clock.borrow_mut();
865 clock.advance_time(UnixNanos::default() + half_interval, true);
866 }
867
868 assert_eq!(throttler.used(), 1.0);
869 assert_eq!(throttler.recv_count, 5);
870 assert_eq!(throttler.sent_count, 5);
871 }
872
873 #[rstest]
874 fn test_buffering_used_before_limit_when_halfway_returns_half(
875 test_throttler_buffered: TestThrottler,
876 ) {
877 let throttler = test_throttler_buffered.get_throttler();
878
879 for _ in 0..3 {
880 throttler.send(42);
881 }
882
883 assert_eq!(throttler.used(), 0.6);
884 assert_eq!(throttler.recv_count, 3);
885 assert_eq!(throttler.sent_count, 3);
886 }
887
888 #[rstest]
889 fn test_try_reserve_counts_messages_without_output(test_throttler_buffered: TestThrottler) {
890 let throttler = test_throttler_buffered.get_throttler();
891
892 assert!(throttler.try_reserve(3));
893
894 assert_eq!(throttler.used(), 0.6);
895 assert_eq!(throttler.recv_count, 3);
896 assert_eq!(throttler.sent_count, 3);
897 assert_eq!(throttler.qsize(), 0);
898 }
899
900 #[rstest]
901 fn test_try_reserve_rejects_when_full_batch_exceeds_limit(
902 test_throttler_buffered: TestThrottler,
903 ) {
904 let throttler = test_throttler_buffered.get_throttler();
905
906 assert!(throttler.try_reserve(3));
907 assert!(!throttler.try_reserve(3));
908
909 assert_eq!(throttler.used(), 0.6);
910 assert_eq!(throttler.recv_count, 6);
911 assert_eq!(throttler.sent_count, 3);
912 assert_eq!(throttler.qsize(), 0);
913 assert!(throttler.is_limiting);
914 assert_eq!(timer_count_with_prefix(throttler, "buffer_timer"), 0);
918
919 assert!(throttler.try_reserve(2));
920
921 assert_eq!(throttler.used(), 1.0);
922 assert_eq!(throttler.recv_count, 8);
923 assert_eq!(throttler.sent_count, 5);
924 assert_eq!(throttler.qsize(), 0);
925 assert!(!throttler.is_limiting);
926 }
927
928 #[rstest]
929 fn test_try_reserve_rejects_batch_larger_than_limit() {
930 let clock: Rc<RefCell<dyn Clock>> = Rc::new(RefCell::new(VirtualClock::new()));
931 let mut throttler = Throttler::<u64, Box<dyn Fn(u64)>>::new(
932 RateLimit::new(5, DurationNanos::new(10)),
933 clock,
934 "reserve_over_limit",
935 Box::new(|_| ()) as Box<dyn Fn(u64)>,
936 None,
937 Ustr::from("reserve-over-limit-actor"),
938 );
939
940 assert!(!throttler.try_reserve(6));
941
942 assert_eq!(throttler.used(), 0.0);
943 assert_eq!(throttler.recv_count, 6);
944 assert_eq!(throttler.sent_count, 0);
945 assert_eq!(throttler.qsize(), 0);
946 assert!(throttler.is_limiting);
947 assert_eq!(throttler.clock.borrow().timer_count(), 0);
948 }
949
950 #[rstest]
951 fn test_try_reserve_zero_count_is_noop() {
952 let clock: Rc<RefCell<dyn Clock>> = Rc::new(RefCell::new(VirtualClock::new()));
953 let mut throttler = Throttler::<u64, Box<dyn Fn(u64)>>::new(
954 RateLimit::new(5, DurationNanos::new(10)),
955 clock,
956 "reserve_zero",
957 Box::new(|_| ()) as Box<dyn Fn(u64)>,
958 None,
959 Ustr::from("reserve-zero-actor"),
960 );
961
962 assert!(throttler.try_reserve(0));
963
964 assert_eq!(throttler.used(), 0.0);
965 assert_eq!(throttler.recv_count, 0);
966 assert_eq!(throttler.sent_count, 0);
967 assert_eq!(throttler.qsize(), 0);
968 assert!(!throttler.is_limiting);
969 assert_eq!(throttler.clock.borrow().timer_count(), 0);
970 }
971
972 #[rstest]
973 fn test_buffering_refresh_when_at_limit_sends_remaining_items(
974 test_throttler_buffered: TestThrottler,
975 ) {
976 let throttler = test_throttler_buffered.get_throttler();
977
978 for _ in 0..6 {
979 throttler.send(42);
980 }
981
982 {
984 let mut clock = test_throttler_buffered.clock.borrow_mut();
985 let time_events = clock.advance_time(
986 UnixNanos::default() + test_throttler_buffered.interval,
987 true,
988 );
989
990 for each_event in clock.match_handlers(time_events) {
991 drop(clock); each_event.callback.call(each_event.event);
994
995 clock = test_throttler_buffered.clock.borrow_mut();
997 }
998 }
999
1000 assert_eq!(throttler.used(), 0.2);
1002 assert_eq!(throttler.recv_count, 6);
1003 assert_eq!(throttler.sent_count, 6);
1004 assert_eq!(throttler.qsize(), 0);
1005 }
1006
1007 #[rstest]
1008 fn test_buffering_send_message_after_buffering_message(test_throttler_buffered: TestThrottler) {
1009 let throttler = test_throttler_buffered.get_throttler();
1010
1011 for _ in 0..6 {
1012 throttler.send(43);
1013 }
1014
1015 {
1017 let mut clock = test_throttler_buffered.clock.borrow_mut();
1018 let time_events = clock.advance_time(
1019 UnixNanos::default() + test_throttler_buffered.interval,
1020 true,
1021 );
1022
1023 for each_event in clock.match_handlers(time_events) {
1024 drop(clock); each_event.callback.call(each_event.event);
1027
1028 clock = test_throttler_buffered.clock.borrow_mut();
1030 }
1031 }
1032
1033 for _ in 0..6 {
1034 throttler.send(42);
1035 }
1036
1037 assert_eq!(throttler.used(), 1.0);
1039 assert_eq!(throttler.recv_count, 12);
1040 assert_eq!(throttler.sent_count, 10);
1041 assert_eq!(throttler.qsize(), 2);
1042 }
1043
1044 #[rstest]
1045 fn test_buffering_send_message_after_halfway_after_buffering_message(
1046 test_throttler_buffered: TestThrottler,
1047 ) {
1048 let throttler = test_throttler_buffered.get_throttler();
1049
1050 for _ in 0..6 {
1051 throttler.send(42);
1052 }
1053
1054 {
1056 let mut clock = test_throttler_buffered.clock.borrow_mut();
1057 let time_events = clock.advance_time(
1058 UnixNanos::default() + test_throttler_buffered.interval,
1059 true,
1060 );
1061
1062 for each_event in clock.match_handlers(time_events) {
1063 drop(clock); each_event.callback.call(each_event.event);
1066
1067 clock = test_throttler_buffered.clock.borrow_mut();
1069 }
1070 }
1071
1072 for _ in 0..3 {
1073 throttler.send(42);
1074 }
1075
1076 assert_eq!(throttler.used(), 0.8);
1078 assert_eq!(throttler.recv_count, 9);
1079 assert_eq!(throttler.sent_count, 9);
1080 assert_eq!(throttler.qsize(), 0);
1081 }
1082
1083 #[rstest]
1084 fn test_dropping_send_sends_message_to_handler(test_throttler_unbuffered: TestThrottler) {
1085 let throttler = test_throttler_unbuffered.get_throttler();
1086 throttler.send(42);
1087
1088 assert!(!throttler.is_limiting);
1089 assert_eq!(throttler.recv_count, 1);
1090 assert_eq!(throttler.sent_count, 1);
1091 }
1092
1093 #[rstest]
1094 fn test_dropping_send_to_limit_drops_message(test_throttler_unbuffered: TestThrottler) {
1095 let throttler = test_throttler_unbuffered.get_throttler();
1096 for _ in 0..6 {
1097 throttler.send(42);
1098 }
1099 assert_eq!(throttler.qsize(), 0);
1100
1101 assert!(throttler.is_limiting);
1102 assert_eq!(throttler.used(), 1.0);
1103 assert_eq!(throttler.clock.borrow().timer_count(), 1);
1104 assert_eq!(timer_count_with_prefix(throttler, "dropper_timer"), 1);
1105 assert_eq!(throttler.recv_count, 6);
1106 assert_eq!(throttler.sent_count, 5);
1107 }
1108
1109 #[rstest]
1110 fn test_dropping_advance_time_when_at_limit_dropped_message(
1111 test_throttler_unbuffered: TestThrottler,
1112 ) {
1113 let throttler = test_throttler_unbuffered.get_throttler();
1114 for _ in 0..6 {
1115 throttler.send(42);
1116 }
1117
1118 {
1120 let mut clock = test_throttler_unbuffered.clock.borrow_mut();
1121 let time_events = clock.advance_time(
1122 UnixNanos::default() + test_throttler_unbuffered.interval,
1123 true,
1124 );
1125
1126 for each_event in clock.match_handlers(time_events) {
1127 drop(clock); each_event.callback.call(each_event.event);
1130
1131 clock = test_throttler_unbuffered.clock.borrow_mut();
1133 }
1134 }
1135
1136 assert_eq!(throttler.clock.borrow().timer_count(), 0);
1137 assert!(!throttler.is_limiting);
1138 assert_eq!(throttler.used(), 0.0);
1139 assert_eq!(throttler.recv_count, 6);
1140 assert_eq!(throttler.sent_count, 5);
1141 }
1142
1143 #[rstest]
1144 fn test_dropping_send_message_after_dropping_message(test_throttler_unbuffered: TestThrottler) {
1145 let throttler = test_throttler_unbuffered.get_throttler();
1146 for _ in 0..6 {
1147 throttler.send(42);
1148 }
1149
1150 {
1152 let mut clock = test_throttler_unbuffered.clock.borrow_mut();
1153 let time_events = clock.advance_time(
1154 UnixNanos::default() + test_throttler_unbuffered.interval,
1155 true,
1156 );
1157
1158 for each_event in clock.match_handlers(time_events) {
1159 drop(clock); each_event.callback.call(each_event.event);
1162
1163 clock = test_throttler_unbuffered.clock.borrow_mut();
1165 }
1166 }
1167
1168 throttler.send(42);
1169
1170 assert_eq!(throttler.used(), 0.2);
1171 assert_eq!(throttler.clock.borrow().timer_count(), 0);
1172 assert!(!throttler.is_limiting);
1173 assert_eq!(throttler.recv_count, 7);
1174 assert_eq!(throttler.sent_count, 6);
1175 }
1176
1177 #[rstest]
1178 fn test_embedded_dropping_auto_resets_after_window_without_actor_callback() {
1179 let clock: Rc<RefCell<VirtualClock>> = Rc::new(RefCell::new(VirtualClock::new()));
1180 let sent = Rc::new(RefCell::new(0));
1181 let dropped = Rc::new(RefCell::new(0));
1182
1183 let sent_cb = {
1184 let sent = Rc::clone(&sent);
1185 Box::new(move |_| *sent.borrow_mut() += 1) as Box<dyn Fn(u64)>
1186 };
1187 let drop_cb = {
1188 let dropped = Rc::clone(&dropped);
1189 Box::new(move |_| *dropped.borrow_mut() += 1) as Box<dyn Fn(u64)>
1190 };
1191
1192 let mut throttler = Throttler::new(
1193 RateLimit::new(5, DurationNanos::new(10)),
1194 Rc::clone(&clock) as Rc<RefCell<dyn Clock>>,
1195 "embedded_drop_timer",
1196 sent_cb,
1197 Some(drop_cb),
1198 Ustr::from("embedded-drop-actor"),
1199 );
1200
1201 for _ in 0..6 {
1202 throttler.send(42);
1203 }
1204 let events = clock.borrow_mut().advance_time(10.into(), true);
1205 throttler.send(42);
1206
1207 assert_eq!(events.len(), 1);
1208 assert_eq!(*sent.borrow(), 6);
1209 assert_eq!(*dropped.borrow(), 1);
1210 assert_eq!(throttler.recv_count, 7);
1211 assert_eq!(throttler.sent_count, 6);
1212 assert!(!throttler.is_limiting);
1213 assert_eq!(throttler.clock.borrow().timer_count(), 0);
1214 }
1215
1216 #[rstest]
1217 fn test_large_limit_fast_path_admits_until_limit_then_limits() {
1218 let clock: Rc<RefCell<dyn Clock>> = Rc::new(RefCell::new(VirtualClock::new()));
1219 let limit = MAX_INITIAL_TIMESTAMPS_CAPACITY + 1;
1220 let mut throttler = Throttler::<u64, Box<dyn Fn(u64)>>::new(
1221 RateLimit::new(limit, DurationNanos::new(10)),
1222 clock,
1223 "large_limit_timer",
1224 Box::new(|_| ()) as Box<dyn Fn(u64)>,
1225 None,
1226 Ustr::from("large-limit-actor"),
1227 );
1228
1229 for _ in 0..limit {
1230 throttler.send(42);
1231 }
1232 throttler.send(42);
1233
1234 assert_eq!(throttler.used(), 1.0);
1235 assert_eq!(throttler.recv_count, limit + 1);
1236 assert_eq!(throttler.sent_count, limit);
1237 assert_eq!(throttler.qsize(), 1);
1238 assert!(throttler.is_limiting);
1239 assert_eq!(throttler.clock.borrow().timer_count(), 1);
1240 }
1241
1242 #[rstest]
1243 fn test_new_preserves_rate_limit() {
1244 let clock: Rc<RefCell<dyn Clock>> = Rc::new(RefCell::new(VirtualClock::new()));
1245 let rate_limit = RateLimit::new(5, DurationNanos::new(10));
1246
1247 let throttler = Throttler::<u64, Box<dyn Fn(u64)>>::new(
1248 rate_limit,
1249 clock,
1250 "rate_limit",
1251 Box::new(|_| ()) as Box<dyn Fn(u64)>,
1252 None,
1253 Ustr::from("rate-limit-actor"),
1254 );
1255
1256 assert_eq!(throttler.rate_limit(), rate_limit);
1257 assert_eq!(throttler.limit(), 5);
1258 assert_eq!(throttler.interval_ns(), DurationNanos::new(10));
1259 }
1260
1261 #[rstest]
1262 fn test_debug_output_includes_identity_and_state() {
1263 let clock: Rc<RefCell<dyn Clock>> = Rc::new(RefCell::new(VirtualClock::new()));
1264 let actor_id = Ustr::from("debug-actor");
1265 let mut throttler = Throttler::<u64, Box<dyn Fn(u64)>>::new(
1266 RateLimit::new(5, DurationNanos::new(10)),
1267 clock,
1268 "debug_timer",
1269 Box::new(|_| ()) as Box<dyn Fn(u64)>,
1270 None,
1271 actor_id,
1272 );
1273
1274 throttler.send(42);
1275
1276 let debug = format!("{throttler:?}");
1277 let timer_name = Ustr::from("debug_timer-debug-actor");
1278
1279 assert!(debug.contains(&format!("actor_id: {actor_id:?}")));
1280 assert!(debug.contains(&format!("timer_name: {timer_name:?}")));
1281 assert!(debug.contains("limit: 5"));
1282 assert!(debug.contains("interval_ns: 10"));
1283 assert!(debug.contains("is_limiting: false"));
1284 assert!(debug.contains("recv_count: 1"));
1285 assert!(debug.contains("sent_count: 1"));
1286 }
1287
1288 #[rstest]
1289 fn test_reset_clears_state_and_cancels_timer(test_throttler_buffered: TestThrottler) {
1290 let throttler = test_throttler_buffered.get_throttler();
1291
1292 for _ in 0..6 {
1293 throttler.send(42);
1294 }
1295 assert_eq!(timer_count_with_prefix(throttler, "buffer_timer"), 1);
1296 assert_eq!(throttler.qsize(), 1);
1297
1298 throttler.reset();
1299
1300 assert_eq!(throttler.qsize(), 0);
1301 assert_eq!(throttler.recv_count, 0);
1302 assert_eq!(throttler.sent_count, 0);
1303 assert!(!throttler.is_limiting);
1304 assert!(throttler.timestamps.is_empty());
1305 assert_eq!(timer_count_with_prefix(throttler, "buffer_timer"), 0);
1306 assert_eq!(throttler.clock.borrow().timer_count(), 0);
1307 }
1308
1309 #[rstest]
1310 fn test_two_throttlers_share_clock_without_timer_collision() {
1311 let clock: Rc<RefCell<VirtualClock>> = Rc::new(RefCell::new(VirtualClock::new()));
1312 let interval = 10u64;
1313
1314 let mk = |base: &str| -> SharedThrottler {
1315 let clock: Rc<RefCell<dyn Clock>> = Rc::clone(&clock) as Rc<RefCell<dyn Clock>>;
1316 Throttler::new(
1317 RateLimit::new(5, DurationNanos::new(interval)),
1318 clock,
1319 base,
1320 Box::new(|_| ()) as Box<dyn Fn(u64)>,
1321 None,
1322 Ustr::from(UUID4::new().as_str()),
1323 )
1324 .to_actor()
1325 };
1326
1327 let t1 = mk("shared_timer");
1328 let t2 = mk("shared_timer");
1329
1330 {
1333 let t1 = access_shared(&t1);
1334
1335 for _ in 0..6 {
1336 t1.send(42);
1337 }
1338 }
1339 {
1340 let t2 = access_shared(&t2);
1341
1342 for _ in 0..6 {
1343 t2.send(42);
1344 }
1345 }
1346
1347 let clock_ref = clock.borrow();
1348 let names = clock_ref.timer_names();
1349 let shared_count = names
1350 .iter()
1351 .filter(|n| n.starts_with("shared_timer"))
1352 .count();
1353 assert_eq!(
1354 shared_count, 2,
1355 "two distinct timers expected, found {names:?}"
1356 );
1357 }
1358
1359 #[rstest]
1360 fn test_try_reserve_then_send_interleaved(test_throttler_buffered: TestThrottler) {
1361 let throttler = test_throttler_buffered.get_throttler();
1362
1363 assert!(throttler.try_reserve(3));
1366 throttler.send(42);
1367
1368 assert_eq!(throttler.recv_count, 4);
1369 assert_eq!(throttler.sent_count, 4);
1370 assert_eq!(throttler.used(), 0.8);
1371 assert!(!throttler.is_limiting);
1372 }
1373
1374 #[rstest]
1375 fn test_dispose_cancels_timer_and_deregisters_endpoint(test_throttler_buffered: TestThrottler) {
1376 let throttler = test_throttler_buffered.get_throttler();
1377
1378 for _ in 0..6 {
1379 throttler.send(42);
1380 }
1381 let actor_id = throttler.actor_id;
1382 let endpoint_name = format!("{actor_id}_process");
1383 assert_eq!(timer_count_with_prefix(throttler, "buffer_timer"), 1);
1384 assert!(msgbus::has_endpoint(&endpoint_name));
1385
1386 throttler.dispose();
1387
1388 assert_eq!(throttler.clock.borrow().timer_count(), 0);
1389 assert!(
1390 !msgbus::has_endpoint(&endpoint_name),
1391 "dispose must deregister the process endpoint"
1392 );
1393 }
1394
1395 #[rstest]
1396 fn test_try_reserve_then_buffered_sends_drain_in_order_after_window() {
1397 let clock: Rc<RefCell<VirtualClock>> = Rc::new(RefCell::new(VirtualClock::new()));
1398 let sent: Rc<RefCell<Vec<u64>>> = Rc::new(RefCell::new(Vec::new()));
1399 let sent_cb = {
1400 let sent = Rc::clone(&sent);
1401 Box::new(move |msg| sent.borrow_mut().push(msg)) as Box<dyn Fn(u64)>
1402 };
1403 let throttler = Throttler::new(
1404 RateLimit::new(5, DurationNanos::new(10)),
1405 Rc::clone(&clock) as Rc<RefCell<dyn Clock>>,
1406 "reserve_drain_timer",
1407 sent_cb,
1408 None,
1409 Ustr::from(UUID4::new().as_str()),
1410 )
1411 .to_actor();
1412 let throttler = access_shared(&throttler);
1413
1414 assert!(throttler.try_reserve(5));
1415 assert!(!throttler.try_reserve(1));
1416 assert!(throttler.is_limiting);
1417
1418 throttler.send(1);
1419 throttler.send(2);
1420 assert_eq!(throttler.qsize(), 2);
1421 assert!(sent.borrow().is_empty());
1422
1423 {
1426 let mut clock_ref = clock.borrow_mut();
1427 let time_events = clock_ref.advance_time(10.into(), true);
1428 for each_event in clock_ref.match_handlers(time_events) {
1429 drop(clock_ref);
1430 each_event.callback.call(each_event.event);
1431 clock_ref = clock.borrow_mut();
1432 }
1433 }
1434 assert!(throttler.is_limiting);
1435 assert_eq!(throttler.qsize(), 2);
1436 assert!(sent.borrow().is_empty());
1437
1438 throttler.send(3);
1440 assert_eq!(throttler.qsize(), 3);
1441 assert!(sent.borrow().is_empty());
1442
1443 for _ in 0..10 {
1444 if throttler.qsize() == 0 && !throttler.is_limiting {
1445 break;
1446 }
1447 let mut clock_ref = clock.borrow_mut();
1448 let current_time = clock_ref.get_time_ns();
1449 let time_events = clock_ref.advance_time(current_time + DurationNanos::new(10), true);
1450 for each_event in clock_ref.match_handlers(time_events) {
1451 drop(clock_ref);
1452 each_event.callback.call(each_event.event);
1453 clock_ref = clock.borrow_mut();
1454 }
1455 }
1456
1457 assert_eq!(*sent.borrow(), vec![1, 2, 3]);
1458 assert_eq!(throttler.qsize(), 0);
1459 assert!(!throttler.is_limiting);
1460 assert_eq!(throttler.sent_count, 8);
1461 assert_eq!(throttler.clock.borrow().timer_count(), 0);
1462 }
1463
1464 use proptest::prelude::*;
1469
1470 #[derive(Clone, Debug)]
1471 enum ThrottlerInput {
1472 SendMessage(u64),
1473 AdvanceClock(u8),
1474 }
1475
1476 fn throttler_input_strategy() -> impl Strategy<Value = ThrottlerInput> {
1478 prop_oneof![
1479 2 => prop::bool::ANY.prop_map(|_| ThrottlerInput::SendMessage(42)),
1480 8 => prop::num::u8::ANY.prop_map(|v| ThrottlerInput::AdvanceClock(v % 5 + 5)),
1481 ]
1482 }
1483
1484 fn throttler_test_strategy() -> impl Strategy<Value = Vec<ThrottlerInput>> {
1486 prop::collection::vec(throttler_input_strategy(), 10..=150)
1487 }
1488
1489 fn test_throttler_with_inputs(inputs: Vec<ThrottlerInput>, test_throttler: &TestThrottler) {
1490 let test_clock = test_throttler.clock.clone();
1491 let interval = test_throttler.interval;
1492 let throttler = test_throttler.get_throttler();
1493 let mut sent_count = 0;
1494
1495 for input in inputs {
1496 match input {
1497 ThrottlerInput::SendMessage(msg) => {
1498 throttler.send(msg);
1499 sent_count += 1;
1500 }
1501 ThrottlerInput::AdvanceClock(duration) => {
1502 let mut clock_ref = test_clock.borrow_mut();
1503 let current_time = clock_ref.get_time_ns();
1504 let time_events = clock_ref
1505 .advance_time(current_time + DurationNanos::new(u64::from(duration)), true);
1506
1507 for each_event in clock_ref.match_handlers(time_events) {
1508 drop(clock_ref);
1509 each_event.callback.call(each_event.event);
1510 clock_ref = test_clock.borrow_mut();
1511 }
1512 }
1513 }
1514
1515 let buffered_messages = throttler.qsize() > 0;
1520 let now = throttler.clock.borrow().timestamp_ns();
1521 let limit_filled_within_interval = throttler
1522 .timestamps
1523 .get(throttler.limit() - 1)
1524 .is_some_and(|&ts| (now - ts) < interval);
1525 let expected_limiting = buffered_messages && limit_filled_within_interval;
1526 assert_eq!(throttler.is_limiting, expected_limiting);
1527
1528 assert_eq!(sent_count, throttler.sent_count + throttler.qsize());
1530 }
1531
1532 for i in 1..=100u64 {
1536 if throttler.qsize() == 0 {
1537 break;
1538 }
1539 let advance_to = UnixNanos::default() + interval * 100 * i;
1540 let time_events = test_clock.borrow_mut().advance_time(advance_to, true);
1541 let mut clock_ref = test_clock.borrow_mut();
1542 for each_event in clock_ref.match_handlers(time_events) {
1543 drop(clock_ref);
1544 each_event.callback.call(each_event.event);
1545 clock_ref = test_clock.borrow_mut();
1546 }
1547 }
1548 assert_eq!(throttler.qsize(), 0);
1549 }
1550
1551 #[rstest]
1552 fn prop_test() {
1553 proptest!(|(inputs in throttler_test_strategy())| {
1556 let test_throttler = test_throttler_buffered();
1557 test_throttler_with_inputs(inputs, &test_throttler);
1558 });
1559 }
1560
1561 #[rstest]
1562 fn prop_test_dropping() {
1563 proptest!(|(inputs in throttler_test_strategy())| {
1567 let clock = Rc::new(RefCell::new(VirtualClock::new()));
1568 let sent: Rc<RefCell<usize>> = Rc::new(RefCell::new(0));
1569 let dropped: Rc<RefCell<usize>> = Rc::new(RefCell::new(0));
1570
1571 let sent_cb = {
1572 let sent = Rc::clone(&sent);
1573 Box::new(move |_| *sent.borrow_mut() += 1) as Box<dyn Fn(u64)>
1574 };
1575 let drop_cb = {
1576 let dropped = Rc::clone(&dropped);
1577 Box::new(move |_| *dropped.borrow_mut() += 1) as Box<dyn Fn(u64)>
1578 };
1579
1580 let interval = 10u64;
1581 let throttler = Throttler::new(
1582 RateLimit::new(5, DurationNanos::new(interval)),
1583 Rc::clone(&clock) as Rc<RefCell<dyn Clock>>,
1584 "prop_drop_timer",
1585 sent_cb,
1586 Some(drop_cb),
1587 Ustr::from(UUID4::new().as_str()),
1588 )
1589 .to_actor();
1590 let throttler = access_shared(&throttler);
1591
1592 for input in inputs {
1593 match input {
1594 ThrottlerInput::SendMessage(msg) => {
1595 throttler.send(msg);
1596 }
1597 ThrottlerInput::AdvanceClock(duration) => {
1598 let mut clock_ref = clock.borrow_mut();
1599 let current_time = clock_ref.get_time_ns();
1600 let time_events = clock_ref.advance_time(
1601 current_time + DurationNanos::new(u64::from(duration)),
1602 true,
1603 );
1604
1605 for each_event in clock_ref.match_handlers(time_events) {
1606 drop(clock_ref);
1607 each_event.callback.call(each_event.event);
1608 clock_ref = clock.borrow_mut();
1609 }
1610 }
1611 }
1612
1613 let sent_now = *sent.borrow();
1614 let dropped_now = *dropped.borrow();
1615 assert_eq!(sent_now + dropped_now, throttler.recv_count);
1617 assert_eq!(throttler.sent_count, sent_now);
1618 assert!(throttler.qsize() == 0, "drop mode must never buffer");
1619 }
1620 });
1621 }
1622
1623 #[derive(Clone, Debug)]
1624 enum ThrottlerReserveInput {
1625 SendMessage(u64),
1626 TryReserve(u8),
1627 AdvanceClock(u8),
1628 }
1629
1630 fn throttler_reserve_input_strategy() -> impl Strategy<Value = ThrottlerReserveInput> {
1631 prop_oneof![
1632 2 => prop::bool::ANY.prop_map(|_| ThrottlerReserveInput::SendMessage(42)),
1633 2 => prop::num::u8::ANY.prop_map(|v| ThrottlerReserveInput::TryReserve(v % 6 + 1)),
1634 6 => prop::num::u8::ANY.prop_map(|v| ThrottlerReserveInput::AdvanceClock(v % 5 + 5)),
1635 ]
1636 }
1637
1638 #[rstest]
1639 fn prop_test_try_reserve_interleaved() {
1640 proptest!(|(inputs in prop::collection::vec(throttler_reserve_input_strategy(), 10..=150))| {
1645 let test_throttler = test_throttler_buffered();
1646 let test_clock = test_throttler.clock.clone();
1647 let interval = test_throttler.interval;
1648 let throttler = test_throttler.get_throttler();
1649 let mut attempted = 0usize;
1650 let mut reserved = 0usize;
1651
1652 for input in inputs {
1653 match input {
1654 ThrottlerReserveInput::SendMessage(msg) => {
1655 throttler.send(msg);
1656 attempted += 1;
1657 }
1658 ThrottlerReserveInput::TryReserve(n) => {
1659 if throttler.try_reserve(usize::from(n)) {
1660 reserved += usize::from(n);
1661 }
1662 }
1663 ThrottlerReserveInput::AdvanceClock(duration) => {
1664 let mut clock_ref = test_clock.borrow_mut();
1665 let current_time = clock_ref.get_time_ns();
1666 let time_events = clock_ref.advance_time(
1667 current_time + DurationNanos::new(u64::from(duration)),
1668 true,
1669 );
1670
1671 for each_event in clock_ref.match_handlers(time_events) {
1672 drop(clock_ref);
1673 each_event.callback.call(each_event.event);
1674 clock_ref = test_clock.borrow_mut();
1675 }
1676 }
1677 }
1678
1679 assert_eq!(throttler.sent_count + throttler.qsize(), attempted + reserved);
1680 }
1681
1682 for i in 1..=100u64 {
1683 if throttler.qsize() == 0 {
1684 break;
1685 }
1686 let advance_to = UnixNanos::default() + interval * 100 * i;
1687 let time_events = test_clock
1688 .borrow_mut()
1689 .advance_time(advance_to, true);
1690 let mut clock_ref = test_clock.borrow_mut();
1691 for each_event in clock_ref.match_handlers(time_events) {
1692 drop(clock_ref);
1693 each_event.callback.call(each_event.event);
1694 clock_ref = test_clock.borrow_mut();
1695 }
1696 }
1697 assert_eq!(throttler.qsize(), 0);
1698 });
1699 }
1700
1701 #[rstest]
1702 fn test_throttler_process_id_returns_ustr() {
1703 let actor_id = Ustr::from("test_throttler");
1706 let process = ThrottlerProcess::<String, fn(String)>::new(actor_id);
1707
1708 let handler_id: Ustr = process.id();
1710
1711 assert!(handler_id.contains("test_throttler_process"));
1713 assert!(!handler_id.is_empty());
1714
1715 let _type_check: Ustr = handler_id;
1717 }
1718}