1use crate::{Error, Result};
7use std::{
8 collections::VecDeque,
9 fmt,
10 future::Future,
11 pin::Pin,
12 sync::{Arc, Mutex},
13 task::{Context, Poll},
14 time::Duration,
15};
16use tokio::{sync::Notify, time::Instant};
17use tower_layer::Layer;
18use tower_service::Service;
19
20pub(crate) const REQUESTS_PER_SECOND: u32 = 120;
21const STREAM_LIMIT: usize = 20;
22const ADMISSION_WINDOW: Duration = Duration::from_secs(1);
23
24#[derive(Clone)]
29pub struct Admission {
30 pub(crate) inner: Arc<AdmissionInner>,
31}
32
33impl fmt::Debug for Admission {
34 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
35 f.debug_struct("Admission").finish_non_exhaustive()
36 }
37}
38
39pub(crate) struct AdmissionInner {
40 pub(crate) rate: Mutex<Rate>,
41 policy_changed: Notify,
42 pub(crate) stream_slots: Mutex<StreamSlots>,
43 connection_gate: Arc<ConnectionGate>,
44}
45
46pub(crate) struct Rate {
47 pub(crate) exit_reserve: u32,
48 pub(crate) admitted: VecDeque<Instant>,
49 pub(crate) cooldown_until: Option<Instant>,
50}
51
52pub(crate) struct StreamSlots {
53 pub(crate) active: usize,
54}
55
56impl Default for Admission {
57 fn default() -> Self {
58 Self::new()
59 }
60}
61
62impl Admission {
63 #[must_use]
65 pub fn new() -> Self {
66 Self::create(0)
67 }
68
69 pub fn with_exit_reserve(requests: u32) -> Result<Self> {
75 if requests >= REQUESTS_PER_SECOND {
76 return Err(Error::InvalidInput(
77 "exit reserve must be less than 120 requests".into(),
78 ));
79 }
80 Ok(Self::create(requests))
81 }
82
83 fn create(requests: u32) -> Self {
84 Self {
85 inner: Arc::new(AdmissionInner {
86 rate: Mutex::new(Rate {
87 exit_reserve: requests,
88 admitted: VecDeque::new(),
89 cooldown_until: None,
90 }),
91 policy_changed: Notify::new(),
92 stream_slots: Mutex::new(StreamSlots { active: 0 }),
93 connection_gate: Arc::new(ConnectionGate {
94 starts: Mutex::new(ConnectionStarts::Empty),
95 #[cfg(test)]
96 wait_observer: Mutex::new(None),
97 }),
98 }),
99 }
100 }
101
102 pub fn set_exit_reserve(&self, requests: u32) -> Result<()> {
108 if requests >= REQUESTS_PER_SECOND {
109 return Err(Error::InvalidInput(
110 "exit reserve must be less than 120 requests".into(),
111 ));
112 }
113 let mut rate = self
114 .inner
115 .rate
116 .lock()
117 .map_err(|_| Error::InvalidInput("rate lock poisoned".into()))?;
118 let changed = rate.exit_reserve != requests;
119 rate.exit_reserve = requests;
120 drop(rate);
121 if changed {
122 self.inner.policy_changed.notify_waiters();
123 }
124 Ok(())
125 }
126
127 pub(crate) fn exit_reserve(&self) -> Result<u32> {
128 self.inner
129 .rate
130 .lock()
131 .map(|rate| rate.exit_reserve)
132 .map_err(|_| Error::InvalidInput("rate lock poisoned".into()))
133 }
134
135 pub(crate) async fn admit(&self, mutation: bool, exit: bool) -> Result<()> {
136 loop {
137 let changed = self.inner.policy_changed.notified();
138 tokio::pin!(changed);
139 changed.as_mut().enable();
140 let wait = {
141 let mut rate = self
142 .inner
143 .rate
144 .lock()
145 .map_err(|_| Error::InvalidInput("rate lock poisoned".into()))?;
146 let capacity = if exit {
147 REQUESTS_PER_SECOND
148 } else {
149 REQUESTS_PER_SECOND - rate.exit_reserve
150 };
151 let now = Instant::now();
152 while rate
153 .admitted
154 .front()
155 .is_some_and(|at| now.duration_since(*at) >= ADMISSION_WINDOW)
156 {
157 rate.admitted.pop_front();
158 }
159 if let Some(until) = rate.cooldown_until.filter(|until| *until > now) {
160 until.duration_since(now)
161 } else if u32::try_from(rate.admitted.len())
162 .map_err(|_| Error::Invariant("rate admission count exceeds u32"))?
163 < capacity
164 {
165 rate.admitted.push_back(now);
166 return Ok(());
167 } else {
168 let oldest = rate
169 .admitted
170 .front()
171 .ok_or(Error::Invariant("full rate window has no admission"))?;
172 ADMISSION_WINDOW.saturating_sub(now.duration_since(*oldest))
173 }
174 };
175 if mutation {
176 return Err(Error::RateLimited);
177 }
178 tokio::select! {
179 () = &mut changed => {},
180 () = tokio::time::sleep(wait) => {},
181 }
182 }
183 }
184
185 pub(crate) fn install_cooldown(&self, cooldown: Duration) {
186 let mut rate = self
187 .inner
188 .rate
189 .lock()
190 .unwrap_or_else(std::sync::PoisonError::into_inner);
191 let until = Instant::now() + cooldown;
192 rate.cooldown_until = Some(
193 rate.cooldown_until
194 .map_or(until, |existing| existing.max(until)),
195 );
196 }
197
198 pub(crate) fn lease_stream(&self) -> Result<StreamLease> {
199 let mut slots = self
200 .inner
201 .stream_slots
202 .lock()
203 .map_err(|_| Error::InvalidInput("stream slots lock poisoned".into()))?;
204 if slots.active >= STREAM_LIMIT {
205 return Err(Error::RateLimited);
206 }
207 slots.active += 1;
208 Ok(StreamLease {
209 owner: self.clone(),
210 })
211 }
212
213 pub(crate) fn connection_layer(&self) -> ConnectionLayer {
214 ConnectionLayer {
215 gate: Arc::clone(&self.inner.connection_gate),
216 }
217 }
218}
219
220pub(crate) struct StreamLease {
221 pub(crate) owner: Admission,
222}
223
224impl Drop for StreamLease {
225 fn drop(&mut self) {
226 let mut slots = self
227 .owner
228 .inner
229 .stream_slots
230 .lock()
231 .unwrap_or_else(std::sync::PoisonError::into_inner);
232 slots.active = slots.active.saturating_sub(1);
233 }
234}
235
236#[derive(Clone)]
237pub(crate) struct ConnectionLayer {
238 gate: Arc<ConnectionGate>,
239}
240
241impl<S> Layer<S> for ConnectionLayer {
242 type Service = ConnectionService<S>;
243
244 fn layer(&self, inner: S) -> Self::Service {
245 ConnectionService {
246 inner,
247 gate: Arc::clone(&self.gate),
248 }
249 }
250}
251
252#[derive(Clone)]
253pub(crate) struct ConnectionService<S> {
254 inner: S,
255 gate: Arc<ConnectionGate>,
256}
257
258struct ConnectionGate {
259 starts: Mutex<ConnectionStarts>,
260 #[cfg(test)]
261 wait_observer: Mutex<Option<tokio::sync::mpsc::UnboundedSender<Instant>>>,
262}
263
264enum ConnectionStarts {
265 Empty,
266 One(Instant),
267 Two { oldest: Instant, newest: Instant },
268}
269
270impl ConnectionStarts {
271 fn admit_or_wait(&mut self, now: Instant) -> Option<Instant> {
272 match self {
273 Self::Empty => {
274 *self = Self::One(now);
275 None
276 }
277 Self::One(oldest) if now.duration_since(*oldest) >= ADMISSION_WINDOW => {
278 *oldest = now;
279 None
280 }
281 Self::One(oldest) => {
282 *self = Self::Two {
283 oldest: *oldest,
284 newest: now,
285 };
286 None
287 }
288 Self::Two { newest, .. } if now.duration_since(*newest) >= ADMISSION_WINDOW => {
289 *self = Self::One(now);
290 None
291 }
292 Self::Two { oldest, newest } if now.duration_since(*oldest) >= ADMISSION_WINDOW => {
293 *oldest = *newest;
294 *newest = now;
295 None
296 }
297 Self::Two { oldest, .. } => Some(*oldest + ADMISSION_WINDOW),
298 }
299 }
300}
301
302impl ConnectionGate {
303 async fn admit_start(&self) {
304 loop {
305 let wait_until = {
306 let mut starts = self
307 .starts
308 .lock()
309 .unwrap_or_else(std::sync::PoisonError::into_inner);
310 starts.admit_or_wait(Instant::now())
311 };
312 if let Some(until) = wait_until {
313 #[cfg(test)]
314 if until > Instant::now() {
315 let observer = self
316 .wait_observer
317 .lock()
318 .unwrap_or_else(std::sync::PoisonError::into_inner);
319 if let Some(observer) = observer.as_ref() {
320 let _ = observer.send(until);
321 }
322 }
323 tokio::time::sleep_until(until).await;
324 } else {
325 return;
326 }
327 }
328 }
329}
330
331impl<S, Request> Service<Request> for ConnectionService<S>
332where
333 S: Service<Request>,
334 S::Future: Send + 'static,
335{
336 type Response = S::Response;
337 type Error = S::Error;
338 type Future =
339 Pin<Box<dyn Future<Output = std::result::Result<Self::Response, Self::Error>> + Send>>;
340
341 fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<std::result::Result<(), Self::Error>> {
342 self.inner.poll_ready(cx)
343 }
344
345 fn call(&mut self, request: Request) -> Self::Future {
346 let connection = self.inner.call(request);
347 let gate = Arc::clone(&self.gate);
348 Box::pin(async move {
349 gate.admit_start().await;
350 connection.await
351 })
352 }
353}
354
355#[cfg(test)]
356#[allow(clippy::expect_used, clippy::unwrap_used)]
358mod tests {
359 use super::*;
360 use crate::{AccountID, Client, Environment};
361 use std::convert::Infallible;
362 use tokio::{
363 io::{AsyncReadExt, AsyncWriteExt},
364 net::{TcpListener, TcpStream},
365 sync::{mpsc, oneshot},
366 };
367 use url::Url;
368
369 struct HeldConnector {
370 id: u8,
371 begun: mpsc::UnboundedSender<u8>,
372 release: Option<oneshot::Receiver<()>>,
373 }
374
375 impl Service<()> for HeldConnector {
376 type Response = ();
377 type Error = Infallible;
378 type Future = Pin<Box<dyn Future<Output = std::result::Result<(), Infallible>> + Send>>;
379
380 fn poll_ready(
381 &mut self,
382 _cx: &mut Context<'_>,
383 ) -> Poll<std::result::Result<(), Self::Error>> {
384 Poll::Ready(Ok(()))
385 }
386
387 fn call(&mut self, _request: ()) -> Self::Future {
388 let begun = self.begun.clone();
389 let release = self.release.take().expect("one call");
390 let id = self.id;
391 Box::pin(async move {
392 begun.send(id).expect("observer");
393 let _ = release.await;
394 Ok(())
395 })
396 }
397 }
398
399 async fn loopback() -> (
400 Url,
401 mpsc::UnboundedReceiver<Instant>,
402 tokio::task::JoinHandle<()>,
403 ) {
404 let listener = TcpListener::bind("127.0.0.1:0").await.expect("bind");
405 let url = Url::parse(&format!(
406 "http://{}/",
407 listener.local_addr().expect("address")
408 ))
409 .expect("URL");
410 let (accepted, observations) = mpsc::unbounded_channel();
411 let task = tokio::spawn(async move {
412 while let Ok((socket, _)) = listener.accept().await {
413 accepted.send(Instant::now()).expect("accept receiver");
414 tokio::spawn(serve_socket(socket));
415 }
416 });
417 (url, observations, task)
418 }
419
420 async fn serve_socket(mut socket: TcpStream) {
421 let mut pending = Vec::new();
422 loop {
423 let mut buffer = [0_u8; 4096];
424 let read = socket.read(&mut buffer).await.expect("read");
425 if read == 0 {
426 return;
427 }
428 pending.extend_from_slice(&buffer[..read]);
429 let Some(end) = pending.windows(4).position(|bytes| bytes == b"\r\n\r\n") else {
430 continue;
431 };
432 let stream = pending[..end]
433 .windows(b"/transactions/stream".len())
434 .any(|part| part == b"/transactions/stream");
435 pending.drain(..end + 4);
436 if stream {
437 socket
438 .write_all(
439 b"HTTP/1.1 200 OK\r\nContent-Type: application/octet-stream\r\nTransfer-Encoding: chunked\r\n\r\n",
440 )
441 .await
442 .expect("stream header");
443 std::future::pending::<()>().await;
444 }
445 let body = br#"{"accounts":[]}"#;
446 let header = format!(
447 "HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: keep-alive\r\n\r\n",
448 body.len()
449 );
450 socket.write_all(header.as_bytes()).await.expect("header");
451 socket.write_all(body).await.expect("body");
452 }
453 }
454
455 fn keep_clock_frozen() -> tokio::task::JoinHandle<()> {
456 tokio::spawn(async {
457 loop {
458 tokio::task::yield_now().await;
459 }
460 })
461 }
462
463 #[tokio::test(start_paused = true)]
464 async fn built_rest_and_stream_clients_share_two_connector_starts_per_second() {
465 let clock_guard = keep_clock_frozen();
466 let (url, mut accepts, server) = loopback().await;
467 let owner = Admission::new();
468 let (wait_sender, mut waits) = mpsc::unbounded_channel();
469 *owner
470 .inner
471 .connection_gate
472 .wait_observer
473 .lock()
474 .expect("wait observer") = Some(wait_sender);
475 let first_client = Client::builder(Environment::Practice, "token-a")
476 .admission(owner.clone())
477 .endpoints(url.clone(), url.clone())
478 .build()
479 .expect("first client");
480 let second_client = Client::builder(Environment::Practice, "token-b")
481 .admission(owner)
482 .endpoints(url.clone(), url)
483 .build()
484 .expect("second client");
485 let (stop_watchdog, watchdog_stopped) = std::sync::mpsc::channel();
486 let (expired, mut watchdog_expired) = oneshot::channel();
487 let watchdog = std::thread::spawn(move || {
488 if matches!(
489 watchdog_stopped.recv_timeout(Duration::from_secs(5)),
490 Err(std::sync::mpsc::RecvTimeoutError::Timeout)
491 ) {
492 let _ = expired.send(());
493 }
494 });
495
496 first_client.list_accounts().await.expect("first REST");
497 let first = accepts.recv().await.expect("first accept");
498 first_client
499 .list_accounts()
500 .await
501 .expect("reused REST socket");
502 assert!(accepts.try_recv().is_err());
503 assert!(waits.try_recv().is_err());
504
505 second_client
506 .list_accounts()
507 .await
508 .expect("second REST response");
509 let second = accepts.recv().await.expect("second accept");
510 assert_eq!(second, first);
511 assert!(waits.try_recv().is_err());
512
513 let account = AccountID::new("101-001-1-001").expect("account");
514 let mut stream_task =
515 tokio::spawn(async move { first_client.stream_transactions(&account).await });
516 let stream_deadline = tokio::select! {
517 observed = waits.recv() => observed,
518 _ = &mut stream_task => None,
519 _ = &mut watchdog_expired => None,
520 }
521 .expect("stream reached the connector start budget");
522 assert_eq!(
523 stream_deadline.duration_since(first),
524 Duration::from_secs(1)
525 );
526 tokio::time::advance(Duration::from_millis(999)).await;
527 assert!(accepts.try_recv().is_err());
528 tokio::time::advance(Duration::from_millis(1)).await;
529 let stream = stream_task
530 .await
531 .expect("stream task")
532 .expect("stream response");
533 let third = accepts.recv().await.expect("stream accept");
534 assert_eq!(third, stream_deadline);
535
536 drop(stream);
537 stop_watchdog.send(()).expect("stop watchdog");
538 watchdog.join().expect("watchdog");
539 clock_guard.abort();
540 server.abort();
541 }
542
543 #[tokio::test(start_paused = true)]
544 async fn two_connectors_start_while_held_and_third_waits_for_a_rolling_second() {
545 let owner = Admission::new();
546 let clock_guard = keep_clock_frozen();
547 let (begun, mut observed) = mpsc::unbounded_channel();
548 let (release_first, first_done) = oneshot::channel();
549 let (release_second, second_done) = oneshot::channel();
550 let (release_third, third_done) = oneshot::channel();
551 let mut first = owner.connection_layer().layer(HeldConnector {
552 id: 1,
553 begun: begun.clone(),
554 release: Some(first_done),
555 });
556 let mut second = owner.connection_layer().layer(HeldConnector {
557 id: 2,
558 begun: begun.clone(),
559 release: Some(second_done),
560 });
561 let mut third = owner.connection_layer().layer(HeldConnector {
562 id: 3,
563 begun,
564 release: Some(third_done),
565 });
566 let started = Instant::now();
567 let first_task = tokio::spawn(async move { first.call(()).await });
568 assert_eq!(observed.recv().await, Some(1));
569 let second_task = tokio::spawn(async move { second.call(()).await });
570 assert_eq!(observed.recv().await, Some(2));
571 assert_eq!(Instant::now(), started);
572 let third_task = tokio::spawn(async move { third.call(()).await });
573 tokio::time::advance(Duration::from_millis(999)).await;
574 assert!(observed.try_recv().is_err());
575 tokio::time::advance(Duration::from_millis(1)).await;
576 assert_eq!(observed.recv().await, Some(3));
577 assert_eq!(Instant::now(), started + Duration::from_secs(1));
578 release_first.send(()).expect("release first");
579 release_second.send(()).expect("release second");
580 release_third.send(()).expect("release third");
581 first_task
582 .await
583 .expect("first task")
584 .expect("first connect");
585 second_task
586 .await
587 .expect("second task")
588 .expect("second connect");
589 third_task
590 .await
591 .expect("third task")
592 .expect("third connect");
593 clock_guard.abort();
594 }
595
596 #[tokio::test(start_paused = true)]
597 async fn cancelled_started_connector_keeps_budget_but_queued_cancellation_spends_none() {
598 let owner = Admission::new();
599 let clock_guard = keep_clock_frozen();
600 let (begun, mut observed) = mpsc::unbounded_channel();
601 let (_release_first, first_done) = oneshot::channel();
602 let (_release_second, second_done) = oneshot::channel();
603 let (_release_third, third_done) = oneshot::channel();
604 let (release_fourth, fourth_done) = oneshot::channel();
605 let mut first = owner.connection_layer().layer(HeldConnector {
606 id: 1,
607 begun: begun.clone(),
608 release: Some(first_done),
609 });
610 let mut second = owner.connection_layer().layer(HeldConnector {
611 id: 2,
612 begun: begun.clone(),
613 release: Some(second_done),
614 });
615 let mut third = owner.connection_layer().layer(HeldConnector {
616 id: 3,
617 begun: begun.clone(),
618 release: Some(third_done),
619 });
620 let mut fourth = owner.connection_layer().layer(HeldConnector {
621 id: 4,
622 begun,
623 release: Some(fourth_done),
624 });
625 let first_task = tokio::spawn(async move { first.call(()).await });
626 assert_eq!(observed.recv().await, Some(1));
627 first_task.abort();
628 let _ = first_task.await;
629 let second_task = tokio::spawn(async move { second.call(()).await });
630 assert_eq!(observed.recv().await, Some(2));
631 let third_task = tokio::spawn(async move { third.call(()).await });
632 tokio::time::advance(Duration::from_millis(999)).await;
633 assert!(observed.try_recv().is_err());
634 third_task.abort();
635 let _ = third_task.await;
636 let fourth_task = tokio::spawn(async move { fourth.call(()).await });
637 assert!(observed.try_recv().is_err());
638 tokio::time::advance(Duration::from_millis(1)).await;
639 assert_eq!(observed.recv().await, Some(4));
640 release_fourth.send(()).expect("release fourth");
641 fourth_task
642 .await
643 .expect("fourth task")
644 .expect("fourth connect");
645 second_task.abort();
646 let _ = second_task.await;
647 clock_guard.abort();
648 }
649
650 #[tokio::test(start_paused = true)]
651 async fn shared_rate_and_cooldown_keep_exit_reserve_spent_across_clones() {
652 let owner = Admission::with_exit_reserve(119).expect("reserve");
653 owner.admit(true, false).await.expect("ordinary capacity");
654 assert!(matches!(
655 owner.clone().admit(true, false).await,
656 Err(Error::RateLimited)
657 ));
658 owner
659 .admit(true, true)
660 .await
661 .expect("reserved exit capacity");
662 for _ in 0..118 {
663 owner
664 .admit(true, true)
665 .await
666 .expect("remaining exit capacity");
667 }
668 owner.set_exit_reserve(0).expect("release reserve");
669 assert!(matches!(
670 owner.admit(true, false).await,
671 Err(Error::RateLimited)
672 ));
673 owner.install_cooldown(Duration::from_secs(2));
674 assert!(matches!(
675 owner.clone().admit(true, true).await,
676 Err(Error::RateLimited)
677 ));
678 assert!(matches!(
679 owner.admit(true, false).await,
680 Err(Error::RateLimited)
681 ));
682 tokio::time::advance(Duration::from_secs(2)).await;
683 owner.admit(true, false).await.expect("cooldown expired");
684 }
685
686 #[tokio::test(start_paused = true)]
687 async fn shared_owner_admits_one_hundred_twenty_rest_requests_and_refuses_the_next_mutation() {
688 let owner = Admission::new();
689 for _ in 0..120 {
690 owner
691 .clone()
692 .admit(true, false)
693 .await
694 .expect("documented REST capacity");
695 }
696 assert!(matches!(
697 owner.admit(true, false).await,
698 Err(Error::RateLimited)
699 ));
700 }
701
702 #[tokio::test(start_paused = true)]
703 async fn rolling_window_does_not_refill_a_late_burst_at_the_first_requests_expiry() {
704 let owner = Admission::new();
705 owner.admit(true, false).await.expect("first request");
706 tokio::time::advance(Duration::from_millis(999)).await;
707 for _ in 0..119 {
708 owner.admit(true, false).await.expect("late burst");
709 }
710 tokio::time::advance(Duration::from_millis(1)).await;
711 owner.admit(true, false).await.expect("one expired request");
712 assert!(matches!(
713 owner.admit(true, false).await,
714 Err(Error::RateLimited)
715 ));
716 tokio::time::advance(Duration::from_millis(998)).await;
717 assert!(matches!(
718 owner.admit(true, false).await,
719 Err(Error::RateLimited)
720 ));
721 tokio::time::advance(Duration::from_millis(1)).await;
722 owner.admit(true, false).await.expect("late burst expired");
723 }
724
725 #[tokio::test(start_paused = true)]
726 async fn reserve_of_one_hundred_nineteen_leaves_one_ordinary_request_and_all_exits() {
727 let owner = Admission::with_exit_reserve(119).expect("largest valid reserve");
728 owner.admit(true, false).await.expect("ordinary request");
729 assert!(matches!(
730 owner.admit(true, false).await,
731 Err(Error::RateLimited)
732 ));
733 for _ in 0..119 {
734 owner.admit(true, true).await.expect("reserved exit");
735 }
736 assert!(matches!(
737 owner.admit(true, true).await,
738 Err(Error::RateLimited)
739 ));
740 assert!(matches!(
741 Admission::with_exit_reserve(120),
742 Err(Error::InvalidInput(_))
743 ));
744 }
745
746 #[tokio::test(start_paused = true)]
747 async fn twenty_stream_leases_start_together_and_the_next_is_unsent() {
748 let owner = Admission::new();
749 let mut leases = Vec::new();
750 for _ in 0..20 {
751 leases.push(owner.lease_stream().expect("active stream lease"));
752 }
753 assert!(matches!(owner.lease_stream(), Err(Error::RateLimited)));
754 drop(leases.pop());
755 leases.push(owner.lease_stream().expect("released lease"));
756 }
757}