1use subms_rate_limiter::{Acquire, RateLimiter};
20
21struct Order {
24 at_ms: u64,
25 symbol: &'static str,
26 action: &'static str,
27 weight: u64,
28}
29
30const TAPE: &[Order] = &[
33 Order {
34 at_ms: 0,
35 symbol: "ESU5",
36 action: "new",
37 weight: 1,
38 },
39 Order {
40 at_ms: 0,
41 symbol: "ESU5",
42 action: "new",
43 weight: 1,
44 },
45 Order {
46 at_ms: 0,
47 symbol: "NQU5",
48 action: "new",
49 weight: 1,
50 },
51 Order {
52 at_ms: 0,
53 symbol: "ESU5",
54 action: "cancel-replace",
55 weight: 3,
56 },
57 Order {
58 at_ms: 1,
59 symbol: "NQU5",
60 action: "new",
61 weight: 1,
62 },
63 Order {
64 at_ms: 1,
65 symbol: "ESU5",
66 action: "new",
67 weight: 1,
68 },
69 Order {
70 at_ms: 4,
71 symbol: "ESU5",
72 action: "new",
73 weight: 1,
74 },
75 Order {
76 at_ms: 9,
77 symbol: "NQU5",
78 action: "cancel-replace",
79 weight: 3,
80 },
81];
82
83const MS: u64 = 1_000_000;
84
85fn main() {
86 session_throttle();
87
88 #[cfg(feature = "keyed")]
89 per_symbol_quota();
90
91 #[cfg(feature = "token-bucket")]
92 weighted_message_budget();
93
94 #[cfg(feature = "hierarchical")]
95 desk_gateway_cap();
96
97 #[cfg(feature = "distributed-backend")]
98 per_account_quota();
99
100 #[cfg(feature = "metrics")]
101 metered_feed_throttle();
102}
103
104fn session_throttle() {
110 println!("== session throttle: 1000 msg/sec, burst 5 ==");
111 let session = RateLimiter::new(1000.0, 5);
112
113 let mut sent = 0u64;
114 let mut units = 0u64;
115 for o in TAPE {
116 let now = o.at_ms * MS;
117 match session.try_acquire_n_with_retry_at(now, o.weight) {
118 Acquire::Ok => {
119 sent += 1;
120 units += o.weight;
121 println!(" t={:>2}ms {:<5} {:<14} sent", o.at_ms, o.symbol, o.action);
122 }
123 Acquire::Retry(wait) => {
124 println!(
125 " t={:>2}ms {:<5} {:<14} throttled, retry after {} us",
126 o.at_ms,
127 o.symbol,
128 o.action,
129 wait.as_micros()
130 );
131 }
132 Acquire::Unattainable { burst_capacity } => {
133 println!(
134 " t={:>2}ms {:<5} {:<14} rejected: weight {} exceeds the burst of {burst_capacity}",
135 o.at_ms, o.symbol, o.action, o.weight
136 );
137 }
138 }
139 }
140 println!(
141 " -> {sent} of {} messages on the wire, {units} units",
142 TAPE.len()
143 );
144 assert_eq!(sent, 7);
145 assert_eq!(units, 9);
146
147 let wait = session
150 .time_until_ready_at(9 * MS, 3)
151 .expect("weight 3 fits a burst of 5");
152 println!(
153 " next weight-3 message conforms in {} us",
154 wait.as_micros()
155 );
156 assert_eq!(wait.as_micros(), 1000);
157 assert!(
158 session.time_until_ready_at(9 * MS, 6).is_none(),
159 "weight 6 can never fit a burst of 5"
160 );
161
162 session.reset();
164 assert_eq!(
165 session.time_until_ready_at(9 * MS, 5),
166 Some(std::time::Duration::ZERO)
167 );
168 println!(" after reconnect: the full burst of 5 is available again");
169}
170
171#[cfg(feature = "keyed")]
175fn per_symbol_quota() {
176 use subms_rate_limiter::KeyedRateLimiter;
177
178 println!("\n== keyed: per-symbol quota, 1000 msg/sec each, burst 2 ==");
179 let per_symbol = KeyedRateLimiter::new(1000.0, 2);
180
181 let mut sent = 0u64;
182 for o in TAPE {
183 let now = o.at_ms * MS;
184 if matches!(per_symbol.try_acquire_at(now, o.symbol, 1), Acquire::Ok) {
185 sent += 1;
186 } else {
187 println!(
188 " t={:>2}ms {:<5} throttled on its own quota",
189 o.at_ms, o.symbol
190 );
191 }
192 }
193 println!(" -> {sent} admitted across {} symbols", per_symbol.len());
194 assert_eq!(sent, 7);
195 assert_eq!(per_symbol.len(), 2);
196
197 let evicted = per_symbol.retain_active_at(20 * MS);
200 println!(
201 " swept at t=20ms: {evicted} idle symbols dropped, {} live",
202 per_symbol.len()
203 );
204 assert_eq!(evicted, 2);
205 assert!(per_symbol.is_empty());
206}
207
208#[cfg(feature = "token-bucket")]
212fn weighted_message_budget() {
213 use std::sync::Arc;
214
215 use subms_rate_limiter::{TestClock, TokenBucket};
216
217 println!("\n== token-bucket: weighted budget that banks idle credit ==");
218 let clock = Arc::new(TestClock::new());
219 let budget = TokenBucket::with_clock(10, 5.0, Box::new(SharedClock(clock.clone())));
221
222 assert!(budget.try_acquire(1), "new order costs 1 unit");
223 assert!(budget.try_acquire(5), "bulk cancel-replace costs 5 units");
224 println!(" after 1 + 5 units: {} left", budget.available());
225 assert_eq!(budget.available(), 4);
226
227 assert!(
229 !budget.try_acquire(5),
230 "insufficient budget rejects the batch"
231 );
232 assert_eq!(budget.available(), 4, "a rejected batch spends nothing");
233
234 clock.advance_ms(1_000); println!(" after 1s idle: {} left", budget.available());
236 assert!(budget.try_acquire(5), "banked credit admits the batch");
237}
238
239#[cfg(feature = "hierarchical")]
243fn desk_gateway_cap() {
244 use std::sync::Arc;
245
246 use subms_rate_limiter::{HierarchicalLimiter, TestClock};
247
248 println!("\n== hierarchical: desk uplink caps two strategies ==");
249 let clock = Arc::new(TestClock::new());
250 let c = clock.clone();
251 let desk =
253 HierarchicalLimiter::with_clock_fn(5, 0.0, 2, 10, 0.0, || Box::new(SharedClock(c.clone())));
254
255 let mut sent = 0usize;
256 for round in 0..10 {
257 if desk.try_acquire(round % 2, 1) {
258 sent += 1;
259 }
260 }
261 println!(" strategies offered 10 orders; uplink admitted {sent}");
262 assert_eq!(sent, 5, "the parent caps the desk aggregate at 5");
263}
264
265#[cfg(feature = "distributed-backend")]
270fn per_account_quota() {
271 use std::sync::Arc;
272
273 use subms_rate_limiter::{DistributedLimiter, InMemoryBackend, TestClock};
274
275 println!("\n== distributed-backend: one account quota, two routers ==");
276 let clock = Arc::new(TestClock::new());
277 let shared = Arc::new(InMemoryBackend::new());
278 let window_ns = 1_000_000_000u64; let router_a = DistributedLimiter::with_clock(
281 Box::new(SharedBackend(shared.clone())),
282 5,
283 window_ns,
284 Box::new(SharedClock(clock.clone())),
285 );
286 let router_b = DistributedLimiter::with_clock(
287 Box::new(SharedBackend(shared.clone())),
288 5,
289 window_ns,
290 Box::new(SharedClock(clock.clone())),
291 );
292
293 let account = "acct-42";
294 let mut admitted = 0usize;
295 for round in 0..8 {
296 let router = if round % 2 == 0 { &router_a } else { &router_b };
297 if router.try_acquire(account) {
298 admitted += 1;
299 }
300 }
301 println!(" 8 orders sprayed across 2 routers: {admitted} admitted (quota 5)");
302 assert_eq!(admitted, 5, "the shared quota holds across both routers");
303}
304
305#[cfg(feature = "metrics")]
309fn metered_feed_throttle() {
310 use std::sync::Arc;
311
312 use subms_rate_limiter::{MeteredTokenBucket, TestClock};
313
314 println!("\n== metrics: self-observing market-data throttle ==");
315 let clock = Arc::new(TestClock::new());
316 let feed = MeteredTokenBucket::with_clock(5, 100.0, Box::new(SharedClock(clock.clone())));
318
319 for _ in 0..8 {
320 feed.try_acquire(1);
321 }
322 let s = feed.snapshot();
323 println!(
324 " granted {}, rejected {}, headroom {}",
325 s.granted, s.rejected, s.available
326 );
327 assert_eq!(s.granted, 5);
328 assert_eq!(s.rejected, 3);
329
330 clock.advance_ms(100); assert!(feed.try_acquire(1), "the refilled feed admits again");
332 let s2 = feed.snapshot();
333 println!(
334 " after refill: granted {}, refills {}",
335 s2.granted, s2.refills
336 );
337 assert_eq!(s2.granted, 6);
338 assert!(s2.refills >= 1, "a refill step was observed");
339}
340
341#[cfg(any(
344 feature = "token-bucket",
345 feature = "hierarchical",
346 feature = "distributed-backend",
347 feature = "metrics",
348))]
349struct SharedClock(std::sync::Arc<subms_rate_limiter::TestClock>);
350
351#[cfg(any(
352 feature = "token-bucket",
353 feature = "hierarchical",
354 feature = "distributed-backend",
355 feature = "metrics",
356))]
357impl subms_rate_limiter::Clock for SharedClock {
358 fn now_ns(&self) -> u64 {
359 self.0.now_ns()
360 }
361}
362
363#[cfg(feature = "distributed-backend")]
366struct SharedBackend(std::sync::Arc<subms_rate_limiter::InMemoryBackend>);
367
368#[cfg(feature = "distributed-backend")]
369impl subms_rate_limiter::Backend for SharedBackend {
370 fn incr(&self, key: &str, window_start_ns: u64, ttl_ns: u64) -> u64 {
371 self.0.incr(key, window_start_ns, ttl_ns)
372 }
373 fn read(&self, key: &str, window_start_ns: u64) -> u64 {
374 self.0.read(key, window_start_ns)
375 }
376}