diff --git a/vortex-layout/src/segments/cache.rs b/vortex-layout/src/segments/cache.rs index 1f7f5e91f5c..edec733b45c 100644 --- a/vortex-layout/src/segments/cache.rs +++ b/vortex-layout/src/segments/cache.rs @@ -148,14 +148,14 @@ impl SegmentCacheSourceAdapter { impl SegmentSource for SegmentCacheSourceAdapter { fn request(&self, id: SegmentId) -> SegmentFuture { let cache = Arc::clone(&self.cache); - let delegate = self.source.request(id); + let source = Arc::clone(&self.source); async move { if let Ok(Some(segment)) = cache.get(id).await { tracing::debug!("Resolved segment {} from cache", id); return Ok(BufferHandle::new_host(segment)); } - let result = delegate.await?; + let result = source.request(id).await?; // Cache only CPU buffers; device buffers are not cached. if let Some(buffer) = result.as_host_opt() && let Err(e) = cache.put(id, buffer.clone()).await @@ -167,3 +167,98 @@ impl SegmentSource for SegmentCacheSourceAdapter { .boxed() } } + +#[cfg(test)] +mod tests { + use std::sync::atomic::AtomicUsize; + use std::sync::atomic::Ordering; + + use vortex_error::vortex_err; + + use super::*; + + #[derive(Default)] + struct CountingSource { + requests: AtomicUsize, + } + + impl SegmentSource for CountingSource { + fn request(&self, _id: SegmentId) -> SegmentFuture { + self.requests.fetch_add(1, Ordering::SeqCst); + futures::future::ready(Ok(BufferHandle::new_host(ByteBuffer::from(vec![ + 1u8, 2, 3, + ])))) + .boxed() + } + } + + #[tokio::test] + async fn cache_hit_does_not_request_the_source() -> VortexResult<()> { + let cache: Arc = Arc::new(MokaSegmentCache::new(1024)); + let id = SegmentId::from(0); + cache.put(id, ByteBuffer::from(vec![9u8, 8, 7])).await?; + let source = Arc::new(CountingSource::default()); + let adapter = SegmentCacheSourceAdapter::new(cache, Arc::::clone(&source)); + + let request = adapter.request(id); + assert_eq!(source.requests.load(Ordering::SeqCst), 0); + let result = request.await?.to_host().await; + assert_eq!(result.as_slice(), &[9, 8, 7]); + assert_eq!(source.requests.load(Ordering::SeqCst), 0); + Ok(()) + } + + #[tokio::test] + async fn miss_reads_and_populates_the_cache() -> VortexResult<()> { + let cache: Arc = Arc::new(MokaSegmentCache::new(1024)); + let source = Arc::new(CountingSource::default()); + let adapter = SegmentCacheSourceAdapter::new(cache, Arc::::clone(&source)); + let id = SegmentId::from(0); + + for _ in 0..2 { + assert_eq!( + adapter.request(id).await?.to_host().await.as_slice(), + &[1, 2, 3] + ); + } + assert_eq!(source.requests.load(Ordering::SeqCst), 1); + Ok(()) + } + + #[test] + fn dropping_an_unpolled_request_does_not_request_the_source() { + let source = Arc::new(CountingSource::default()); + let adapter = SegmentCacheSourceAdapter::new( + Arc::new(NoOpSegmentCache), + Arc::::clone(&source), + ); + drop(adapter.request(SegmentId::from(0))); + assert_eq!(source.requests.load(Ordering::SeqCst), 0); + } + + struct UnavailableCache; + + #[async_trait] + impl SegmentCache for UnavailableCache { + async fn get(&self, _id: SegmentId) -> VortexResult> { + Err(vortex_err!("cache unavailable")) + } + + async fn put(&self, _id: SegmentId, _buffer: ByteBuffer) -> VortexResult<()> { + Err(vortex_err!("cache unavailable")) + } + } + + #[tokio::test] + async fn unavailable_cache_preserves_source_results() -> VortexResult<()> { + let source = Arc::new(CountingSource::default()); + let adapter = SegmentCacheSourceAdapter::new( + Arc::new(UnavailableCache), + Arc::::clone(&source), + ); + let result = adapter.request(SegmentId::from(0)).await?.to_host().await; + assert_eq!(result.as_slice(), &[1, 2, 3]); + assert_eq!(source.requests.load(Ordering::SeqCst), 1); + Ok(()) + } +}