diff --git a/data_plane/src/stores/sketch_db/index/mod.rs b/data_plane/src/stores/sketch_db/index/mod.rs index 5a95d4ba..acf0fb63 100644 --- a/data_plane/src/stores/sketch_db/index/mod.rs +++ b/data_plane/src/stores/sketch_db/index/mod.rs @@ -487,6 +487,20 @@ impl AggPayload { AggPayload::Sketch(_) => None, } } + + /// Approximate in-memory byte footprint of the payload — used by + /// the persistence layer's memory-pressure trigger. Sketch + /// payloads report their byte buffer length (the dominant cost); + /// precompute payloads forward to the accumulator's own + /// `approx_memory_bytes()`. + pub fn approx_bytes(&self) -> usize { + match self { + AggPayload::Sketch(s) => { + s.bytes.len() + std::mem::size_of::() + } + AggPayload::Precompute(p) => p.approx_memory_bytes(), + } + } } /// Per-sid storage value — wraps `SidStoreData` in an `RwLock` so the @@ -888,6 +902,151 @@ impl SketchIndex { } } +// ── Phase 5 M2.3.6b — EpochSource impl ────────────────────────────────────── +// +// Lets the existing persistence flusher (`store/persistence/flusher.rs`) +// drive `SketchIndex` instead of `SketchStorePerKey`. The `agg_id: u64` +// field on `SealedEpochRef` / `EpochSnapshot` carries a `sid` here — +// the trait keeps the historical name so the flusher / manifest / +// part-writer stay untouched. +impl crate::stores::sketch_db::store::persistence::EpochSource for SketchIndex { + fn list_sealed_epochs( + &self, + ) -> Vec { + use crate::stores::sketch_db::store::persistence::SealedEpochRef; + let mut out = Vec::new(); + for entry in self.series.iter() { + let sid = *entry.key(); + let Ok(data) = entry.value().read() else { + continue; + }; + for (epoch_id, epoch) in data.sealed_epochs.iter() { + if let Some((_, max_end)) = epoch.time_bounds() { + let approx_bytes: usize = + epoch.entries.iter().map(|(_, _, p)| p.approx_bytes()).sum(); + out.push(SealedEpochRef { + agg_id: sid, + epoch_id: *epoch_id, + end_ts: max_end, + approx_bytes, + }); + } + } + } + out + } + + fn snapshot_sealed_epoch( + &self, + sid: u64, + epoch_id: u64, + ) -> crate::stores::sketch_db::store::persistence::PersistResult< + Option, + > { + use crate::stores::sketch_db::store::persistence::source::{ + EpochSnapshot, EpochSnapshotEntry, + }; + use crate::stores::sketch_db::store::persistence::PersistError; + + let Some(store_ref) = self.series.get(&sid) else { + return Ok(None); + }; + let data = store_ref + .read() + .map_err(|_| PersistError::Internal(format!("sid {sid}: read lock poisoned")))?; + let Some(epoch) = data.sealed_epochs.get(&epoch_id) else { + return Ok(None); + }; + let Some((min_ts, max_ts)) = epoch.time_bounds() else { + return Ok(None); + }; + + // Resolve sketch_kind once from instance metadata so sketch + // payloads can label their bytes for read-back dispatch. + // Precompute payloads pull their type_name directly from the + // accumulator trait. + let sketch_kind_label: Option = { + self.instances + .read() + .ok() + .and_then(|g| g.get(&sid).cloned()) + .and_then(|m| match &m.agg_kind { + AggKind::Sketch { kind, .. } => Some(format!("{:?}", kind)), + AggKind::Precompute { .. } => None, + }) + }; + + let mut entries = Vec::with_capacity(epoch.entries.len()); + let mut approx_bytes: usize = 0; + for (window, label_id, payload) in &epoch.entries { + let label_map = data.intern.resolve(*label_id).cloned(); + let label_kv = label_map.and_then(|m| { + if m.is_empty() { + None + } else { + Some(crate::stores::types::KeyByLabelValues { + labels: m.values().cloned().collect(), + }) + } + }); + let (type_name, bytes) = match payload { + AggPayload::Sketch(s) => ( + sketch_kind_label + .clone() + .unwrap_or_else(|| "UnknownSketch".to_string()), + s.bytes.clone(), + ), + AggPayload::Precompute(p) => (p.type_name().to_string(), { + use asap_types::traits::SerializableToSink; + p.serialize_to_bytes() + }), + }; + approx_bytes += payload.approx_bytes(); + entries.push(EpochSnapshotEntry { + start_ts: window.0, + end_ts: window.1, + label: label_kv, + sketch_type_name: type_name, + sketch_bytes: bytes, + }); + } + + Ok(Some(EpochSnapshot { + agg_id: sid, + epoch_id, + min_ts, + max_ts, + entries, + approx_bytes, + })) + } + + fn evict_sealed_epoch(&self, sid: u64, epoch_id: u64) { + let Some(store_ref) = self.series.get(&sid) else { + return; + }; + let Ok(mut data) = store_ref.write() else { + return; + }; + data.sealed_epochs.remove(&epoch_id); + } + + fn approx_memory_bytes(&self) -> usize { + let mut total = 0usize; + for entry in self.series.iter() { + let Ok(data) = entry.value().read() else { + continue; + }; + for epoch in data.sealed_epochs.values() { + for (_, _, payload) in &epoch.entries { + total += payload.approx_bytes(); + } + } + } + total + } +} + #[cfg(test)] mod tests { use super::*; @@ -1338,6 +1497,93 @@ mod tests { assert!(result.is_empty(), "sketch sids must not surface"); } + /// Helper: append one sample to materialize the SidStoreData, then + /// set its epoch_capacity so subsequent appends rotate aggressively. + fn with_tight_rotation(idx: &SketchIndex, sid: u64) { + idx.append_sample(sid, BTreeMap::new(), (0, 10), sample(0)); + if let Some(s) = idx.series.get(&sid) { + let mut g = s.write().unwrap(); + g.epoch_capacity = Some(1); + g.max_epochs = 8; + } + } + + #[test] + fn epoch_source_lists_only_sealed_epochs() { + use crate::stores::sketch_db::store::persistence::EpochSource; + let idx = SketchIndex::new(); + idx.register(meta(13)); + with_tight_rotation(&idx, 13); + idx.append_sample(13, BTreeMap::new(), (10, 20), sample(2)); + idx.append_sample(13, BTreeMap::new(), (20, 30), sample(3)); + + let refs = idx.list_sealed_epochs(); + assert!( + !refs.is_empty(), + "rotation should have produced at least one sealed epoch" + ); + assert!( + refs.iter().all(|r| r.agg_id == 13), + "all sealed-epoch refs come from sid=13" + ); + } + + #[test] + fn epoch_source_snapshot_round_trips_sketch_payload() { + use crate::stores::sketch_db::store::persistence::EpochSource; + let idx = SketchIndex::new(); + idx.register(meta(21)); + with_tight_rotation(&idx, 21); + idx.append_sample(21, BTreeMap::new(), (1000, 2000), sample(0xAB)); + idx.append_sample(21, BTreeMap::new(), (2000, 3000), sample(0xCD)); + + let refs = idx.list_sealed_epochs(); + let first = refs.first().expect("a sealed epoch exists"); + let snap = idx + .snapshot_sealed_epoch(first.agg_id, first.epoch_id) + .expect("snapshot ok") + .expect("populated"); + assert_eq!(snap.agg_id, 21); + assert!(!snap.entries.is_empty()); + let entry = &snap.entries[0]; + assert!( + entry.sketch_type_name.starts_with("DDSketch"), + "sketch_type_name should reflect sid metadata's sketch_kind: {}", + entry.sketch_type_name + ); + assert_eq!(entry.sketch_bytes.len(), 1, "single-byte sample bytes carry"); + } + + #[test] + fn epoch_source_evict_drops_the_epoch() { + use crate::stores::sketch_db::store::persistence::EpochSource; + let idx = SketchIndex::new(); + idx.register(meta(31)); + with_tight_rotation(&idx, 31); + idx.append_sample(31, BTreeMap::new(), (10, 20), sample(2)); + + let refs = idx.list_sealed_epochs(); + let one = refs.first().cloned().expect("populated"); + idx.evict_sealed_epoch(one.agg_id, one.epoch_id); + let after = idx.list_sealed_epochs(); + assert!( + !after.iter().any(|r| r.epoch_id == one.epoch_id), + "the evicted epoch must no longer appear" + ); + } + + #[test] + fn epoch_source_approx_memory_bytes_grows_with_sealed_state() { + use crate::stores::sketch_db::store::persistence::EpochSource; + let idx = SketchIndex::new(); + let before = idx.approx_memory_bytes(); + idx.register(meta(41)); + with_tight_rotation(&idx, 41); + idx.append_sample(41, BTreeMap::new(), (10, 20), sample(2)); + let after = idx.approx_memory_bytes(); + assert!(after > before, "sealed state contributes to memory total"); + } + #[test] fn agg_payload_accessors_are_disjoint() { let sketch = AggPayload::Sketch(SketchSampleState {