congestion-limiter 0.1.0

Dynamic congestion-based concurrency limits for controlling backpressure
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117

use std::{
    sync::{
        atomic::{self, AtomicUsize},
        Arc,
    },
    time::Duration,
};

use tokio::{sync::OwnedSemaphorePermit, time::Instant};

use super::partitioning::Scheduler;

/// A concurrency token, required to run a job.
///
/// Release the token back to the [Limiter](crate::limiter::Limiter) after the job is complete.
#[derive(Debug)]
pub struct Token {
    inner: Option<TokenInner>,
    partition: Option<Partition>,

    start: Instant,
    #[cfg(test)]
    latency: Duration,
}

#[derive(Debug)]
pub(crate) struct TokenInner {
    _permit: OwnedSemaphorePermit,
    in_flight: Arc<AtomicUsize>,
}

#[derive(Debug)]
pub(crate) struct Partition {
    in_flight: Arc<AtomicUsize>,
    scheduler: Arc<Scheduler>,
}

impl Token {
    pub(crate) fn new(permit: OwnedSemaphorePermit, in_flight: Arc<AtomicUsize>) -> Self {
        in_flight.fetch_add(1, atomic::Ordering::SeqCst);
        Self {
            inner: Some(TokenInner {
                _permit: permit,
                in_flight,
            }),
            partition: None,
            start: Instant::now(),
            #[cfg(test)]
            latency: Duration::ZERO,
        }
    }

    pub(crate) fn new_from_inner(inner: TokenInner) -> Self {
        Self {
            inner: Some(inner),
            partition: None,
            start: Instant::now(),
            #[cfg(test)]
            latency: Duration::ZERO,
        }
    }

    pub(crate) fn for_partition(mut self, partition: Partition) -> Self {
        partition.in_flight.fetch_add(1, atomic::Ordering::SeqCst);
        self.partition = Some(partition);
        self
    }

    #[cfg(test)]
    pub(crate) fn set_latency(&mut self, latency: Duration) {
        use std::ops::Sub;

        use tokio::time::Instant;

        self.start = Instant::now().sub(latency);
        self.latency = latency;
    }

    #[cfg(test)]
    pub(crate) fn latency(&self) -> Duration {
        self.latency
    }

    #[cfg(not(test))]
    pub(crate) fn latency(&self) -> Duration {
        self.start.elapsed()
    }
}

impl Drop for Token {
    /// Reduces the number of jobs in flight and releases the token back to the available pool.
    fn drop(&mut self) {
        if let Some(partition) = self.partition.take() {
            partition.in_flight.fetch_sub(1, atomic::Ordering::SeqCst);
            partition.scheduler.reuse_permit(
                self.inner
                    .take()
                    .expect("TokenInner should always be present until drop"),
            );
        }
    }
}

impl Drop for TokenInner {
    fn drop(&mut self) {
        self.in_flight.fetch_sub(1, atomic::Ordering::SeqCst);
    }
}

impl Partition {
    pub(crate) fn new(in_flight: Arc<AtomicUsize>, scheduler: Arc<Scheduler>) -> Self {
        Self {
            in_flight,
            scheduler,
        }
    }
}