congestion_limiter/limiter/
token.rs

1use std::{
2    sync::{
3        atomic::{self, AtomicUsize},
4        Arc,
5    },
6    time::Duration,
7};
8
9use tokio::{sync::OwnedSemaphorePermit, time::Instant};
10
11use super::partitioning::Scheduler;
12
13/// A concurrency token, required to run a job.
14///
15/// Release the token back to the [Limiter](crate::limiter::Limiter) after the job is complete.
16#[derive(Debug)]
17pub struct Token {
18    inner: Option<TokenInner>,
19    partition: Option<Partition>,
20
21    start: Instant,
22    #[cfg(test)]
23    latency: Duration,
24}
25
26#[derive(Debug)]
27pub(crate) struct TokenInner {
28    _permit: OwnedSemaphorePermit,
29    in_flight: Arc<AtomicUsize>,
30}
31
32#[derive(Debug)]
33pub(crate) struct Partition {
34    in_flight: Arc<AtomicUsize>,
35    scheduler: Arc<Scheduler>,
36}
37
38impl Token {
39    pub(crate) fn new(permit: OwnedSemaphorePermit, in_flight: Arc<AtomicUsize>) -> Self {
40        in_flight.fetch_add(1, atomic::Ordering::SeqCst);
41        Self {
42            inner: Some(TokenInner {
43                _permit: permit,
44                in_flight,
45            }),
46            partition: None,
47            start: Instant::now(),
48            #[cfg(test)]
49            latency: Duration::ZERO,
50        }
51    }
52
53    pub(crate) fn new_from_inner(inner: TokenInner) -> Self {
54        Self {
55            inner: Some(inner),
56            partition: None,
57            start: Instant::now(),
58            #[cfg(test)]
59            latency: Duration::ZERO,
60        }
61    }
62
63    pub(crate) fn for_partition(mut self, partition: Partition) -> Self {
64        partition.in_flight.fetch_add(1, atomic::Ordering::SeqCst);
65        self.partition = Some(partition);
66        self
67    }
68
69    #[cfg(test)]
70    pub(crate) fn set_latency(&mut self, latency: Duration) {
71        use std::ops::Sub;
72
73        use tokio::time::Instant;
74
75        self.start = Instant::now().sub(latency);
76        self.latency = latency;
77    }
78
79    #[cfg(test)]
80    pub(crate) fn latency(&self) -> Duration {
81        self.latency
82    }
83
84    #[cfg(not(test))]
85    pub(crate) fn latency(&self) -> Duration {
86        self.start.elapsed()
87    }
88}
89
90impl Drop for Token {
91    /// Reduces the number of jobs in flight and releases the token back to the available pool.
92    fn drop(&mut self) {
93        if let Some(partition) = self.partition.take() {
94            partition.in_flight.fetch_sub(1, atomic::Ordering::SeqCst);
95            partition.scheduler.reuse_permit(
96                self.inner
97                    .take()
98                    .expect("TokenInner should always be present until drop"),
99            );
100        }
101    }
102}
103
104impl Drop for TokenInner {
105    fn drop(&mut self) {
106        self.in_flight.fetch_sub(1, atomic::Ordering::SeqCst);
107    }
108}
109
110impl Partition {
111    pub(crate) fn new(in_flight: Arc<AtomicUsize>, scheduler: Arc<Scheduler>) -> Self {
112        Self {
113            in_flight,
114            scheduler,
115        }
116    }
117}
congestion_limiter/limiter/token.rs

congestion_limiter/limiter/
token.rs
