@@ -26,6 +26,26 @@ use crate::segments::SegmentSource;
2626///
2727/// Caches are optional and operate above a [`SegmentSource`]. They should only store host buffers:
2828/// device buffers and other non-host handles should be passed through uncached.
29+ ///
30+ /// To bring your own cache, implement this trait and give it to the file with
31+ /// `VortexOpenOptions::with_segment_cache`. A [`SegmentId`] is unique only within one file, so a
32+ /// cache that many files share must also key its entries by file. Give each file its own
33+ /// implementation of this trait that holds the [`SegmentSourceId`] of that file and adds it to
34+ /// the key, as [`MokaSegmentCache::for_file`] does with [`FileSegmentCache`].
35+ ///
36+ /// The Python bindings accept only a [`MokaSegmentCache`]. A cache written in Python can use the
37+ /// same design as `PyReadable` in `vortex-python`:
38+ ///
39+ /// - The Python object has the methods `get(key: str, segment_id: int) -> Buffer | None` and
40+ /// `put(key: str, segment_id: int, data: Buffer) -> None`. `vortex.open` and
41+ /// `vortex.open_readable` accept it with the `cache_key` that they accept now, and the stubs
42+ /// declare it as a `typing.Protocol`. Then users can use `cachetools`, `diskcache`, Redis, or a
43+ /// cache that many processes share.
44+ /// - A Rust adapter holds the Python object and the [`SegmentSourceId`] of one file, and
45+ /// implements this trait. `get` calls Python through `spawn_blocking` and `Python::attach`, as
46+ /// `PyReadable::read_at` does, and wraps the buffer that Python returns without a copy.
47+ /// `put` gives Python a read-only object that owns a clone of the [`ByteBuffer`] and exports
48+ /// it through `__getbuffer__`, as `ReadBuffer` does, so Python can keep it without a copy.
2949#[ async_trait]
3050pub trait SegmentCache : Send + Sync {
3151 /// Return a cached segment, or `None` on cache miss.
@@ -48,36 +68,130 @@ impl SegmentCache for NoOpSegmentCache {
4868 }
4969}
5070
51- /// A [`SegmentCache`] based around an in-memory Moka cache.
52- pub struct MokaSegmentCache ( Cache < SegmentId , ByteBuffer , FxBuildHasher > ) ;
71+ /// Which segments a [`MokaSegmentCache`] evicts when it is full.
72+ #[ derive( Clone , Copy , Debug , Default , PartialEq , Eq ) ]
73+ pub enum SegmentEviction {
74+ /// Admit every new segment and evict the least recently used. Suits a cache shared by readers
75+ /// that move from file to file, where a newly opened file's segments must displace older ones.
76+ #[ default]
77+ Lru ,
78+ /// Admit a new segment only if it is likely to be used more often than the one it would
79+ /// evict. Suits repeated scans of one file larger than the cache, which LRU would evict
80+ /// entirely on every pass, but rejects a new file's segments while older ones are popular.
81+ TinyLfu ,
82+ }
83+
84+ /// Identifies the contents of one segment source, usually a file, in a cache that many sources
85+ /// share.
86+ ///
87+ /// Sources with the same ID share cached segments, so an ID must identify the contents of the
88+ /// source, not only its location. A file that has changed must get a new ID.
89+ #[ derive( Clone , Debug , PartialEq , Eq , Hash ) ]
90+ pub struct SegmentSourceId ( Arc < str > ) ;
91+
92+ impl SegmentSourceId {
93+ /// Construct an ID from a string that identifies the contents of a source.
94+ pub fn new ( id : impl Into < Arc < str > > ) -> Self {
95+ Self ( id. into ( ) )
96+ }
97+ }
98+
99+ impl From < & str > for SegmentSourceId {
100+ fn from ( id : & str ) -> Self {
101+ Self :: new ( id)
102+ }
103+ }
104+
105+ impl From < String > for SegmentSourceId {
106+ fn from ( id : String ) -> Self {
107+ Self :: new ( id)
108+ }
109+ }
110+
111+ impl From < Arc < str > > for SegmentSourceId {
112+ fn from ( id : Arc < str > ) -> Self {
113+ Self ( id)
114+ }
115+ }
116+
117+ /// An in-memory Moka cache of segments, capped by total buffer bytes, that any number of files can
118+ /// share.
119+ ///
120+ /// A [`SegmentId`] is unique only within one file, so files use the cache through a
121+ /// [`FileSegmentCache`] from [`Self::for_file`], which adds a [`SegmentSourceId`]. Opening the
122+ /// same file again with the same ID reuses the segments an earlier open read.
123+ #[ derive( Clone ) ]
124+ pub struct MokaSegmentCache ( Cache < ( SegmentSourceId , SegmentId ) , ByteBuffer , FxBuildHasher > ) ;
53125
54126impl MokaSegmentCache {
55127 /// Construct a Moka-backed cache capped by total buffer bytes.
56- pub fn new ( max_capacity_bytes : u64 ) -> Self {
128+ pub fn new ( max_capacity_bytes : u64 , eviction : SegmentEviction ) -> Self {
129+ let eviction_policy = match eviction {
130+ SegmentEviction :: Lru => EvictionPolicy :: lru ( ) ,
131+ SegmentEviction :: TinyLfu => EvictionPolicy :: tiny_lfu ( ) ,
132+ } ;
57133 Self (
58134 CacheBuilder :: new ( max_capacity_bytes)
59135 . name ( "vortex-segment-cache" )
60136 // Weight each segment by the number of bytes in the buffer.
61137 . weigher ( |_, buffer : & ByteBuffer | {
62138 u32:: try_from ( buffer. len ( ) . min ( u32:: MAX as usize ) ) . vortex_expect ( "must fit" )
63139 } )
64- // We configure LFU (vs LRU) since the cache is mostly used when re-reading the
65- // same file - it is _not_ used when reading the same segments during a single
66- // scan.
67- . eviction_policy ( EvictionPolicy :: tiny_lfu ( ) )
140+ . eviction_policy ( eviction_policy)
68141 . build_with_hasher ( FxBuildHasher ) ,
69142 )
70143 }
144+
145+ /// The view of this cache for the file identified by `source_id`.
146+ pub fn for_file ( & self , source_id : impl Into < SegmentSourceId > ) -> FileSegmentCache {
147+ FileSegmentCache {
148+ cache : self . clone ( ) ,
149+ source_id : source_id. into ( ) ,
150+ }
151+ }
152+
153+ /// The total bytes of the cached segments.
154+ ///
155+ /// Recent inserts and evictions may not be counted yet; call [`Self::run_pending_tasks`] first
156+ /// for an exact value.
157+ pub fn weighted_size ( & self ) -> u64 {
158+ self . 0 . weighted_size ( )
159+ }
160+
161+ /// The number of cached segments, with the same caveat as [`Self::weighted_size`].
162+ pub fn entry_count ( & self ) -> u64 {
163+ self . 0 . entry_count ( )
164+ }
165+
166+ /// Apply pending inserts and evictions, so that the counts are exact.
167+ pub async fn run_pending_tasks ( & self ) {
168+ self . 0 . run_pending_tasks ( ) . await ;
169+ }
170+
171+ /// Remove every cached segment.
172+ pub fn invalidate_all ( & self ) {
173+ self . 0 . invalidate_all ( ) ;
174+ }
175+ }
176+
177+ /// One file's view of a [`MokaSegmentCache`]; see [`MokaSegmentCache::for_file`].
178+ #[ derive( Clone ) ]
179+ pub struct FileSegmentCache {
180+ cache : MokaSegmentCache ,
181+ source_id : SegmentSourceId ,
71182}
72183
73184#[ async_trait]
74- impl SegmentCache for MokaSegmentCache {
185+ impl SegmentCache for FileSegmentCache {
75186 async fn get ( & self , id : SegmentId ) -> VortexResult < Option < ByteBuffer > > {
76- Ok ( self . 0 . get ( & id ) . await )
187+ Ok ( self . cache . 0 . get ( & ( self . source_id . clone ( ) , id ) ) . await )
77188 }
78189
79190 async fn put ( & self , id : SegmentId , buffer : ByteBuffer ) -> VortexResult < ( ) > {
80- self . 0 . insert ( id, buffer) . await ;
191+ self . cache
192+ . 0
193+ . insert ( ( self . source_id . clone ( ) , id) , buffer)
194+ . await ;
81195 Ok ( ( ) )
82196 }
83197}
@@ -167,3 +281,59 @@ impl SegmentSource for SegmentCacheSourceAdapter {
167281 . boxed ( )
168282 }
169283}
284+
285+ #[ cfg( test) ]
286+ mod tests {
287+ use rstest:: rstest;
288+ use vortex_buffer:: ByteBuffer ;
289+
290+ use super :: * ;
291+
292+ #[ rstest]
293+ #[ case:: lru( SegmentEviction :: Lru ) ]
294+ #[ case:: tiny_lfu( SegmentEviction :: TinyLfu ) ]
295+ #[ tokio:: test]
296+ async fn shared_cache_separates_files ( #[ case] eviction : SegmentEviction ) -> VortexResult < ( ) > {
297+ let shared = MokaSegmentCache :: new ( 1 << 20 , eviction) ;
298+ let a = shared. for_file ( "a" ) ;
299+ let b = shared. for_file ( "b" ) ;
300+ let id = SegmentId :: from ( 0 ) ;
301+
302+ a. put ( id, ByteBuffer :: copy_from ( b"from a" ) ) . await ?;
303+ assert_eq ! ( a. get( id) . await ?. as_deref( ) , Some ( b"from a" . as_slice( ) ) ) ;
304+ assert ! ( b. get( id) . await ?. is_none( ) ) ;
305+
306+ // A later view with the same source ID sees what the earlier one stored.
307+ let a_again = shared. for_file ( "a" ) ;
308+ assert_eq ! (
309+ a_again. get( id) . await ?. as_deref( ) ,
310+ Some ( b"from a" . as_slice( ) )
311+ ) ;
312+ Ok ( ( ) )
313+ }
314+
315+ #[ rstest]
316+ #[ case:: lru( SegmentEviction :: Lru ) ]
317+ #[ case:: tiny_lfu( SegmentEviction :: TinyLfu ) ]
318+ #[ tokio:: test]
319+ async fn shared_cache_is_capped_by_bytes (
320+ #[ case] eviction : SegmentEviction ,
321+ ) -> VortexResult < ( ) > {
322+ let shared = MokaSegmentCache :: new ( 1000 , eviction) ;
323+ let file = shared. for_file ( "file" ) ;
324+ for i in 0 ..10u32 {
325+ file. put ( SegmentId :: from ( i) , ByteBuffer :: copy_from ( vec ! [ 0u8 ; 400 ] ) )
326+ . await ?;
327+ }
328+
329+ shared. run_pending_tasks ( ) . await ;
330+ assert ! ( shared. weighted_size( ) <= 1000 ) ;
331+ assert ! ( shared. entry_count( ) <= 2 ) ;
332+
333+ // LRU admits every segment, so the most recently stored one survives.
334+ if eviction == SegmentEviction :: Lru {
335+ assert ! ( file. get( SegmentId :: from( 9 ) ) . await ?. is_some( ) ) ;
336+ }
337+ Ok ( ( ) )
338+ }
339+ }
0 commit comments