tower_memlim/
layer.rs

1use std::time::Duration;
2
3use crate::{
4    error::BoxError,
5    memory::{AvailableMemory, Threshold},
6    service::MemoryLimit,
7};
8
9use tower_layer::Layer;
10
11/// Enforces a limit on the underlying service when a memory [Threshold] is met.
12///
13/// A common use case is to load shed (reject requests), once the threshold is met. For doing so you must add tower's `load_shed` layer.
14/// Otherwise the service enqueues requests until the memory is available again.
15#[derive(Debug, Clone)]
16pub struct MemoryLimitLayer<M>
17where
18    M: AvailableMemory,
19{
20    threshold: Threshold,
21    /// Interval in which the next memory check is performed, if the threshold is exceeded.
22    ///
23    /// The `retryìnterval` has no effect if this layer is wrapped within a load shed layer.
24    retry_interval: std::time::Duration,
25    /// Memory stats provider
26    provider: M,
27}
28
29impl<M> MemoryLimitLayer<M>
30where
31    M: AvailableMemory,
32{
33    /// Get the available memory. Can be used to validate the
34    /// inner memory stat provder before adding the layer.
35    pub fn available_memory(&self) -> Result<usize, BoxError> {
36        self.provider.available_memory()
37    }
38
39    /// Create a new concurrency limit layer with a default [MemoryLimitLayer::retry_interval] of 50ms.
40    pub const fn new(threshold: Threshold, provider: M) -> Self {
41        MemoryLimitLayer {
42            threshold,
43            provider,
44            retry_interval: std::time::Duration::from_millis(50),
45        }
46    }
47
48    /// Set a custom [MemoryLimitLayer::retry_interval] which determines when the next memory check is performed, if the threshold is exceeded.
49    pub fn with_retry_interval(self, retry_interval: Duration) -> Self {
50        MemoryLimitLayer {
51            threshold: self.threshold,
52            provider: self.provider,
53            retry_interval,
54        }
55    }
56}
57
58impl<S, M> Layer<S> for MemoryLimitLayer<M>
59where
60    M: AvailableMemory,
61{
62    type Service = MemoryLimit<S, M>;
63
64    fn layer(&self, service: S) -> Self::Service {
65        MemoryLimit::new(
66            service,
67            self.threshold.clone(),
68            self.provider.clone(),
69            self.retry_interval,
70        )
71    }
72}
tower_memlim/layer.rs

tower_memlim/
layer.rs
