vortex_layout/segments/
cache.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::sync::Arc;
5use std::task::Poll;
6
7use async_trait::async_trait;
8use futures::FutureExt;
9use futures::future::poll_fn;
10use moka::future::Cache;
11use moka::future::CacheBuilder;
12use moka::policy::EvictionPolicy;
13use rustc_hash::FxBuildHasher;
14use vortex_array::buffer::BufferHandle;
15use vortex_buffer::ByteBuffer;
16use vortex_error::VortexExpect;
17use vortex_error::VortexResult;
18use vortex_metrics::Counter;
19use vortex_metrics::Label;
20use vortex_metrics::MetricBuilder;
21use vortex_metrics::MetricsRegistry;
22
23use crate::segments::SegmentFuture;
24use crate::segments::SegmentId;
25use crate::segments::SegmentSource;
26
27/// Cache for individual segment byte buffers.
28///
29/// Caches are optional and operate above a [`SegmentSource`]. They should only store host buffers:
30/// device buffers and other non-host handles should be passed through uncached.
31///
32/// To bring your own cache, implement this trait and give it to the file with
33/// `VortexOpenOptions::with_segment_cache`. A [`SegmentId`] is unique only within one file, so a
34/// cache that many files share must also key its entries by file. Give each file its own
35/// implementation of this trait that holds the [`SegmentSourceId`] of that file and adds it to
36/// the key, as [`MokaSegmentCache::for_file`] does with [`FileSegmentCache`].
37///
38/// The Python bindings accept only a [`MokaSegmentCache`]. A cache written in Python can use the
39/// same design as `PyReadable` in `vortex-python`:
40///
41/// - The Python object has the methods `get(key: str, segment_id: int) -> Buffer | None` and
42///   `put(key: str, segment_id: int, data: Buffer) -> None`. `vortex.open` and
43///   `vortex.open_readable` accept it with the `cache_key` that they accept now, and the stubs
44///   declare it as a `typing.Protocol`. Then users can use `cachetools`, `diskcache`, Redis, or a
45///   cache that many processes share.
46/// - A Rust adapter holds the Python object and the [`SegmentSourceId`] of one file, and
47///   implements this trait. `get` calls Python through `spawn_blocking` and `Python::attach`, as
48///   `PyReadable::read_at` does, and wraps the buffer that Python returns without a copy.
49///   `put` gives Python a read-only object that owns a clone of the [`ByteBuffer`] and exports
50///   it through `__getbuffer__`, as `ReadBuffer` does, so Python can keep it without a copy.
51#[async_trait]
52pub trait SegmentCache: Send + Sync {
53    /// Return a cached segment, or `None` on cache miss.
54    ///
55    /// Lookups may be polled while a segment request is constructed, before its returned
56    /// future is polled. Ready misses preserve the source's eager read registration.
57    /// Pending lookups delay registration until they resolve and the caller polls the request,
58    /// which can reduce read coalescing. This prevents cache hits from registering source reads.
59    async fn get(&self, id: SegmentId) -> VortexResult<Option<ByteBuffer>>;
60    /// Store a segment in the cache.
61    async fn put(&self, id: SegmentId, buffer: ByteBuffer) -> VortexResult<()>;
62}
63
64/// Segment cache implementation that never stores anything.
65pub struct NoOpSegmentCache;
66
67#[async_trait]
68impl SegmentCache for NoOpSegmentCache {
69    async fn get(&self, _id: SegmentId) -> VortexResult<Option<ByteBuffer>> {
70        Ok(None)
71    }
72
73    async fn put(&self, _id: SegmentId, _buffer: ByteBuffer) -> VortexResult<()> {
74        Ok(())
75    }
76}
77
78/// Which segments a [`MokaSegmentCache`] evicts when it is full.
79#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
80pub enum SegmentEviction {
81    /// Admit every new segment and evict the least recently used. Suits a cache shared by readers
82    /// that move from file to file, where a newly opened file's segments must displace older ones.
83    #[default]
84    Lru,
85    /// Admit a new segment only if it is likely to be used more often than the one it would
86    /// evict. Suits repeated scans of one file larger than the cache, which LRU would evict
87    /// entirely on every pass, but rejects a new file's segments while older ones are popular.
88    TinyLfu,
89}
90
91/// Identifies the contents of one segment source, usually a file, in a cache that many sources
92/// share.
93///
94/// Sources with the same ID share cached segments, so an ID must identify the contents of the
95/// source, not only its location. A file that has changed must get a new ID.
96#[derive(Clone, Debug, PartialEq, Eq, Hash)]
97pub struct SegmentSourceId(Arc<str>);
98
99impl SegmentSourceId {
100    /// Construct an ID from a string that identifies the contents of a source.
101    pub fn new(id: impl Into<Arc<str>>) -> Self {
102        Self(id.into())
103    }
104}
105
106impl From<&str> for SegmentSourceId {
107    fn from(id: &str) -> Self {
108        Self::new(id)
109    }
110}
111
112impl From<String> for SegmentSourceId {
113    fn from(id: String) -> Self {
114        Self::new(id)
115    }
116}
117
118impl From<Arc<str>> for SegmentSourceId {
119    fn from(id: Arc<str>) -> Self {
120        Self(id)
121    }
122}
123
124/// An in-memory Moka cache of segments, capped by total buffer bytes, that any number of files can
125/// share.
126///
127/// A [`SegmentId`] is unique only within one file, so files use the cache through a
128/// [`FileSegmentCache`] from [`Self::for_file`], which adds a [`SegmentSourceId`]. Opening the
129/// same file again with the same ID reuses the segments an earlier open read.
130#[derive(Clone)]
131pub struct MokaSegmentCache(Cache<(SegmentSourceId, SegmentId), ByteBuffer, FxBuildHasher>);
132
133impl MokaSegmentCache {
134    /// Construct a Moka-backed cache capped by total buffer bytes.
135    pub fn new(max_capacity_bytes: u64, eviction: SegmentEviction) -> Self {
136        let eviction_policy = match eviction {
137            SegmentEviction::Lru => EvictionPolicy::lru(),
138            SegmentEviction::TinyLfu => EvictionPolicy::tiny_lfu(),
139        };
140        Self(
141            CacheBuilder::new(max_capacity_bytes)
142                .name("vortex-segment-cache")
143                // Weight each segment by the number of bytes in the buffer.
144                .weigher(|_, buffer: &ByteBuffer| {
145                    u32::try_from(buffer.len().min(u32::MAX as usize)).vortex_expect("must fit")
146                })
147                .eviction_policy(eviction_policy)
148                .build_with_hasher(FxBuildHasher),
149        )
150    }
151
152    /// The view of this cache for the file identified by `source_id`.
153    pub fn for_file(&self, source_id: impl Into<SegmentSourceId>) -> FileSegmentCache {
154        FileSegmentCache {
155            cache: self.clone(),
156            source_id: source_id.into(),
157        }
158    }
159
160    /// The total bytes of the cached segments.
161    ///
162    /// Recent inserts and evictions may not be counted yet; call [`Self::run_pending_tasks`] first
163    /// for an exact value.
164    pub fn weighted_size(&self) -> u64 {
165        self.0.weighted_size()
166    }
167
168    /// The number of cached segments, with the same caveat as [`Self::weighted_size`].
169    pub fn entry_count(&self) -> u64 {
170        self.0.entry_count()
171    }
172
173    /// Apply pending inserts and evictions, so that the counts are exact.
174    pub async fn run_pending_tasks(&self) {
175        self.0.run_pending_tasks().await;
176    }
177
178    /// Remove every cached segment.
179    pub fn invalidate_all(&self) {
180        self.0.invalidate_all();
181    }
182}
183
184/// One file's view of a [`MokaSegmentCache`]; see [`MokaSegmentCache::for_file`].
185#[derive(Clone)]
186pub struct FileSegmentCache {
187    cache: MokaSegmentCache,
188    source_id: SegmentSourceId,
189}
190
191#[async_trait]
192impl SegmentCache for FileSegmentCache {
193    async fn get(&self, id: SegmentId) -> VortexResult<Option<ByteBuffer>> {
194        Ok(self.cache.0.get(&(self.source_id.clone(), id)).await)
195    }
196
197    async fn put(&self, id: SegmentId, buffer: ByteBuffer) -> VortexResult<()> {
198        self.cache
199            .0
200            .insert((self.source_id.clone(), id), buffer)
201            .await;
202        Ok(())
203    }
204}
205
206/// Wrapper for [`SegmentCache`] that tracks its hit rate.
207pub struct InstrumentedSegmentCache<C> {
208    segment_cache: C,
209
210    hits: Counter,
211    misses: Counter,
212    stores: Counter,
213}
214
215impl<C: SegmentCache> InstrumentedSegmentCache<C> {
216    /// Wrap a segment cache and record hit/miss/store metrics with the supplied labels.
217    pub fn new(
218        segment_cache: C,
219        metrics_registry: &dyn MetricsRegistry,
220        labels: Vec<Label>,
221    ) -> Self {
222        Self {
223            segment_cache,
224            hits: MetricBuilder::new(metrics_registry)
225                .add_labels(labels.clone())
226                .counter("vortex.file.segments.cache.hits"),
227            misses: MetricBuilder::new(metrics_registry)
228                .add_labels(labels.clone())
229                .counter("vortex.file.segments.cache.misses"),
230            stores: MetricBuilder::new(metrics_registry)
231                .add_labels(labels)
232                .counter("vortex.file.segments.cache.stores"),
233        }
234    }
235}
236
237#[async_trait]
238impl<C: SegmentCache> SegmentCache for InstrumentedSegmentCache<C> {
239    async fn get(&self, id: SegmentId) -> VortexResult<Option<ByteBuffer>> {
240        let result = self.segment_cache.get(id).await?;
241        if result.is_some() {
242            self.hits.add(1);
243        } else {
244            self.misses.add(1);
245        }
246        Ok(result)
247    }
248
249    async fn put(&self, id: SegmentId, buffer: ByteBuffer) -> VortexResult<()> {
250        self.segment_cache.put(id, buffer).await?;
251        self.stores.add(1);
252        Ok(())
253    }
254}
255
256/// [`SegmentSource`] wrapper that consults a [`SegmentCache`] before the underlying source.
257pub struct SegmentCacheSourceAdapter {
258    cache: Arc<dyn SegmentCache>,
259    source: Arc<dyn SegmentSource>,
260}
261
262impl SegmentCacheSourceAdapter {
263    /// Construct a cache-fronted source.
264    pub fn new(cache: Arc<dyn SegmentCache>, source: Arc<dyn SegmentSource>) -> Self {
265        Self { cache, source }
266    }
267}
268
269impl SegmentSource for SegmentCacheSourceAdapter {
270    fn request(&self, id: SegmentId) -> SegmentFuture {
271        let cache = Arc::clone(&self.cache);
272        let source = Arc::clone(&self.source);
273        let mut request = async move {
274            if let Ok(Some(segment)) = cache.get(id).await {
275                tracing::debug!("Resolved segment {} from cache", id);
276                return Ok(BufferHandle::new_host(segment));
277            }
278            let delegate = source.request(id);
279            // The eager poll may register a miss for coalescing, but only the caller may
280            // demand the read. Pending lookups reaching this point later must wake that caller.
281            let mut defer = true;
282            poll_fn(move |cx| {
283                if std::mem::take(&mut defer) {
284                    cx.waker().wake_by_ref();
285                    Poll::Pending
286                } else {
287                    Poll::Ready(())
288                }
289            })
290            .await;
291            let result = delegate.await?;
292            // Cache only CPU buffers; device buffers are not cached.
293            if let Some(buffer) = result.as_host_opt()
294                && let Err(e) = cache.put(id, buffer.clone()).await
295            {
296                tracing::warn!("Failed to store segment {} in cache: {}", id, e);
297            }
298            Ok(result)
299        }
300        .boxed();
301
302        // Resolve ready lookups while requests are constructed, retaining eager registration
303        // of cold misses without allocating a separate future to own the cache lookup.
304        match request.as_mut().now_or_never() {
305            Some(result) => async move { result }.boxed(),
306            None => request,
307        }
308    }
309}
310
311#[cfg(test)]
312mod tests {
313    use std::sync::atomic::AtomicUsize;
314    use std::sync::atomic::Ordering;
315
316    use futures::channel::oneshot;
317    use parking_lot::Mutex;
318    use rstest::rstest;
319    use vortex_buffer::ByteBuffer;
320
321    use super::*;
322
323    #[rstest]
324    #[case::lru(SegmentEviction::Lru)]
325    #[case::tiny_lfu(SegmentEviction::TinyLfu)]
326    #[tokio::test]
327    async fn shared_cache_separates_files(#[case] eviction: SegmentEviction) -> VortexResult<()> {
328        let shared = MokaSegmentCache::new(1 << 20, eviction);
329        let a = shared.for_file("a");
330        let b = shared.for_file("b");
331        let id = SegmentId::from(0);
332
333        a.put(id, ByteBuffer::copy_from(b"from a")).await?;
334        assert_eq!(a.get(id).await?.as_deref(), Some(b"from a".as_slice()));
335        assert!(b.get(id).await?.is_none());
336
337        // A later view with the same source ID sees what the earlier one stored.
338        let a_again = shared.for_file("a");
339        assert_eq!(
340            a_again.get(id).await?.as_deref(),
341            Some(b"from a".as_slice())
342        );
343        Ok(())
344    }
345
346    #[rstest]
347    #[case::lru(SegmentEviction::Lru)]
348    #[case::tiny_lfu(SegmentEviction::TinyLfu)]
349    #[tokio::test]
350    async fn shared_cache_is_capped_by_bytes(
351        #[case] eviction: SegmentEviction,
352    ) -> VortexResult<()> {
353        let shared = MokaSegmentCache::new(1000, eviction);
354        let file = shared.for_file("file");
355        for i in 0..10u32 {
356            file.put(SegmentId::from(i), ByteBuffer::copy_from(vec![0u8; 400]))
357                .await?;
358        }
359
360        shared.run_pending_tasks().await;
361        assert!(shared.weighted_size() <= 1000);
362        assert!(shared.entry_count() <= 2);
363
364        // LRU admits every segment, so the most recently stored one survives.
365        if eviction == SegmentEviction::Lru {
366            assert!(file.get(SegmentId::from(9)).await?.is_some());
367        }
368        Ok(())
369    }
370
371    #[derive(Default)]
372    struct CountingSource(AtomicUsize, Arc<AtomicUsize>);
373
374    impl SegmentSource for CountingSource {
375        fn request(&self, _id: SegmentId) -> SegmentFuture {
376            self.0.fetch_add(1, Ordering::SeqCst);
377            let demanded = Arc::clone(&self.1);
378            async move {
379                demanded.fetch_add(1, Ordering::SeqCst);
380                Ok(BufferHandle::new_host(ByteBuffer::copy_from(b"source")))
381            }
382            .boxed()
383        }
384    }
385
386    #[tokio::test]
387    async fn cache_lookup_preserves_registration_and_demand_order() -> VortexResult<()> {
388        let cache = MokaSegmentCache::new(1024, SegmentEviction::Lru).for_file("file");
389        let hit = SegmentId::from(0);
390        cache.put(hit, ByteBuffer::copy_from(b"cached")).await?;
391        let source = Arc::new(CountingSource::default());
392        let cached =
393            SegmentCacheSourceAdapter::new(Arc::new(cache), Arc::<CountingSource>::clone(&source));
394
395        assert_eq!(
396            cached.request(hit).await?.unwrap_host(),
397            ByteBuffer::copy_from(b"cached")
398        );
399        assert_eq!(source.0.load(Ordering::SeqCst), 0);
400        let miss = cached.request(SegmentId::from(1));
401        assert_eq!(source.0.load(Ordering::SeqCst), 1);
402        assert_eq!(source.1.load(Ordering::SeqCst), 0);
403        assert_eq!(miss.await?.unwrap_host(), ByteBuffer::copy_from(b"source"));
404        assert_eq!(source.1.load(Ordering::SeqCst), 1);
405        Ok(())
406    }
407
408    struct DelayedCache(Mutex<Option<oneshot::Receiver<Option<ByteBuffer>>>>);
409
410    #[async_trait]
411    impl SegmentCache for DelayedCache {
412        async fn get(&self, _id: SegmentId) -> VortexResult<Option<ByteBuffer>> {
413            let lookup = self
414                .0
415                .lock()
416                .take()
417                .vortex_expect("lookup is not restarted");
418            Ok(lookup.await.expect("test supplies lookup result"))
419        }
420
421        async fn put(&self, _id: SegmentId, _buffer: ByteBuffer) -> VortexResult<()> {
422            Ok(())
423        }
424    }
425
426    #[rstest]
427    #[case::hit(true)]
428    #[case::miss(false)]
429    #[tokio::test]
430    async fn pending_lookup_finishes_before_registering_reads(
431        #[case] hit: bool,
432    ) -> VortexResult<()> {
433        let (sender, receiver) = oneshot::channel();
434        let cache = Arc::new(DelayedCache(Mutex::new(Some(receiver))));
435        let source = Arc::new(CountingSource::default());
436        let cached = SegmentCacheSourceAdapter::new(cache, Arc::<CountingSource>::clone(&source));
437        let mut request = cached.request(SegmentId::from(0));
438        assert_eq!(source.0.load(Ordering::SeqCst), 0);
439        assert!(futures::poll!(request.as_mut()).is_pending());
440        assert_eq!(source.0.load(Ordering::SeqCst), 0);
441        sender
442            .send(hit.then(|| ByteBuffer::copy_from(b"cached")))
443            .unwrap();
444        let expected: &[u8] = if hit { b"cached" } else { b"source" };
445        assert_eq!(request.await?.unwrap_host().as_ref(), expected);
446        assert_eq!(source.0.load(Ordering::SeqCst), usize::from(!hit));
447        Ok(())
448    }
449}
vortex_layout/segments/cache.rs

vortex_layout/segments/
cache.rs
