tokio_too_busy/
lib.rs

1//! Identify when the Tokio workers are too busy to handle more work.
2//!
3//! ### Monitor the Tokio workers' busy ratio
4//! [TooBusy] allows you to track the busy ratio of the Tokio workers and react accordingly
5//! when they are too busy.
6//!
7//! In the example below, a [`TooBusy`] is [built][TooBusy::builder] and used to
8//! [evaluate][TooBusy::eval] the busy ratio of the Tokio workers and progressively reject requests
9//! when the busy ratio goes above the [low watermark][TooBusyBuilder::low_watermark] threshold:
10//! ```
11//!use axum::{
12//!   extract::{Request, State},
13//!   http::StatusCode,
14//!   middleware::{self, Next},
15//!   response::{IntoResponse, Response},
16//!   routing::get,
17//!   Router,
18//!};
19//!use tokio::time::Duration;
20//!use tokio_too_busy::*;
21//!
22//!#[derive(Clone)]
23//!struct AppState {
24//!    too_busy: TooBusy,
25//!}
26//!
27//!async fn my_middleware(State(state): State<AppState>, request: Request, next: Next) -> Response {
28//!
29//!   if state.too_busy.eval() {
30//!       return (StatusCode::SERVICE_UNAVAILABLE, "Server is too busy").into_response();
31//!   }
32//!   next.run(request).await
33//!}
34//!
35//!async fn root() -> &'static str {
36//!    "Hello, World!"
37//!}
38//!
39//!#[tokio::main(flavor = "multi_thread", worker_threads = 1)]
40//!async fn main() {
41//!   let too_busy = TooBusy::builder()
42//!       .interval(Duration::from_secs(1))
43//!       .ewma_alpha(0.1)
44//!       .high_watermark(95)
45//!       .low_watermark(90)
46//!       .build();
47//!
48//!   let state = AppState { too_busy };
49//!
50//!   let app = Router::new()
51//!      .route("/", get(root))
52//!      .route_layer(middleware::from_fn_with_state(state.clone(), my_middleware))
53//!      .with_state(state);
54//!
55//!  let listener = tokio::net::TcpListener::bind("0.0.0.0:8080").await.unwrap();
56//!  axum::serve(listener, app).await.unwrap();
57//!}
58//! ```
59//! Check the full example in the [examples] directory.
60//!
61//! Tokio too busy leverages the [worker_total_busy_duration] interface provided by the Tokio [Runtime] for knowledge.
62//! and its only available when Tokio is compiled using the following flags ```RUSTFLAGS="--cfg tokio_unstable"```.
63//!
64//! Busy ratio is calculated by averaging the busy duration of all workers since latest iteration and using an
65//! exponential moving average (EWMA).
66//!
67//! [examples]: https://github.com/pfreixes/tokio-too-busy/tree/main/examples
68//! [worker_total_busy_duration]: https://docs.rs/tokio/1.0.1/tokio/runtime/struct.Runtime.html#method.worker_total_busy_duration
69//! [Runtime]: https://docs.rs/tokio/1.0.1/tokio/runtime/struct.Runtime.html
70mod inner;
71use crate::inner::{LoadFeeder, TooBusyShared};
72use std::sync::Arc;
73use tokio::time::Duration;
74
75/// Track how busy are the tokio workers.
76///
77/// This type is internally reference-counted and can be freely cloned.
78pub struct TooBusy {
79    inner: Arc<TooBusyShared>,
80}
81
82impl TooBusy {
83    /// Create a new builder to configure the [`TooBusy`] instance.
84    ///
85    /// # Examples
86    ///
87    /// ```
88    /// use tokio_too_busy::TooBusy;
89    /// use std::time::Duration;
90    ///
91    /// let too_busy = TooBusy::builder()
92    ///     .interval(Duration::from_secs(1))
93    ///     .ewma_alpha(0.1)
94    ///     .high_watermark(95)
95    ///     .low_watermark(90)
96    ///     .build();
97    /// ```
98    pub fn builder() -> TooBusyBuilder {
99        TooBusyBuilder::default()
100    }
101
102    /// Get the current busy ratio of the Tokio workers.
103    pub fn ratio_busy_ewma(&self) -> u32 {
104        self.inner.ratio_busy_ewma()
105    }
106
107    /// Evaluate if the Tokio workers are too busy.
108    ///
109    /// Returns `true` if the busy ratio is above the low watermark threshold using
110    /// a progressive rejection strategy until reach the high watermark threshold, above
111    /// the high watermark all evaluations will return `true`. Otherwise, it returns `false`.
112    pub fn eval(&self) -> bool {
113        self.inner.eval()
114    }
115}
116
117impl Clone for TooBusy {
118    fn clone(&self) -> Self {
119        TooBusy {
120            inner: self.inner.clone(),
121        }
122    }
123}
124
125/// This builder allows you to configure the [`TooBusy`] instance.
126pub struct TooBusyBuilder {
127    interval: Duration,
128    low_watermark: u32,
129    high_watermark: u32,
130    ewma_alpha: f32,
131}
132
133impl TooBusyBuilder {
134    /// Set the interval to update the busy ratio.
135    ///
136    /// By default, it is 1 seconds.
137    pub fn interval(mut self, interval: Duration) -> TooBusyBuilder {
138        self.interval = interval;
139        self
140    }
141    /// Set the low watermark threshold to start considering the workers too busy.
142    ///
143    /// By default, it is 85%.
144    pub fn low_watermark(mut self, low_watermark: u32) -> TooBusyBuilder {
145        assert!(
146            low_watermark < self.high_watermark,
147            "low_watermark must be lower than high_watermark!"
148        );
149        self.low_watermark = low_watermark;
150        self
151    }
152    /// Set the high watermark threshold to start considering the workers too busy all the time.
153    ///
154    /// By default, it is 95%.
155    pub fn high_watermark(mut self, high_watermark: u32) -> TooBusyBuilder {
156        assert!(
157            high_watermark > self.low_watermark,
158            "high_watermark must be greater than low_watermark!"
159        );
160        self.high_watermark = high_watermark;
161        self
162    }
163    /// Set the alpha value for the EWMA calculation.
164    ///
165    /// By default, it is 0.1 (10%) for having quick reactions to changes. If you want to smooth the changes
166    /// you can set a higher value.
167    pub fn ewma_alpha(mut self, ewma_alpha: f32) -> TooBusyBuilder {
168        assert!(
169            ewma_alpha > 0.0 && ewma_alpha < 1.0,
170            "ewma_alpha must be between the range (0, 1)!"
171        );
172        self.ewma_alpha = ewma_alpha;
173        self
174    }
175    /// Build the [`TooBusy`] instance.
176    pub fn build(self) -> TooBusy {
177        let inner = Arc::new(TooBusyShared::new(self.low_watermark, self.high_watermark));
178        let weak_reference = Arc::downgrade(&inner);
179
180        tokio::spawn(async move {
181            LoadFeeder::new(weak_reference, self.interval, self.ewma_alpha)
182                .run()
183                .await;
184        });
185
186        TooBusy { inner }
187    }
188}
189
190impl Default for TooBusyBuilder {
191    fn default() -> Self {
192        TooBusyBuilder {
193            interval: Duration::from_secs(1),
194            low_watermark: 85,
195            high_watermark: 95,
196            ewma_alpha: 0.1,
197        }
198    }
199}
200
201#[cfg(test)]
202mod tests {
203    use super::*;
204
205    #[test]
206    #[should_panic]
207    fn too_busy_builder_invalid_low_watermark() {
208        TooBusyBuilder::default()
209            .high_watermark(90)
210            .low_watermark(90);
211    }
212
213    #[test]
214    #[should_panic]
215    fn too_busy_builder_invalid_high_watermark() {
216        TooBusyBuilder::default()
217            .low_watermark(90)
218            .high_watermark(90);
219    }
220
221    #[test]
222    #[should_panic]
223    fn too_busy_builder_invalid_ewma_alpha_too_low() {
224        TooBusyBuilder::default().ewma_alpha(0.0);
225    }
226
227    #[test]
228    #[should_panic]
229    fn too_busy_builder_invalid_ewma_alpha_too_high() {
230        TooBusyBuilder::default().ewma_alpha(1.0);
231    }
232}
tokio_too_busy/lib.rs

tokio_too_busy/
lib.rs
