risingwave_common/memory/
mem_context.rs

1// Copyright 2023 RisingWave Labs
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::ops::Deref;
16use std::sync::Arc;
17
18use prometheus::core::Atomic;
19use risingwave_common_metrics::TrAdderAtomic;
20
21use super::MonitoredGlobalAlloc;
22use crate::metrics::{LabelGuardedIntGauge, TrAdderGauge};
23
24pub trait MemCounter: Send + Sync + 'static {
25    fn add(&self, bytes: i64);
26    fn get_bytes_used(&self) -> i64;
27}
28
29impl MemCounter for TrAdderGauge {
30    fn add(&self, bytes: i64) {
31        self.add(bytes)
32    }
33
34    fn get_bytes_used(&self) -> i64 {
35        self.get()
36    }
37}
38
39impl MemCounter for TrAdderAtomic {
40    fn add(&self, bytes: i64) {
41        self.inc_by(bytes)
42    }
43
44    fn get_bytes_used(&self) -> i64 {
45        self.get()
46    }
47}
48
49impl MemCounter for LabelGuardedIntGauge {
50    fn add(&self, bytes: i64) {
51        self.deref().add(bytes)
52    }
53
54    fn get_bytes_used(&self) -> i64 {
55        self.get()
56    }
57}
58
59struct MemoryContextInner {
60    counter: Box<dyn MemCounter>,
61    parent: Option<MemoryContext>,
62    mem_limit: u64,
63}
64
65#[derive(Clone)]
66pub struct MemoryContext {
67    /// Add None op mem context, so that we don't need to return [`Option`] in
68    /// `BatchTaskContext`. This helps with later `Allocator` implementation.
69    inner: Option<Arc<MemoryContextInner>>,
70}
71
72impl MemoryContext {
73    pub fn new(parent: Option<MemoryContext>, counter: impl MemCounter) -> Self {
74        let mem_limit = parent.as_ref().map_or_else(|| u64::MAX, |p| p.mem_limit());
75        Self::new_with_mem_limit(parent, counter, mem_limit)
76    }
77
78    pub fn new_with_mem_limit(
79        parent: Option<MemoryContext>,
80        counter: impl MemCounter,
81        mem_limit: u64,
82    ) -> Self {
83        let c = Box::new(counter);
84        Self {
85            inner: Some(Arc::new(MemoryContextInner {
86                counter: c,
87                parent,
88                mem_limit,
89            })),
90        }
91    }
92
93    /// Creates a noop memory context.
94    pub fn none() -> Self {
95        Self { inner: None }
96    }
97
98    pub fn root(counter: impl MemCounter, mem_limit: u64) -> Self {
99        Self::new_with_mem_limit(None, counter, mem_limit)
100    }
101
102    pub fn for_spill_test() -> Self {
103        Self::new_with_mem_limit(None, TrAdderAtomic::new(0), 0)
104    }
105
106    /// Attempts to charge `bytes` against this context and its ancestors.
107    /// Returns `false`, without updating counters, if a positive charge would exceed a limit.
108    /// Negative values release previously recorded usage, even while over budget.
109    /// Use [`Self::add_unchecked`] to record allocations that proceed regardless of the budget.
110    pub fn add(&self, bytes: i64) -> bool {
111        if let Some(inner) = &self.inner {
112            // Releasing memory must succeed even if concurrent admissions exceeded the limit.
113            if bytes > 0 && (inner.counter.get_bytes_used() + bytes) as u64 > inner.mem_limit {
114                return false;
115            }
116            if let Some(parent) = &inner.parent {
117                if parent.add(bytes) {
118                    inner.counter.add(bytes);
119                } else {
120                    return false;
121                }
122            } else {
123                inner.counter.add(bytes);
124            }
125        }
126        true
127    }
128
129    /// Records a memory-usage change without enforcing this context's or its ancestors' limits.
130    ///
131    /// Use this for allocations that have already succeeded and their corresponding releases.
132    /// Every recorded allocation must have a matching release. This does not allocate/free memory
133    /// or clamp the counter; callers remain responsible for balanced accounting.
134    /// For admission control, use [`Self::add`] instead. To inspect the resulting budget state,
135    /// use [`Self::check_memory_usage`].
136    pub fn add_unchecked(&self, bytes: i64) {
137        if let Some(inner) = &self.inner {
138            if let Some(parent) = &inner.parent {
139                parent.add_unchecked(bytes);
140            }
141            inner.counter.add(bytes);
142        }
143    }
144
145    pub fn get_bytes_used(&self) -> i64 {
146        if let Some(inner) = &self.inner {
147            inner.counter.get_bytes_used()
148        } else {
149            0
150        }
151    }
152
153    pub fn mem_limit(&self) -> u64 {
154        if let Some(inner) = &self.inner {
155            inner.mem_limit
156        } else {
157            u64::MAX
158        }
159    }
160
161    /// Check if the memory usage exceeds the limit.
162    /// Returns `false` if the memory usage exceeds the limit.
163    pub fn check_memory_usage(&self) -> bool {
164        if let Some(inner) = &self.inner {
165            if inner.counter.get_bytes_used() as u64 > inner.mem_limit {
166                return false;
167            }
168            if let Some(parent) = &inner.parent {
169                return parent.check_memory_usage();
170            }
171        }
172
173        true
174    }
175
176    /// Creates a new global allocator that reports memory usage to this context.
177    pub fn global_allocator(&self) -> MonitoredGlobalAlloc {
178        MonitoredGlobalAlloc::with_memory_context(self.clone())
179    }
180}
181
182impl Drop for MemoryContextInner {
183    fn drop(&mut self) {
184        if let Some(p) = &self.parent {
185            p.add_unchecked(-self.counter.get_bytes_used());
186        }
187    }
188}
risingwave_common/memory/mem_context.rs

risingwave_common/memory/
mem_context.rs
