source_budget.rs (5511B)
1 use std::sync::Mutex; 2 3 pub(crate) const MAX_EVENT_BYTES: usize = radroots_event_codec::decode::MAX_EVENT_JSON_BYTES; 4 pub(crate) const MAX_FETCH_BYTES: usize = 8 * 1024 * 1024; 5 pub(crate) const MAX_FETCH_EVENTS: usize = 4096; 6 pub(crate) const MAX_FETCH_NOTIFICATIONS: usize = 8192; 7 8 #[derive(Debug, Default)] 9 struct Usage { 10 bytes: usize, 11 events: usize, 12 notifications: usize, 13 wire_bytes: usize, 14 wire_messages: usize, 15 wire_data: usize, 16 exhausted: bool, 17 } 18 19 /// One monotonic inventory shared by all relay batches. Reservations are not 20 /// refunded when a duplicate, malformed event or completed batch is discarded. 21 #[derive(Debug, Default)] 22 pub(crate) struct FetchBudget(Mutex<Usage>, tokio::sync::Notify); 23 24 impl FetchBudget { 25 pub(crate) fn wire(&self, bytes: usize, frames: usize, data: usize) -> bool { 26 let Ok(mut usage) = self.0.lock() else { 27 return false; 28 }; 29 if usage.exhausted 30 || bytes > MAX_FETCH_BYTES - usage.wire_bytes 31 || frames > MAX_FETCH_NOTIFICATIONS - usage.wire_messages 32 || data > MAX_FETCH_EVENTS - usage.wire_data 33 { 34 usage.exhausted = true; 35 self.1.notify_waiters(); 36 return false; 37 } 38 usage.wire_bytes += bytes; 39 usage.wire_messages += frames; 40 usage.wire_data += data; 41 true 42 } 43 44 pub(crate) fn exhausted(&self) -> bool { 45 self.0.lock().map_or(true, |usage| usage.exhausted) 46 } 47 48 pub(super) async fn wait_exhausted(&self) { 49 loop { 50 let notified = self.1.notified(); 51 tokio::pin!(notified); 52 notified.as_mut().enable(); 53 if self.exhausted() { 54 return; 55 } 56 notified.await; 57 } 58 } 59 60 pub(super) fn notification(&self) -> bool { 61 let Ok(mut usage) = self.0.lock() else { 62 return false; 63 }; 64 if usage.exhausted || usage.notifications == MAX_FETCH_NOTIFICATIONS { 65 usage.exhausted = true; 66 self.1.notify_waiters(); 67 return false; 68 } 69 usage.notifications += 1; 70 true 71 } 72 73 pub(super) fn event(&self, bytes: usize) -> bool { 74 let Ok(mut usage) = self.0.lock() else { 75 return false; 76 }; 77 if usage.exhausted 78 || bytes > MAX_EVENT_BYTES 79 || usage.events == MAX_FETCH_EVENTS 80 || bytes > MAX_FETCH_BYTES - usage.bytes 81 { 82 usage.exhausted = true; 83 self.1.notify_waiters(); 84 return false; 85 } 86 usage.events += 1; 87 usage.bytes += bytes; 88 true 89 } 90 } 91 92 #[cfg(test)] 93 mod tests { 94 use super::*; 95 96 #[test] 97 fn each_maximum_is_accepted_and_the_next_unit_is_rejected() { 98 let bytes = FetchBudget::default(); 99 assert!(!FetchBudget::default().event(MAX_EVENT_BYTES + 1)); 100 for _ in 0..MAX_FETCH_BYTES / MAX_EVENT_BYTES { 101 assert!(bytes.event(MAX_EVENT_BYTES)); 102 } 103 assert!(!bytes.event(1)); 104 let events = FetchBudget::default(); 105 for _ in 0..MAX_FETCH_EVENTS { 106 assert!(events.event(1)); 107 } 108 assert!(!events.event(1)); 109 let notifications = FetchBudget::default(); 110 for _ in 0..MAX_FETCH_NOTIFICATIONS { 111 assert!(notifications.notification()); 112 } 113 assert!(!notifications.notification()); 114 } 115 116 #[test] 117 fn competing_relays_cannot_overreserve_the_aggregate_budget() { 118 let budget = FetchBudget::default(); 119 let accepted = std::thread::scope(|scope| { 120 let tasks = (0..8) 121 .map(|_| scope.spawn(|| (0..32).filter(|_| budget.event(MAX_EVENT_BYTES)).count())) 122 .collect::<Vec<_>>(); 123 tasks 124 .into_iter() 125 .map(|task| task.join().unwrap()) 126 .sum::<usize>() 127 }); 128 assert_eq!(accepted * MAX_EVENT_BYTES, MAX_FETCH_BYTES); 129 assert!(!budget.event(1)); 130 } 131 132 #[tokio::test] 133 async fn raw_bytes_frames_and_data_attempts_have_sticky_independent_limits() { 134 for (bytes, frames, data) in [ 135 (MAX_FETCH_BYTES, 0, 0), 136 (0, MAX_FETCH_NOTIFICATIONS, 0), 137 (0, 0, MAX_FETCH_EVENTS), 138 ] { 139 let budget = FetchBudget::default(); 140 assert!(budget.wire(bytes, frames, data)); 141 assert!(!budget.exhausted()); 142 assert!(!budget.wire( 143 usize::from(bytes > 0), 144 usize::from(frames > 0), 145 usize::from(data > 0) 146 )); 147 assert!(budget.exhausted()); 148 assert!(!budget.wire(0, 0, 0)); 149 assert!(!budget.event(0)); 150 assert!(!budget.notification()); 151 tokio::time::timeout(std::time::Duration::from_secs(1), budget.wait_exhausted()) 152 .await 153 .unwrap(); 154 } 155 } 156 157 #[tokio::test] 158 async fn exhaustion_wakes_every_waiter_without_lost_notifications() { 159 let budget = FetchBudget::default(); 160 let first = budget.wait_exhausted(); 161 let second = budget.wait_exhausted(); 162 tokio::pin!(first, second); 163 assert!(futures::poll!(&mut first).is_pending()); 164 assert!(futures::poll!(&mut second).is_pending()); 165 assert!(!budget.wire(MAX_FETCH_BYTES + 1, 0, 0)); 166 assert!(futures::poll!(&mut first).is_ready()); 167 assert!(futures::poll!(&mut second).is_ready()); 168 } 169 }