-
Notifications
You must be signed in to change notification settings - Fork 13
Expand file tree
/
Copy pathmemory.rs
More file actions
208 lines (186 loc) · 8.47 KB
/
Copy pathmemory.rs
File metadata and controls
208 lines (186 loc) · 8.47 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
//! Session-wide memory accounting wrapper around any `MemoryPool`.
use std::collections::HashMap;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, Mutex, OnceLock};
use datafusion::common::Result;
use datafusion::execution::memory_pool::{MemoryLimit, MemoryPool, MemoryReservation};
use jni::sys::jlong;
/// Wraps an inner `MemoryPool` with two counters: total bytes currently held
/// and the peak observed since the wrapper was constructed. Bookkeeping is two
/// atomic loads/stores per `grow`/`shrink` and adds no contention because the
/// inner pool already serializes its own state.
#[derive(Debug)]
pub struct TrackingMemoryPool {
inner: Arc<dyn MemoryPool>,
current_bytes: AtomicU64,
peak_bytes: AtomicU64,
}
impl TrackingMemoryPool {
pub fn new(inner: Arc<dyn MemoryPool>) -> Self {
Self {
inner,
current_bytes: AtomicU64::new(0),
peak_bytes: AtomicU64::new(0),
}
}
// Note: there are intentionally no separate `current_bytes()` /
// `peak_bytes()` accessors. Reading them independently allows a torn
// observation where `current > peak` while a concurrent `record_grow` is
// mid-flight. Callers must use `snapshot()` instead, which clamps.
/// Returns `(current_bytes, peak_bytes)` as a consistent snapshot. Two
/// invariants are maintained:
///
/// 1. **Within a snapshot:** `peak >= current`. This matters when a
/// concurrent `record_grow` is in-flight between its `fetch_add` of
/// `current_bytes` and its `fetch_max` of `peak_bytes` -- a naive
/// snapshot could observe the higher current but the stale peak.
/// 2. **Across snapshots:** `peak` is monotonically non-decreasing for
/// every observer. If reader R1 ever sees peak = X, no later reader R2
/// (on any thread) sees peak < X.
///
/// The read-side clamp `max(stored_peak, observed_current)` alone gives
/// (1) but not (2): it's possible for R1 to *synthesize* a peak of X
/// from a transiently-high current, and then for a `record_shrink` plus
/// the slow path of `record_grow`'s missing `fetch_max` to leave the
/// stored peak briefly below X while R2 polls. R2 would observe a
/// smaller current and a smaller stored peak -- non-monotonic.
///
/// The fix is to persist the clamped value back into `peak_bytes` via
/// `fetch_max` *before* returning. Doing so guarantees every snapshot's
/// returned peak is also reflected in the global atomic, so any
/// subsequent observer's `peak_bytes.load()` is at least as high as the
/// largest peak any earlier observer returned.
///
/// Both atomics are `Relaxed`-ordered; we tolerate slightly stale values
/// in exchange for cheaper writes.
pub fn snapshot(&self) -> (u64, u64) {
// Order matters on the load side too: read `peak` first, then
// `current`. If a `record_grow` is in flight, this order ensures
// that whenever we observe a fresh `current` we also observe at
// least the corresponding (or older) `peak` -- so the eventual
// clamp `max(peak, current)` produces the right value.
let peak_observed = self.peak_bytes.load(Ordering::Relaxed);
let current = self.current_bytes.load(Ordering::Relaxed);
let clamped_peak = peak_observed.max(current);
// Persist the clamped peak so future observers can never see a
// smaller peak than the one we just returned. fetch_max is
// monotonic, so this is a no-op when another thread (or this very
// thread on a previous call) already pushed peak past clamped_peak.
let stored_peak = self.peak_bytes.fetch_max(clamped_peak, Ordering::Relaxed);
(current, stored_peak.max(clamped_peak))
}
fn record_grow(&self, additional: usize) {
// We bump `current` first to avoid double-counting in `peak` if a
// concurrent reader sees a stale current and a fresh peak. The
// snapshot path's read-side `max(peak, current)` clamp closes the
// window where `current` is up-to-date but `peak` is not.
let now = self
.current_bytes
.fetch_add(additional as u64, Ordering::Relaxed)
.saturating_add(additional as u64);
// fetch_max is monotonic so concurrent grows still produce the right peak.
self.peak_bytes.fetch_max(now, Ordering::Relaxed);
}
fn record_shrink(&self, shrink: usize) {
// Saturating: if the inner pool ever calls shrink with a value larger
// than the running total we'd rather report 0 than wrap around.
let prev = self
.current_bytes
.fetch_sub(shrink as u64, Ordering::Relaxed);
if prev < shrink as u64 {
// Restore to zero. Concurrent grows that interleaved are fine --
// they'll just see a transient under-count.
self.current_bytes.store(0, Ordering::Relaxed);
}
}
}
/// DataFusion 54 added `Display` as a supertrait of [`MemoryPool`] so pools
/// can render themselves in resource-exhausted messages. Mirror upstream's
/// wrapper convention: name the wrapper, then defer to the inner pool for the
/// limit/usage detail, and add the counters this wrapper exists to expose.
impl std::fmt::Display for TrackingMemoryPool {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
let (current, peak) = self.snapshot();
write!(
f,
"{}(current: {current}, peak: {peak}, inner: {})",
self.name(),
self.inner
)
}
}
impl MemoryPool for TrackingMemoryPool {
fn name(&self) -> &str {
"tracking"
}
fn register(&self, consumer: &datafusion::execution::memory_pool::MemoryConsumer) {
self.inner.register(consumer);
}
fn unregister(&self, consumer: &datafusion::execution::memory_pool::MemoryConsumer) {
self.inner.unregister(consumer);
}
fn grow(&self, reservation: &MemoryReservation, additional: usize) {
self.inner.grow(reservation, additional);
self.record_grow(additional);
}
fn shrink(&self, reservation: &MemoryReservation, shrink: usize) {
self.inner.shrink(reservation, shrink);
self.record_shrink(shrink);
}
fn try_grow(&self, reservation: &MemoryReservation, additional: usize) -> Result<()> {
self.inner.try_grow(reservation, additional)?;
self.record_grow(additional);
Ok(())
}
fn reserved(&self) -> usize {
self.inner.reserved()
}
fn memory_limit(&self) -> MemoryLimit {
self.inner.memory_limit()
}
}
/// Process-wide registry mapping a `SessionContext`'s native handle to its
/// `TrackingMemoryPool`. Populated at session creation, drained at session
/// close. We can't downcast `Arc<dyn MemoryPool>` back to the concrete type
/// without an `Any` bound (which the trait does not require), so instead we
/// keep a parallel `Arc` indexed by the JNI handle Java already threads
/// through every call.
fn registry() -> &'static Mutex<HashMap<jlong, Arc<TrackingMemoryPool>>> {
static REGISTRY: OnceLock<Mutex<HashMap<jlong, Arc<TrackingMemoryPool>>>> = OnceLock::new();
REGISTRY.get_or_init(|| Mutex::new(HashMap::new()))
}
pub fn register(handle: jlong, pool: Arc<TrackingMemoryPool>) {
registry()
.lock()
.expect("memory registry lock poisoned")
.insert(handle, pool);
}
pub fn lookup(handle: jlong) -> Option<Arc<TrackingMemoryPool>> {
registry()
.lock()
.expect("memory registry lock poisoned")
.get(&handle)
.cloned()
}
pub fn unregister(handle: jlong) {
registry()
.lock()
.expect("memory registry lock poisoned")
.remove(&handle);
}