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;
16use std::sync::atomic::{AtomicU64, Ordering};
17
18use snafu::ensure;
19use tokio::sync::{OwnedSemaphorePermit, Semaphore, TryAcquireError};
20
21use crate::error::{
22    MemoryAcquireTimeoutSnafu, MemoryLimitExceededSnafu, MemorySemaphoreClosedSnafu, Result,
23};
24use crate::granularity::PermitGranularity;
25use crate::guard::MemoryGuard;
26use crate::policy::OnExhaustedPolicy;
27
28/// Trait for recording memory usage metrics.
29pub trait MemoryMetrics: Clone + Send + Sync + 'static {
30    fn set_limit(&self, bytes: i64);
31    fn set_in_use(&self, bytes: i64);
32    fn inc_rejected(&self, reason: &str);
33}
34
35/// Generic memory manager for quota-controlled operations.
36#[derive(Clone)]
37pub struct MemoryManager<M: MemoryMetrics> {
38    quota: MemoryQuotaState<M>,
39}
40
41impl<M: MemoryMetrics + Default> Default for MemoryManager<M> {
42    fn default() -> Self {
43        Self::new(0, M::default())
44    }
45}
46
47#[derive(Clone)]
48pub(crate) struct MemoryQuota<M: MemoryMetrics> {
49    pub(crate) semaphore: Arc<Semaphore>,
50    pub(crate) limit_permits: u32,
51    pub(crate) granularity: PermitGranularity,
52    pub(crate) metrics: M,
53}
54
55#[derive(Clone)]
56pub(crate) struct UnlimitedMemoryQuota<M: MemoryMetrics> {
57    pub(crate) current_bytes: Arc<AtomicU64>,
58    pub(crate) metrics: M,
59}
60
61#[derive(Clone)]
62pub(crate) enum MemoryQuotaState<M: MemoryMetrics> {
63    Unlimited(UnlimitedMemoryQuota<M>),
64    Limited(MemoryQuota<M>),
65}
66
67impl<M: MemoryMetrics> MemoryManager<M> {
68    /// Creates a new memory manager with the given limit in bytes.
69    /// `limit_bytes = 0` disables the limit.
70    pub fn new(limit_bytes: u64, metrics: M) -> Self {
71        Self::with_granularity(limit_bytes, PermitGranularity::default(), metrics)
72    }
73
74    /// Creates a new memory manager with specified granularity.
75    pub fn with_granularity(limit_bytes: u64, granularity: PermitGranularity, metrics: M) -> Self {
76        if limit_bytes == 0 {
77            metrics.set_limit(0);
78            return Self {
79                quota: MemoryQuotaState::Unlimited(UnlimitedMemoryQuota {
80                    current_bytes: Arc::new(AtomicU64::new(0)),
81                    metrics,
82                }),
83            };
84        }
85
86        let limit_permits = granularity.bytes_to_permits(limit_bytes);
87        let limit_aligned_bytes = granularity.permits_to_bytes(limit_permits);
88        metrics.set_limit(limit_aligned_bytes as i64);
89
90        Self {
91            quota: MemoryQuotaState::Limited(MemoryQuota {
92                semaphore: Arc::new(Semaphore::new(limit_permits as usize)),
93                limit_permits,
94                granularity,
95                metrics,
96            }),
97        }
98    }
99
100    /// Returns the configured limit in bytes (0 if unlimited).
101    pub fn limit_bytes(&self) -> u64 {
102        match &self.quota {
103            MemoryQuotaState::Unlimited(_) => 0,
104            MemoryQuotaState::Limited(quota) => quota.permits_to_bytes(quota.limit_permits),
105        }
106    }
107
108    /// Returns currently used bytes.
109    pub fn used_bytes(&self) -> u64 {
110        match &self.quota {
111            MemoryQuotaState::Unlimited(quota) => quota.current_bytes.load(Ordering::Acquire),
112            MemoryQuotaState::Limited(quota) => quota.permits_to_bytes(quota.used_permits()),
113        }
114    }
115
116    /// Returns available bytes.
117    ///
118    /// Unlimited managers report `u64::MAX`.
119    pub fn available_bytes(&self) -> u64 {
120        match &self.quota {
121            MemoryQuotaState::Unlimited(_) => u64::MAX,
122            MemoryQuotaState::Limited(quota) => {
123                quota.permits_to_bytes(quota.available_permits_clamped())
124            }
125        }
126    }
127
128    /// Acquires memory, waiting if necessary until enough is available.
129    ///
130    /// # Errors
131    /// - Returns error if requested bytes exceed the total limit
132    /// - Returns error if the semaphore is unexpectedly closed
133    pub async fn acquire(&self, bytes: u64) -> Result<MemoryGuard<M>> {
134        match &self.quota {
135            MemoryQuotaState::Unlimited(quota) => Ok(MemoryGuard::unlimited(quota.clone(), bytes)),
136            MemoryQuotaState::Limited(quota) => {
137                let permits = quota.bytes_to_permits(bytes);
138
139                ensure!(
140                    permits <= quota.limit_permits,
141                    MemoryLimitExceededSnafu {
142                        requested_bytes: bytes,
143                        limit_bytes: self.limit_bytes()
144                    }
145                );
146
147                let permit = quota
148                    .semaphore
149                    .clone()
150                    .acquire_many_owned(permits)
151                    .await
152                    .map_err(|_| MemorySemaphoreClosedSnafu.build())?;
153                quota.update_in_use_metric();
154                Ok(MemoryGuard::limited(quota.clone(), permit))
155            }
156        }
157    }
158
159    /// Tries to acquire memory. Returns Some(guard) on success, None if insufficient.
160    pub fn try_acquire(&self, bytes: u64) -> Option<MemoryGuard<M>> {
161        match &self.quota {
162            MemoryQuotaState::Unlimited(quota) => {
163                Some(MemoryGuard::unlimited(quota.clone(), bytes))
164            }
165            MemoryQuotaState::Limited(quota) => {
166                let permits = quota.bytes_to_permits(bytes);
167
168                match quota.semaphore.clone().try_acquire_many_owned(permits) {
169                    Ok(permit) => {
170                        quota.update_in_use_metric();
171                        Some(MemoryGuard::limited(quota.clone(), permit))
172                    }
173                    Err(TryAcquireError::NoPermits) | Err(TryAcquireError::Closed) => {
174                        quota.metrics.inc_rejected("try_acquire");
175                        None
176                    }
177                }
178            }
179        }
180    }
181
182    /// Acquires memory based on the given policy.
183    ///
184    /// - For `OnExhaustedPolicy::Wait`: Waits up to the timeout duration for memory to become available
185    /// - For `OnExhaustedPolicy::Fail`: Returns immediately if memory is not available
186    ///
187    /// # Errors
188    /// - `MemoryLimitExceeded`: Requested bytes exceed the total limit (both policies), or memory is currently exhausted (Fail policy only)
189    /// - `MemoryAcquireTimeout`: Timeout elapsed while waiting for memory (Wait policy only)
190    /// - `MemorySemaphoreClosed`: The internal semaphore is unexpectedly closed (rare, indicates system issue)
191    pub async fn acquire_with_policy(
192        &self,
193        bytes: u64,
194        policy: OnExhaustedPolicy,
195    ) -> Result<MemoryGuard<M>> {
196        match policy {
197            OnExhaustedPolicy::Wait { timeout } => {
198                match tokio::time::timeout(timeout, self.acquire(bytes)).await {
199                    Ok(Ok(guard)) => Ok(guard),
200                    Ok(Err(e)) => Err(e),
201                    Err(_elapsed) => {
202                        // Timeout elapsed while waiting
203                        MemoryAcquireTimeoutSnafu {
204                            requested_bytes: bytes,
205                            waited: timeout,
206                        }
207                        .fail()
208                    }
209                }
210            }
211            OnExhaustedPolicy::Fail => self.try_acquire(bytes).ok_or_else(|| {
212                MemoryLimitExceededSnafu {
213                    requested_bytes: bytes,
214                    limit_bytes: self.limit_bytes(),
215                }
216                .build()
217            }),
218        }
219    }
220}
221
222impl<M: MemoryMetrics> MemoryQuota<M> {
223    pub(crate) fn bytes_to_permits(&self, bytes: u64) -> u32 {
224        self.granularity.bytes_to_permits(bytes)
225    }
226
227    pub(crate) fn permits_to_bytes(&self, permits: u32) -> u64 {
228        self.granularity.permits_to_bytes(permits)
229    }
230
231    pub(crate) fn used_permits(&self) -> u32 {
232        self.limit_permits
233            .saturating_sub(self.available_permits_clamped())
234    }
235
236    pub(crate) fn available_permits_clamped(&self) -> u32 {
237        self.semaphore
238            .available_permits()
239            .min(self.limit_permits as usize) as u32
240    }
241
242    pub(crate) fn update_in_use_metric(&self) {
243        let bytes = self.permits_to_bytes(self.used_permits());
244        self.metrics.set_in_use(bytes as i64);
245    }
246
247    pub(crate) fn release_permit(&self, permit: OwnedSemaphorePermit) {
248        drop(permit);
249        self.update_in_use_metric();
250    }
251}
252
253impl<M: MemoryMetrics> UnlimitedMemoryQuota<M> {
254    pub(crate) fn add_in_use(&self, bytes: u64) {
255        if bytes == 0 {
256            return;
257        }
258
259        let previous = self
260            .current_bytes
261            .fetch_update(Ordering::AcqRel, Ordering::Acquire, |current| {
262                Some(current.saturating_add(bytes))
263            })
264            .unwrap();
265        let new_total = previous.saturating_add(bytes);
266        debug_assert!(
267            new_total >= previous,
268            "unlimited memory usage counter overflowed"
269        );
270        self.metrics.set_in_use(new_total as i64);
271    }
272
273    pub(crate) fn sub_in_use(&self, bytes: u64) {
274        if bytes == 0 {
275            return;
276        }
277
278        let previous = self
279            .current_bytes
280            .fetch_update(Ordering::AcqRel, Ordering::Acquire, |current| {
281                Some(current.saturating_sub(bytes))
282            })
283            .unwrap();
284        debug_assert!(
285            previous >= bytes,
286            "unlimited memory usage counter underflowed: current={previous}, release={bytes}"
287        );
288        let new_total = previous.saturating_sub(bytes);
289        self.metrics.set_in_use(new_total as i64);
290    }
291}
common_memory_manager/manager.rs

common_memory_manager/
manager.rs
