diff --git a/src/core/src/cache/budget.rs b/src/core/src/cache/budget.rs index 2b7bc5d0..ef2ca9fa 100644 --- a/src/core/src/cache/budget.rs +++ b/src/core/src/cache/budget.rs @@ -9,10 +9,55 @@ pub struct BudgetAccounting { max_memory_bytes: usize, max_disk_bytes: usize, used_memory_bytes: AtomicUsize, + in_flight_memory_bytes: AtomicUsize, + peak_in_flight_memory_bytes: AtomicUsize, used_disk_bytes: AtomicUsize, observer: Arc, } +/// Bytes that are materialized in memory but not yet indexed. +/// +/// The cache holds one of these for every intermediate the +/// hydrate -> insert -> squeeze cycle creates outside the index: a disk entry +/// being decoded, a squeeze output waiting to be written and inserted, and an +/// entry pending admission while room is made for it. Nothing counted them +/// before, which is why a tier reporting itself at its limit could sit inside a +/// process holding several times that. +/// +/// These bytes do not gate admission: they already exist by the time their size +/// is known, so refusing them would free nothing. They are reported, and they +/// bound how many transcodes [`super::core::LiquidCache`] runs at once. +#[derive(Debug)] +#[must_use = "dropping the reservation immediately releases the bytes"] +pub(super) struct InFlightReservation<'a> { + budget: &'a BudgetAccounting, + bytes: usize, +} + +impl InFlightReservation<'_> { + /// Change the reserved amount, for when the true size is only known part + /// way through: a disk read reserves the encoded buffer before decoding and + /// resizes to the decoded array, a transcode reserves its input size as an + /// upper bound and resizes to the compressed result. + pub(super) fn resize(&mut self, bytes: usize) { + self.budget.release_in_flight(self.bytes); + self.budget.reserve_in_flight_bytes(bytes); + self.bytes = bytes; + } + + /// Stop counting these bytes as in-flight, because the caller is about to + /// account for them another way: by inserting them into the index. + pub(super) fn release(self) { + drop(self); + } +} + +impl Drop for InFlightReservation<'_> { + fn drop(&mut self) { + self.budget.release_in_flight(self.bytes); + } +} + impl BudgetAccounting { pub(super) fn new( max_memory_bytes: usize, @@ -23,6 +68,8 @@ impl BudgetAccounting { max_memory_bytes, max_disk_bytes, used_memory_bytes: AtomicUsize::new(0), + in_flight_memory_bytes: AtomicUsize::new(0), + peak_in_flight_memory_bytes: AtomicUsize::new(0), used_disk_bytes: AtomicUsize::new(0), observer, } @@ -31,6 +78,7 @@ impl BudgetAccounting { pub(super) fn reset_usage(&self) { self.used_memory_bytes.store(0, Ordering::Relaxed); self.used_disk_bytes.store(0, Ordering::Relaxed); + self.peak_in_flight_memory_bytes.store(0, Ordering::Relaxed); } /// Try to reserve memory in the cache. @@ -70,10 +118,53 @@ impl BudgetAccounting { } } + /// Reserve bytes that are live in memory but not yet in the index. + /// + /// Cannot fail: see [`InFlightReservation`]. + pub(super) fn reserve_in_flight(&self, bytes: usize) -> InFlightReservation<'_> { + self.reserve_in_flight_bytes(bytes); + InFlightReservation { + budget: self, + bytes, + } + } + + fn reserve_in_flight_bytes(&self, bytes: usize) { + let total = self + .in_flight_memory_bytes + .fetch_add(bytes, Ordering::Relaxed) + + bytes; + self.peak_in_flight_memory_bytes + .fetch_max(total, Ordering::Relaxed); + } + + fn release_in_flight(&self, bytes: usize) { + self.in_flight_memory_bytes + .fetch_sub(bytes, Ordering::Relaxed); + } + + /// Bytes held by the index. pub fn memory_usage_bytes(&self) -> usize { self.used_memory_bytes.load(Ordering::Relaxed) } + /// Bytes materialized in memory right now but not yet in the index. + /// + /// Read alongside [`Self::memory_usage_bytes`]: that one reports the end + /// state of the hydrate -> insert -> squeeze cycle, this one reports what + /// the cycle is holding on the way there. + pub fn in_flight_memory_bytes(&self) -> usize { + self.in_flight_memory_bytes.load(Ordering::Relaxed) + } + + /// High water mark of [`Self::in_flight_memory_bytes`]. + /// + /// Transients live and die between two scrapes of a gauge, so this is the + /// number to look at when a process holds more than its tier reports. + pub fn peak_in_flight_memory_bytes(&self) -> usize { + self.peak_in_flight_memory_bytes.load(Ordering::Relaxed) + } + pub fn disk_usage_bytes(&self) -> usize { self.used_disk_bytes.load(Ordering::Relaxed) } @@ -129,6 +220,34 @@ mod tests { assert_eq!(config.memory_usage_bytes(), 0); } + #[test] + fn in_flight_reservations_are_reported_and_released() { + let budget = test_budget(1000, usize::MAX); + + let small = budget.reserve_in_flight(100); + let mut large = budget.reserve_in_flight(300); + assert_eq!(budget.in_flight_memory_bytes(), 400); + + // A reservation taken as an upper bound is trued up once the real size + // is known, and the peak remembers the bound that was held. + large.resize(50); + assert_eq!(budget.in_flight_memory_bytes(), 150); + assert_eq!(budget.peak_in_flight_memory_bytes(), 400); + + // In-flight bytes are reported, not charged: they say what the cache is + // holding outside the index, and admission does not consult them. + assert!(budget.try_reserve_memory(1000).is_ok()); + + drop(large); + small.release(); + assert_eq!(budget.in_flight_memory_bytes(), 0); + assert_eq!( + budget.peak_in_flight_memory_bytes(), + 400, + "the high water mark outlives the reservations that set it" + ); + } + #[test] fn test_concurrent_memory_operations() { test_concurrent_memory_budget(); diff --git a/src/core/src/cache/core.rs b/src/core/src/cache/core.rs index f26f2877..26655de0 100644 --- a/src/core/src/cache/core.rs +++ b/src/core/src/cache/core.rs @@ -7,7 +7,7 @@ use bytes::Bytes; use futures::StreamExt; use super::{ - budget::BudgetAccounting, + budget::{BudgetAccounting, InFlightReservation}, builders::{EvaluatePredicate, Get, Insert}, cached_batch::{CacheEntry, CachedBatchType}, io_context::{EntryMetadata, entry_id_to_key}, @@ -111,6 +111,8 @@ impl LiquidCache { memory_liquid_bytes, memory_squeezed_liquid_bytes, memory_usage_bytes, + in_flight_memory_bytes: self.budget.in_flight_memory_bytes(), + peak_in_flight_memory_bytes: self.budget.peak_in_flight_memory_bytes(), disk_usage_bytes, max_memory_bytes: self.config.max_memory_bytes(), max_disk_bytes: self.config.max_disk_bytes(), @@ -156,14 +158,15 @@ impl LiquidCache { match batch.as_ref() { CacheEntry::MemoryLiquid(array) => Some(array.clone()), entry @ CacheEntry::DiskLiquid { .. } => { - let liquid = self.read_disk_liquid_array(entry_id).await; + let (liquid, reservation) = self.read_disk_liquid_array(entry_id).await; + reservation.release(); self.maybe_hydrate(entry_id, entry, MaterializedEntry::Liquid(&liquid), None) .await; Some(liquid) } CacheEntry::MemorySqueezedLiquid(array) => match array.disk_backing() { SqueezedBacking::Liquid(_) => { - let liquid = self.read_disk_liquid_array(entry_id).await; + let (liquid, _reservation) = self.read_disk_liquid_array(entry_id).await; Some(liquid) } SqueezedBacking::Arrow(_) => None, @@ -350,6 +353,13 @@ impl LiquidCache { kind: CachedBatchType::from(¬_inserted), }); + // The entry is materialized but not indexed for as long as we are + // making room for it, and that room-making is itself what holds the + // most memory. Count it while it waits. + let pending = self + .budget + .reserve_in_flight(not_inserted.memory_usage_bytes()); + let victims = self.cache_policy.find_memory_victim(8); if victims.is_empty() { // no advice, because the cache is already empty @@ -358,10 +368,12 @@ impl LiquidCache { let on_disk_batch = self .write_in_memory_batch_to_disk(entry_id, not_inserted) .await?; + pending.release(); batch_to_cache = on_disk_batch; continue; } self.squeeze_victims(victims).await?; + pending.release(); batch_to_cache = not_inserted; crate::utils::yield_now_if_shuttle(); @@ -495,12 +507,15 @@ impl LiquidCache { victims: victims.clone(), }); if self.squeeze_victims_concurrently { - let results = futures::stream::iter(victims) - .map(|victim| self.squeeze_victim_inner(victim)) - .buffer_unordered(usize::MAX) - .collect::>() - .await; - results.into_iter().collect::, _>>()?; + for group in self.group_victims_by_in_flight_bytes(victims) { + let concurrency = group.len(); + let results = futures::stream::iter(group) + .map(|victim| self.squeeze_victim_inner(victim)) + .buffer_unordered(concurrency) + .collect::>() + .await; + results.into_iter().collect::, _>>()?; + } } else { for victim in victims { self.squeeze_victim_inner(victim).await?; @@ -509,6 +524,55 @@ impl LiquidCache { Ok(()) } + /// Split victims into groups whose squeezes may be in flight together. + /// + /// Running a whole batch of victims concurrently puts every one of their + /// outputs in memory at once, which is unbounded in bytes however few + /// victims there are: eight 60 MB batches is half a gigabyte of transcode + /// output outside the index. A squeeze's output cannot exceed its input, so + /// each victim's indexed size bounds what it will add, and grouping by that + /// bound caps the pile-up at [`Self::squeeze_in_flight_limit`]. + /// + /// A group always contains at least one victim, so a victim larger than the + /// whole limit proceeds alone rather than stalling the cache. + fn group_victims_by_in_flight_bytes(&self, victims: Vec) -> Vec> { + // Other tasks may already be squeezing, so spend only what is left of + // the limit rather than the whole limit again. + let limit = self + .squeeze_in_flight_limit() + .saturating_sub(self.budget.in_flight_memory_bytes()); + + let mut groups: Vec> = Vec::new(); + let mut group_bytes = 0usize; + for victim in victims { + let victim_bytes = self + .index + .get(&victim) + .map_or(0, |entry| entry.memory_usage_bytes()); + match groups.last_mut() { + Some(group) if group_bytes + victim_bytes <= limit => { + group_bytes += victim_bytes; + group.push(victim); + } + _ => { + group_bytes = victim_bytes; + groups.push(vec![victim]); + } + } + } + groups + } + + /// Ceiling on the squeeze output the cache lets accumulate outside the + /// index at one time, as a fraction of the memory tier. + /// + /// The fraction, rather than a constant, is what keeps this meaningful + /// across tier sizes: the transients of making room have to be small + /// against the room being made. + fn squeeze_in_flight_limit(&self) -> usize { + self.config.max_memory_bytes() / 32 + } + async fn squeeze_victim_inner(&self, to_squeeze: EntryID) -> Result<(), CacheFull> { let Some(mut to_squeeze_batch) = self.index.get(&to_squeeze) else { return Ok(()); @@ -524,6 +588,14 @@ impl LiquidCache { )); loop { + // A squeeze cannot produce more bytes than it consumes, so the + // entry's indexed size is an upper bound on the compressed output + // and the disk buffer the squeeze is about to hold outside the + // index. Reserve the bound first, then true it up to the result. + let mut reservation = self + .budget + .reserve_in_flight(to_squeeze_batch.memory_usage_bytes()); + let outcome = self.squeeze_policy.squeeze( to_squeeze_batch.as_ref(), compressor.as_ref(), @@ -536,10 +608,16 @@ impl LiquidCache { entry: new_batch, bytes_to_write, } => { + reservation.resize( + new_batch.memory_usage_bytes() + + bytes_to_write.as_ref().map_or(0, Bytes::len), + ); if let Some(bytes_to_write) = bytes_to_write { self.write_batch_to_disk(to_squeeze, &new_batch, bytes_to_write) .await?; } + // About to be indexed, so stop counting it as a transient. + reservation.release(); match self.try_insert(to_squeeze, new_batch) { Ok(()) => { break; @@ -550,6 +628,7 @@ impl LiquidCache { } } SqueezeOutcome::Remove => { + reservation.release(); self.remove_disk_entry(to_squeeze).await; break; } @@ -647,7 +726,8 @@ impl LiquidCache { { return Some(arrow::array::new_empty_array(data_type)); } - let full_array = self.read_disk_arrow_array(entry_id).await; + let (full_array, reservation) = self.read_disk_arrow_array(entry_id).await; + reservation.release(); self.maybe_hydrate( entry_id, entry, @@ -669,7 +749,8 @@ impl LiquidCache { { return Some(arrow::array::new_empty_array(data_type)); } - let liquid = self.read_disk_liquid_array(entry_id).await; + let (liquid, reservation) = self.read_disk_liquid_array(entry_id).await; + reservation.release(); self.maybe_hydrate( entry_id, entry, @@ -762,7 +843,8 @@ impl LiquidCache { let full_array = if !all_paths_present { let batch = CacheEntry::MemorySqueezedLiquid(array.clone()); self.observer.on_get_squeezed_needs_io(); - let full_array = self.read_disk_arrow_array(entry_id).await; + let (full_array, reservation) = self.read_disk_arrow_array(entry_id).await; + reservation.release(); self.maybe_hydrate( entry_id, &batch, @@ -818,29 +900,47 @@ impl LiquidCache { Ok(()) } - async fn read_disk_arrow_array(&self, entry_id: &EntryID) -> ArrayRef { + /// Read an on-disk Arrow entry, reserving what the decode holds outside the + /// index. + /// + /// The reservation comes back with the array because the caller is the one + /// that knows when these bytes stop being a transient: it releases just + /// before hydrating them, which counts the same bytes as indexed instead. + async fn read_disk_arrow_array( + &self, + entry_id: &EntryID, + ) -> (ArrayRef, InFlightReservation<'_>) { let bytes = self .store .get(&entry_id_to_key(entry_id)) .await .expect("read failed"); let bytes_len = bytes.len(); + // The encoded buffer is already resident; the decode adds the array on + // top of it, and the IPC reader holds both until this call returns. + let mut reservation = self.budget.reserve_in_flight(bytes_len); let cursor = std::io::Cursor::new(bytes); let mut reader = arrow::ipc::reader::StreamReader::try_new(cursor, None).expect("create reader failed"); let batch = reader.next().unwrap().expect("read batch failed"); let array = batch.column(0).clone(); + // Arrow IPC copies rather than slicing, so the encoded buffer and the + // decoded array are both resident here. + reservation.resize(bytes_len + array.get_array_memory_size()); self.trace(InternalEvent::IoReadArrow { entry: *entry_id, bytes: bytes_len, }); - array + (array, reservation) } + /// Read an on-disk Liquid entry, reserving what the decode holds outside + /// the index. See [`Self::read_disk_arrow_array`] for the reservation's + /// contract. async fn read_disk_liquid_array( &self, entry_id: &EntryID, - ) -> crate::liquid_array::LiquidArrayRef { + ) -> (crate::liquid_array::LiquidArrayRef, InFlightReservation<'_>) { let bytes = self .store .get(&entry_id_to_key(entry_id)) @@ -850,13 +950,19 @@ impl LiquidCache { entry: *entry_id, bytes: bytes.len(), }); + let mut reservation = self.budget.reserve_in_flight(bytes.len()); let compressor_states = self.metadata.get_compressor(entry_id); let compressor = compressor_states.fsst_compressor(); - (crate::liquid_array::ipc::read_from_bytes( + let array = crate::liquid_array::ipc::read_from_bytes( Bytes::from(bytes), &crate::liquid_array::ipc::LiquidIPCContext::new(compressor), - )) as _ + ); + // A liquid array decodes zero-copy over the encoded buffer, so its + // reported size already covers what was reserved above rather than + // adding to it. + reservation.resize(array.get_array_memory_size()); + (array as _, reservation) } pub(crate) async fn eval_predicate_internal( @@ -890,7 +996,8 @@ impl LiquidCache { Some(self.eval_predicate_on_array(filtered, predicate)) } entry @ CacheEntry::DiskArrow { .. } => { - let array = self.read_disk_arrow_array(entry_id).await; + let (array, reservation) = self.read_disk_arrow_array(entry_id).await; + reservation.release(); self.maybe_hydrate(entry_id, entry, MaterializedEntry::Arrow(&array), None) .await; let mut owned = None; @@ -912,7 +1019,8 @@ impl LiquidCache { Some(array.try_eval_predicate(predicate, selection)) } entry @ CacheEntry::DiskLiquid { .. } => { - let liquid = self.read_disk_liquid_array(entry_id).await; + let (liquid, reservation) = self.read_disk_liquid_array(entry_id).await; + reservation.release(); self.maybe_hydrate(entry_id, entry, MaterializedEntry::Liquid(&liquid), None) .await; let mut owned = None; diff --git a/src/core/src/cache/observer/stats.rs b/src/core/src/cache/observer/stats.rs index fa0c3d9a..69cc5023 100644 --- a/src/core/src/cache/observer/stats.rs +++ b/src/core/src/cache/observer/stats.rs @@ -142,8 +142,17 @@ pub struct CacheStats { pub memory_liquid_bytes: usize, /// Total size of in-memory Squeezed-Liquid entries in bytes. pub memory_squeezed_liquid_bytes: usize, - /// Total memory usage of the cache. + /// Total memory usage of the cache: the bytes held by the index. pub memory_usage_bytes: usize, + /// Bytes materialized in memory but not yet indexed: entries being decoded + /// off disk, squeeze outputs awaiting insertion, and entries pending + /// admission. Read with `memory_usage_bytes`, which covers only what the + /// hydrate/insert/squeeze cycle has already landed. + pub in_flight_memory_bytes: usize, + /// High water mark of `in_flight_memory_bytes`, since the last cache reset. + /// Transients come and go between gauge scrapes, so this is the figure that + /// explains a process holding more than `memory_usage_bytes` reports. + pub peak_in_flight_memory_bytes: usize, /// Total disk usage of the cache. pub disk_usage_bytes: usize, /// Maximum memory size. diff --git a/src/core/src/cache/tests/in_flight.rs b/src/core/src/cache/tests/in_flight.rs new file mode 100644 index 00000000..b9612f81 --- /dev/null +++ b/src/core/src/cache/tests/in_flight.rs @@ -0,0 +1,176 @@ +//! What the memory tier holds outside its index, and what bounds it. +//! +//! The budget used to count an entry only once it had landed in the index, so +//! every intermediate the hydrate -> insert -> squeeze cycle created on the way +//! there was invisible to it. A tier could report itself exactly at its limit +//! while the process holding it had several times that resident. + +use crate::cache::{ + CacheEntry, EntryID, LiquidCache, LiquidCacheBuilder, LiquidCompressorStates, LiquidPolicy, + TranscodeSqueezeEvict, transcode_liquid_inner, + utils::{create_cache_store, create_test_arrow_array}, +}; +use crate::sync::Arc; + +/// A cache whose memory tier holds exactly `entries` arrays of `rows` rows, so +/// that the next insert has to make room for itself. +async fn cache_holding( + entries: usize, + rows: usize, + squeeze_concurrently: bool, +) -> Arc { + let entry_bytes = create_test_arrow_array(rows).get_array_memory_size(); + LiquidCacheBuilder::new() + .with_max_memory_bytes(entry_bytes * entries) + .with_cache_policy(Box::new(LiquidPolicy::new())) + .with_squeeze_policy(Box::new(TranscodeSqueezeEvict)) + .with_squeeze_victims_concurrently(squeeze_concurrently) + .build() + .await +} + +/// Sum of what the indexed entries actually hold in memory, which is what the +/// tier's tally is supposed to be tracking. +fn indexed_memory_bytes(cache: &LiquidCache) -> usize { + let mut total = 0; + cache.for_each_entry(|_, entry| total += entry.memory_usage_bytes()); + total +} + +#[tokio::test] +async fn decoding_a_disk_entry_is_counted_while_it_is_in_flight() { + let cache = create_cache_store(1 << 20, Box::new(LiquidPolicy::new())).await; + let entry_id = EntryID::from(700usize); + let array = create_test_arrow_array(4096); + cache.insert(entry_id, array.clone()).await.unwrap(); + cache.flush_all_to_disk().await.unwrap(); + + let at_rest = cache.stats(); + assert_eq!(at_rest.in_flight_memory_bytes, 0); + assert_eq!( + at_rest.peak_in_flight_memory_bytes, 0, + "an insert that fits holds nothing outside the index" + ); + + cache.get(&entry_id).read().await.unwrap(); + + let stats = cache.stats(); + assert_eq!( + stats.in_flight_memory_bytes, 0, + "every reservation is released by the time the read returns" + ); + assert!( + stats.peak_in_flight_memory_bytes >= array.get_array_memory_size(), + "reading the entry back off disk decodes a full copy of it, so the peak \ + should be at least the entry's size, but it was {} against {}", + stats.peak_in_flight_memory_bytes, + array.get_array_memory_size() + ); +} + +/// Victims already in liquid form are the ones that pile up: squeezing them +/// writes their backing to disk, and that write is an await, so every victim in +/// a concurrently squeezed group is holding its output while the others run. +/// (An arrow victim transcodes without ever awaiting, so its future runs to +/// completion in a single poll and it never overlaps with a sibling.) +#[tokio::test] +async fn squeezing_victims_concurrently_bounds_what_they_hold_at_once() { + let rows = 4096; + let array = create_test_arrow_array(rows); + let compressor = LiquidCompressorStates::new(); + let liquid = transcode_liquid_inner(&array, &compressor).expect("int64 transcodes"); + let entry_bytes = liquid.get_array_memory_size(); + + let victims = 8; + let cache = LiquidCacheBuilder::new() + .with_max_memory_bytes(entry_bytes * victims) + .with_cache_policy(Box::new(LiquidPolicy::new())) + .with_squeeze_policy(Box::new(TranscodeSqueezeEvict)) + .with_squeeze_victims_concurrently(true) + .build() + .await; + + for i in 0..victims { + cache + .insert_inner(EntryID::from(i), CacheEntry::memory_liquid(liquid.clone())) + .await + .unwrap(); + } + assert_eq!( + cache.stats().peak_in_flight_memory_bytes, + 0, + "the tier is full but nothing has had to make room yet" + ); + + // This insert finds no room, so the policy hands back all eight resident + // entries and each is squeezed to disk to make space. + cache + .insert_inner( + EntryID::from(victims), + CacheEntry::memory_liquid(liquid.clone()), + ) + .await + .unwrap(); + + let peak = cache.stats().peak_in_flight_memory_bytes; + assert!( + peak <= 3 * entry_bytes, + "a squeeze holds at most its input's worth of output, so the pending \ + entry plus one victim bounds this; all eight victims at once would peak \ + near {}, and this run peaked at {peak} against an entry size of \ + {entry_bytes}", + (victims + 1) * entry_bytes + ); +} + +/// The path that ships is the concurrent one (`builders.rs` turns it on +/// everywhere except unit tests), so it gets the same invariants the sequential +/// path has always been held to: nothing is lost, nothing is miscounted, and +/// nothing is left reserved. +#[tokio::test] +async fn squeezing_victims_holds_its_invariants_on_both_paths() { + // A tier this much larger than a single entry lets several victims share a + // group, so the concurrent path really does run transcodes together instead + // of degenerating into the sequential one. + let rows = 64; + let array = create_test_arrow_array(rows); + let entries = 300; + + for concurrently in [false, true] { + let cache = cache_holding(256, rows, concurrently).await; + for i in 0..entries { + cache.insert(EntryID::from(i), array.clone()).await.unwrap(); + } + + for i in 0..entries { + let read = cache + .get(&EntryID::from(i)) + .read() + .await + .expect("every inserted entry is still readable"); + assert_eq!( + read.as_ref(), + array.as_ref(), + "entry {i} came back changed (squeeze_victims_concurrently={concurrently})" + ); + } + + let stats = cache.stats(); + assert_eq!(stats.total_entries, entries); + assert_eq!( + stats.in_flight_memory_bytes, 0, + "every reservation taken while making room has been released" + ); + assert_eq!( + stats.memory_usage_bytes, + indexed_memory_bytes(&cache), + "the tier's tally has drifted from what its entries hold" + ); + assert!(stats.memory_usage_bytes <= stats.max_memory_bytes); + assert!( + stats.peak_in_flight_memory_bytes > 0, + "making room for this many entries has to hold something outside \ + the index at some point" + ); + } +} diff --git a/src/core/src/cache/tests/mod.rs b/src/core/src/cache/tests/mod.rs index 86beb85c..4af8439b 100644 --- a/src/core/src/cache/tests/mod.rs +++ b/src/core/src/cache/tests/mod.rs @@ -1,2 +1,3 @@ +mod in_flight; mod policies; mod squeezed; diff --git a/src/datafusion-server/src/admin_server/handlers.rs b/src/datafusion-server/src/admin_server/handlers.rs index 1be5b0be..89e01913 100644 --- a/src/datafusion-server/src/admin_server/handlers.rs +++ b/src/datafusion-server/src/admin_server/handlers.rs @@ -123,6 +123,8 @@ pub(crate) struct CacheInfo { batch_size: usize, max_memory_bytes: u64, memory_usage_bytes: u64, + in_flight_memory_bytes: u64, + peak_in_flight_memory_bytes: u64, disk_usage_bytes: u64, } @@ -132,11 +134,15 @@ pub(crate) async fn get_cache_info_handler(State(state): State>) - let batch_size = cache.batch_size(); let max_memory_bytes = cache.max_memory_bytes() as u64; let memory_usage_bytes = cache.memory_usage_bytes() as u64; + let in_flight_memory_bytes = cache.in_flight_memory_bytes() as u64; + let peak_in_flight_memory_bytes = cache.peak_in_flight_memory_bytes() as u64; let disk_usage_bytes = cache.disk_usage_bytes() as u64; Json(CacheInfo { batch_size, max_memory_bytes, memory_usage_bytes, + in_flight_memory_bytes, + peak_in_flight_memory_bytes, disk_usage_bytes, }) } diff --git a/src/datafusion/src/cache/mod.rs b/src/datafusion/src/cache/mod.rs index 814d0d32..fc3de32f 100644 --- a/src/datafusion/src/cache/mod.rs +++ b/src/datafusion/src/cache/mod.rs @@ -374,11 +374,30 @@ impl LiquidCacheParquet { self.cache_store.config().max_disk_bytes() } - /// Get the memory usage of the cache in bytes. + /// Get the memory usage of the cache in bytes: the bytes held by the index. pub fn memory_usage_bytes(&self) -> usize { self.cache_store.budget().memory_usage_bytes() } + /// Get the bytes the cache currently holds in memory but has not indexed: + /// entries being decoded off disk, squeeze output awaiting insertion, and + /// entries pending admission. + /// + /// Export this next to [`Self::memory_usage_bytes`]. That figure alone can + /// sit exactly at the limit while the cycle behind it holds a multiple of + /// the limit, which is invisible from the index. + pub fn in_flight_memory_bytes(&self) -> usize { + self.cache_store.budget().in_flight_memory_bytes() + } + + /// Get the high water mark of [`Self::in_flight_memory_bytes`]. + /// + /// Transients are born and freed between gauge scrapes, so the peak is what + /// a scrape can actually catch. + pub fn peak_in_flight_memory_bytes(&self) -> usize { + self.cache_store.budget().peak_in_flight_memory_bytes() + } + /// Get the disk usage of the cache in bytes. pub fn disk_usage_bytes(&self) -> usize { self.cache_store.budget().disk_usage_bytes()