Skip to main content

fizzy_sdk/resilience/
bulkhead.rs

1use std::sync::Arc;
2use std::time::Duration;
3
4use tokio::sync::{OwnedSemaphorePermit, Semaphore};
5
6use crate::error::Error;
7
8/// How many calls of one kind may run at once, and how long a caller waits for room.
9///
10/// `max_concurrent` of zero reads as "leave it at the default", the same normalising Go
11/// does. A `max_wait` of zero means a caller never waits: a full bulkhead refuses at once.
12#[derive(Debug, Clone)]
13pub struct BulkheadConfig {
14    /// How many calls of one kind may run at once.
15    pub max_concurrent: usize,
16    /// How long [`Bulkhead::acquire`] waits for a permit, and so how long the client's own
17    /// gate holds a call back before refusing it.
18    pub max_wait: Duration,
19}
20
21impl Default for BulkheadConfig {
22    fn default() -> BulkheadConfig {
23        BulkheadConfig {
24            max_concurrent: 10,
25            max_wait: Duration::from_secs(5),
26        }
27    }
28}
29
30/// One scope's ration of calls in flight. Every call holds a [`BulkheadPermit`] for as long
31/// as it runs, and the scope refuses the rest until one is dropped, so a slow operation
32/// cannot take every connection the client has.
33pub struct Bulkhead {
34    max_concurrent: usize,
35    max_wait: Duration,
36    permits: Arc<Semaphore>,
37}
38
39/// A place in the bulkhead, given back by dropping it.
40pub type BulkheadPermit = OwnedSemaphorePermit;
41
42impl Bulkhead {
43    /// A bulkhead of the configured width.
44    #[allow(clippy::needless_pass_by_value)] // the config is consumed by convention, like the others
45    pub fn new(config: BulkheadConfig) -> Bulkhead {
46        let max_concurrent = match config.max_concurrent {
47            0 => BulkheadConfig::default().max_concurrent,
48            max_concurrent => max_concurrent,
49        };
50        Bulkhead {
51            max_concurrent,
52            max_wait: config.max_wait,
53            permits: Arc::new(Semaphore::new(max_concurrent)),
54        }
55    }
56
57    /// A permit, waiting up to [`BulkheadConfig::max_wait`] for one to free up. Answers
58    /// [`Error::bulkhead_full`] when the wait runs out, or straight away when there is no
59    /// wait to spend.
60    pub async fn acquire(&self) -> Result<BulkheadPermit, Error> {
61        match self.try_acquire() {
62            Some(permit) => Ok(permit),
63            None if self.max_wait.is_zero() => Err(Error::bulkhead_full()),
64            None => tokio::time::timeout(self.max_wait, self.permits.clone().acquire_owned())
65                .await
66                .map_err(|_| Error::bulkhead_full())?
67                .map_err(|_| Error::bulkhead_full()),
68        }
69    }
70
71    /// A permit if the scope has room this instant, and nothing if it does not.
72    pub fn try_acquire(&self) -> Option<BulkheadPermit> {
73        self.permits.clone().try_acquire_owned().ok()
74    }
75
76    /// How many more calls this scope will take right now.
77    pub fn available(&self) -> usize {
78        self.permits.available_permits()
79    }
80
81    /// How many calls this scope is running right now.
82    pub fn in_use(&self) -> usize {
83        self.max_concurrent - self.available()
84    }
85}
86
87#[cfg(test)]
88#[allow(clippy::unwrap_used)]
89mod tests {
90    use super::*;
91
92    #[tokio::test]
93    async fn a_full_bulkhead_refuses_until_a_call_finishes() {
94        let bulkhead = Bulkhead::new(BulkheadConfig {
95            max_concurrent: 2,
96            max_wait: Duration::ZERO,
97        });
98
99        let first = bulkhead.acquire().await.unwrap();
100        let second = bulkhead.acquire().await.unwrap();
101        assert_eq!(2, bulkhead.in_use());
102        assert_eq!(0, bulkhead.available());
103
104        let refused = bulkhead.acquire().await.unwrap_err();
105        assert_eq!(Some(crate::error::Refusal::BulkheadFull), refused.refusal());
106        assert_eq!("bulkhead is full", refused.message());
107
108        drop(first);
109        assert_eq!(1, bulkhead.in_use());
110
111        let third = bulkhead.acquire().await.unwrap();
112        drop((second, third));
113        assert_eq!(0, bulkhead.in_use());
114    }
115
116    #[test]
117    fn try_acquire_answers_without_waiting() {
118        let bulkhead = Bulkhead::new(BulkheadConfig {
119            max_concurrent: 1,
120            max_wait: Duration::from_secs(5),
121        });
122
123        let permit = bulkhead.try_acquire();
124        assert!(permit.is_some());
125        assert!(bulkhead.try_acquire().is_none());
126
127        drop(permit);
128        assert!(bulkhead.try_acquire().is_some());
129    }
130
131    #[tokio::test]
132    async fn a_caller_that_may_wait_gets_the_permit_the_call_before_it_gives_back() {
133        let bulkhead = Arc::new(Bulkhead::new(BulkheadConfig {
134            max_concurrent: 1,
135            max_wait: Duration::from_millis(500),
136        }));
137
138        let held = bulkhead.acquire().await.unwrap();
139        tokio::spawn(async move {
140            tokio::time::sleep(Duration::from_millis(10)).await;
141            drop(held);
142        });
143
144        assert!(bulkhead.acquire().await.is_ok());
145    }
146
147    #[tokio::test]
148    async fn a_caller_that_waits_too_long_is_refused() {
149        let bulkhead = Bulkhead::new(BulkheadConfig {
150            max_concurrent: 1,
151            max_wait: Duration::from_millis(10),
152        });
153        let _held = bulkhead.acquire().await.unwrap();
154
155        let refused = bulkhead.acquire().await.unwrap_err();
156
157        assert_eq!(Some(crate::error::Refusal::BulkheadFull), refused.refusal());
158    }
159
160    #[test]
161    fn a_config_of_zero_falls_back_to_the_default_width() {
162        let bulkhead = Bulkhead::new(BulkheadConfig {
163            max_concurrent: 0,
164            max_wait: Duration::ZERO,
165        });
166
167        assert_eq!(10, bulkhead.available());
168    }
169}