common_memory_manager/
manager.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::sync::Arc;
16
17use snafu::ensure;
18use tokio::sync::{Semaphore, TryAcquireError};
19
20use crate::error::{
21    MemoryAcquireTimeoutSnafu, MemoryLimitExceededSnafu, MemorySemaphoreClosedSnafu, Result,
22};
23use crate::granularity::PermitGranularity;
24use crate::guard::MemoryGuard;
25use crate::policy::OnExhaustedPolicy;
26
27/// Trait for recording memory usage metrics.
28pub trait MemoryMetrics: Clone + Send + Sync + 'static {
29    fn set_limit(&self, bytes: i64);
30    fn set_in_use(&self, bytes: i64);
31    fn inc_rejected(&self, reason: &str);
32}
33
34/// Generic memory manager for quota-controlled operations.
35#[derive(Clone)]
36pub struct MemoryManager<M: MemoryMetrics> {
37    quota: Option<MemoryQuota<M>>,
38}
39
40#[derive(Clone)]
41pub(crate) struct MemoryQuota<M: MemoryMetrics> {
42    pub(crate) semaphore: Arc<Semaphore>,
43    pub(crate) limit_permits: u32,
44    pub(crate) granularity: PermitGranularity,
45    pub(crate) metrics: M,
46}
47
48impl<M: MemoryMetrics> MemoryManager<M> {
49    /// Creates a new memory manager with the given limit in bytes.
50    /// `limit_bytes = 0` disables the limit.
51    pub fn new(limit_bytes: u64, metrics: M) -> Self {
52        Self::with_granularity(limit_bytes, PermitGranularity::default(), metrics)
53    }
54
55    /// Creates a new memory manager with specified granularity.
56    pub fn with_granularity(limit_bytes: u64, granularity: PermitGranularity, metrics: M) -> Self {
57        if limit_bytes == 0 {
58            metrics.set_limit(0);
59            return Self { quota: None };
60        }
61
62        let limit_permits = granularity.bytes_to_permits(limit_bytes);
63        let limit_aligned_bytes = granularity.permits_to_bytes(limit_permits);
64        metrics.set_limit(limit_aligned_bytes as i64);
65
66        Self {
67            quota: Some(MemoryQuota {
68                semaphore: Arc::new(Semaphore::new(limit_permits as usize)),
69                limit_permits,
70                granularity,
71                metrics,
72            }),
73        }
74    }
75
76    /// Returns the configured limit in bytes (0 if unlimited).
77    pub fn limit_bytes(&self) -> u64 {
78        self.quota
79            .as_ref()
80            .map(|quota| quota.permits_to_bytes(quota.limit_permits))
81            .unwrap_or(0)
82    }
83
84    /// Returns currently used bytes.
85    pub fn used_bytes(&self) -> u64 {
86        self.quota
87            .as_ref()
88            .map(|quota| quota.permits_to_bytes(quota.used_permits()))
89            .unwrap_or(0)
90    }
91
92    /// Returns available bytes.
93    pub fn available_bytes(&self) -> u64 {
94        self.quota
95            .as_ref()
96            .map(|quota| quota.permits_to_bytes(quota.available_permits_clamped()))
97            .unwrap_or(0)
98    }
99
100    /// Acquires memory, waiting if necessary until enough is available.
101    ///
102    /// # Errors
103    /// - Returns error if requested bytes exceed the total limit
104    /// - Returns error if the semaphore is unexpectedly closed
105    pub async fn acquire(&self, bytes: u64) -> Result<MemoryGuard<M>> {
106        match &self.quota {
107            None => Ok(MemoryGuard::unlimited()),
108            Some(quota) => {
109                let permits = quota.bytes_to_permits(bytes);
110
111                ensure!(
112                    permits <= quota.limit_permits,
113                    MemoryLimitExceededSnafu {
114                        requested_bytes: bytes,
115                        limit_bytes: self.limit_bytes()
116                    }
117                );
118
119                let permit = quota
120                    .semaphore
121                    .clone()
122                    .acquire_many_owned(permits)
123                    .await
124                    .map_err(|_| MemorySemaphoreClosedSnafu.build())?;
125                quota.update_in_use_metric();
126                Ok(MemoryGuard::limited(permit, quota.clone()))
127            }
128        }
129    }
130
131    /// Tries to acquire memory. Returns Some(guard) on success, None if insufficient.
132    pub fn try_acquire(&self, bytes: u64) -> Option<MemoryGuard<M>> {
133        match &self.quota {
134            None => Some(MemoryGuard::unlimited()),
135            Some(quota) => {
136                let permits = quota.bytes_to_permits(bytes);
137
138                match quota.semaphore.clone().try_acquire_many_owned(permits) {
139                    Ok(permit) => {
140                        quota.update_in_use_metric();
141                        Some(MemoryGuard::limited(permit, quota.clone()))
142                    }
143                    Err(TryAcquireError::NoPermits) | Err(TryAcquireError::Closed) => {
144                        quota.metrics.inc_rejected("try_acquire");
145                        None
146                    }
147                }
148            }
149        }
150    }
151
152    /// Acquires memory based on the given policy.
153    ///
154    /// - For `OnExhaustedPolicy::Wait`: Waits up to the timeout duration for memory to become available
155    /// - For `OnExhaustedPolicy::Fail`: Returns immediately if memory is not available
156    ///
157    /// # Errors
158    /// - `MemoryLimitExceeded`: Requested bytes exceed the total limit (both policies), or memory is currently exhausted (Fail policy only)
159    /// - `MemoryAcquireTimeout`: Timeout elapsed while waiting for memory (Wait policy only)
160    /// - `MemorySemaphoreClosed`: The internal semaphore is unexpectedly closed (rare, indicates system issue)
161    pub async fn acquire_with_policy(
162        &self,
163        bytes: u64,
164        policy: OnExhaustedPolicy,
165    ) -> Result<MemoryGuard<M>> {
166        match policy {
167            OnExhaustedPolicy::Wait { timeout } => {
168                match tokio::time::timeout(timeout, self.acquire(bytes)).await {
169                    Ok(Ok(guard)) => Ok(guard),
170                    Ok(Err(e)) => Err(e),
171                    Err(_elapsed) => {
172                        // Timeout elapsed while waiting
173                        MemoryAcquireTimeoutSnafu {
174                            requested_bytes: bytes,
175                            waited: timeout,
176                        }
177                        .fail()
178                    }
179                }
180            }
181            OnExhaustedPolicy::Fail => self.try_acquire(bytes).ok_or_else(|| {
182                MemoryLimitExceededSnafu {
183                    requested_bytes: bytes,
184                    limit_bytes: self.limit_bytes(),
185                }
186                .build()
187            }),
188        }
189    }
190}
191
192impl<M: MemoryMetrics> MemoryQuota<M> {
193    pub(crate) fn bytes_to_permits(&self, bytes: u64) -> u32 {
194        self.granularity.bytes_to_permits(bytes)
195    }
196
197    pub(crate) fn permits_to_bytes(&self, permits: u32) -> u64 {
198        self.granularity.permits_to_bytes(permits)
199    }
200
201    pub(crate) fn used_permits(&self) -> u32 {
202        self.limit_permits
203            .saturating_sub(self.available_permits_clamped())
204    }
205
206    pub(crate) fn available_permits_clamped(&self) -> u32 {
207        self.semaphore
208            .available_permits()
209            .min(self.limit_permits as usize) as u32
210    }
211
212    pub(crate) fn update_in_use_metric(&self) {
213        let bytes = self.permits_to_bytes(self.used_permits());
214        self.metrics.set_in_use(bytes as i64);
215    }
216}
common_memory_manager/manager.rs

common_memory_manager/
manager.rs
