diff --git a/src/cache/cache.rs b/src/cache/cache.rs index a57bde6e9..c11dddfc1 100644 --- a/src/cache/cache.rs +++ b/src/cache/cache.rs @@ -82,6 +82,28 @@ pub trait Storage: Send + Sync { /// return a `Cache::Hit`. async fn get(&self, key: &str) -> Result; + /// Get a parsed cache entry and, when supported, the same entry's raw bytes. + /// + /// Multi-level caches use the raw bytes to backfill faster levels. The + /// default preserves compatibility with existing backends by retaining + /// the historical second `get_raw()` call; raw-capable backends can + /// override this to avoid reading a hit twice. + async fn get_with_raw(&self, key: &str) -> Result<(Cache, Option)> { + let cache = self.get(key).await?; + let raw = if matches!(&cache, Cache::Hit(_)) { + match self.get_raw(key).await { + Ok(raw) => raw, + Err(error) => { + debug!("Failed to get raw bytes for cache backfill: {}", error); + None + } + } + } else { + None + }; + Ok((cache, raw)) + } + /// Put `entry` in the cache under `key`. /// /// Returns a `Future` that will provide the result or error when the put is @@ -235,15 +257,20 @@ impl RemoteStorage { #[async_trait] impl Storage for RemoteStorage { async fn get(&self, key: &str) -> Result { + Ok(self.get_with_raw(key).await?.0) + } + + async fn get_with_raw(&self, key: &str) -> Result<(Cache, Option)> { match self.operator.read(&normalize_key(key)).await { Ok(res) => { - let hit = CacheRead::from(io::Cursor::new(res.to_bytes()))?; - Ok(Cache::Hit(hit)) + let data = res.to_bytes(); + let hit = CacheRead::from(io::Cursor::new(data.clone()))?; + Ok((Cache::Hit(hit), Some(data))) } - Err(e) if e.kind() == opendal::ErrorKind::NotFound => Ok(Cache::Miss), + Err(e) if e.kind() == opendal::ErrorKind::NotFound => Ok((Cache::Miss, None)), Err(e) => { warn!("Got unexpected error: {:?}", e); - Ok(Cache::Miss) + Ok((Cache::Miss, None)) } } } diff --git a/src/cache/disk.rs b/src/cache/disk.rs index d1269e694..13910dc53 100644 --- a/src/cache/disk.rs +++ b/src/cache/disk.rs @@ -103,6 +103,16 @@ impl Storage for DiskCache { .await? } + async fn get_with_raw(&self, key: &str) -> Result<(Cache, Option)> { + match self.get_raw(key).await? { + Some(data) => { + let hit = CacheRead::from(std::io::Cursor::new(data.clone()))?; + Ok((Cache::Hit(hit), Some(data))) + } + None => Ok((Cache::Miss, None)), + } + } + async fn get_raw(&self, key: &str) -> Result> { trace!("DiskCache::get_raw({})", key); let path = make_key_path(key); diff --git a/src/cache/multilevel.rs b/src/cache/multilevel.rs index d7ad127e2..0b2a300f9 100644 --- a/src/cache/multilevel.rs +++ b/src/cache/multilevel.rs @@ -639,7 +639,20 @@ impl Storage for MultiLevelStorage { async fn get(&self, key: &str) -> Result { for (idx, level) in self.levels.iter().enumerate() { let start = Instant::now(); - match level.get(key).await { + let mut raw_bytes_for_backfill = None; + let cache_result = if idx > 0 { + match level.get_with_raw(key).await { + Ok((cache, raw_bytes)) => { + raw_bytes_for_backfill = raw_bytes; + Ok(cache) + } + Err(error) => Err(error), + } + } else { + level.get(key).await + }; + + match cache_result { Ok(Cache::Hit(entry)) => { let duration = start.elapsed(); debug!("Cache hit at level {} in {:?}", idx, duration); @@ -661,9 +674,10 @@ impl Storage for MultiLevelStorage { let key_str = key.to_string(); let hit_level = idx; - // Try to get raw bytes for backfilling - match level.get_raw(key).await { - Ok(Some(raw_bytes)) => { + // Raw bytes obtained above are reused for backfilling; + // no second read is needed for a raw-capable level. + match raw_bytes_for_backfill { + Some(raw_bytes) => { // Update backfill stats inc_stat!( self.atomic_stats.get(hit_level), @@ -704,18 +718,12 @@ impl Storage for MultiLevelStorage { }); } } - Ok(None) => { + None => { debug!( "Cache backend at level {} does not support get_raw(), skipping backfill", hit_level ); } - Err(e) => { - debug!( - "Failed to get raw bytes from level {} for backfill: {}", - hit_level, e - ); - } } } diff --git a/src/cache/multilevel_test.rs b/src/cache/multilevel_test.rs index b75f6d3fa..582e310ce 100644 --- a/src/cache/multilevel_test.rs +++ b/src/cache/multilevel_test.rs @@ -197,6 +197,7 @@ fn test_multi_level_storage_backfill_on_hit() { struct InMemoryStorage { data: Arc>>>, access_log: Arc>>, + raw_access_log: Arc>>, } impl InMemoryStorage { @@ -204,12 +205,17 @@ impl InMemoryStorage { Self { data: Arc::new(Mutex::new(HashMap::new())), access_log: Arc::new(Mutex::new(Vec::new())), + raw_access_log: Arc::new(Mutex::new(Vec::new())), } } fn get_access_log(&self) -> Arc>> { Arc::clone(&self.access_log) } + + fn get_raw_access_log(&self) -> Arc>> { + Arc::clone(&self.raw_access_log) + } } #[async_trait] @@ -230,6 +236,17 @@ impl Storage for InMemoryStorage { } } + async fn get_with_raw(&self, key: &str) -> Result<(Cache, Option)> { + let Some(bytes) = self.get_raw(key).await? else { + return Ok((Cache::Miss, None)); + }; + let raw_bytes = bytes.clone(); + match CacheRead::from(Cursor::new(bytes)) { + Ok(entry) => Ok((Cache::Hit(entry), Some(raw_bytes))), + Err(error) => Err(error), + } + } + async fn put(&self, key: &str, entry: CacheWrite) -> Result { self.access_log.lock().await.push(format!("put:{}", key)); @@ -258,6 +275,10 @@ impl Storage for InMemoryStorage { /// This simulates the behavior of real remote backends (S3, Redis, etc.) that /// can efficiently return raw serialized cache entries for backfilling. async fn get_raw(&self, key: &str) -> Result> { + self.raw_access_log + .lock() + .await + .push(format!("get_raw:{}", key)); Ok(self.data.lock().await.get(key).cloned().map(Bytes::from)) } @@ -271,6 +292,44 @@ impl Storage for InMemoryStorage { } } +#[test] +fn test_multilevel_raw_hit_reads_backend_once() { + let runtime = RuntimeBuilder::new_multi_thread() + .enable_all() + .worker_threads(1) + .build() + .unwrap(); + + let l0 = Arc::new(InMemoryStorage::new()); + let l1 = Arc::new(InMemoryStorage::new()); + let storage = MultiLevelStorage::new(vec![ + l0.clone() as Arc, + l1.clone() as Arc, + ]); + + runtime.block_on(async { + let entry = CacheWrite::default(); + l1.put("single_read_key", entry).await.unwrap(); + + // Ignore setup writes and assert the lookup path itself. A raw-capable + // level should be read once, then the same bytes should feed both + // parsing and the asynchronous backfill. + l0.get_access_log().lock().await.clear(); + l1.get_access_log().lock().await.clear(); + l1.get_raw_access_log().lock().await.clear(); + + assert!(matches!( + storage.get("single_read_key").await.unwrap(), + Cache::Hit(_) + )); + + assert_eq!( + l1.get_raw_access_log().lock().await.as_slice(), + &["get_raw:single_read_key"] + ); + }); +} + #[test] fn test_disk_plus_remote_to_remote_backfill() { let runtime = RuntimeBuilder::new_multi_thread() @@ -987,6 +1046,8 @@ fn test_sequential_read_order() { let l0_log = l0.get_access_log(); let l1_log = l1.get_access_log(); let l2_log = l2.get_access_log(); + let l1_raw_log = l1.get_raw_access_log(); + let l2_raw_log = l2.get_raw_access_log(); // Put data only in L2 (slowest level) let key = "test_key_12345678901234567890"; @@ -1012,15 +1073,22 @@ fn test_sequential_read_order() { let l1_accesses = l1_log.lock().await; let l2_accesses = l2_log.lock().await; - // Each level should have been accessed exactly once for get + // L0 still uses get(); raw-capable lower levels use one get_raw() to + // both parse and backfill the hit. assert_eq!(l0_accesses.len(), 1, "L0 should be checked first"); - assert_eq!(l1_accesses.len(), 1, "L1 should be checked second"); - assert_eq!(l2_accesses.len(), 2, "L2: put (setup) + get (check)"); + assert_eq!(l1_accesses.len(), 0, "L1 should use raw lookup"); + assert_eq!(l2_accesses.len(), 1, "L2 should contain setup put only"); + assert_eq!( + l1_raw_log.lock().await.as_slice(), + &[format!("get_raw:{}", key)] + ); + assert_eq!( + l2_raw_log.lock().await.as_slice(), + &[format!("get_raw:{}", key)] + ); assert_eq!(l0_accesses[0], format!("get:{}", key)); - assert_eq!(l1_accesses[0], format!("get:{}", key)); assert_eq!(l2_accesses[0], format!("put:{}", key)); // from setup - assert_eq!(l2_accesses[1], format!("get:{}", key)); // from sequential check }); } @@ -1039,6 +1107,7 @@ fn test_read_stops_at_first_hit_not_parallel() { let l0_log = l0.get_access_log(); let l1_log = l1.get_access_log(); let l2_log = l2.get_access_log(); + let l1_raw_log = l1.get_raw_access_log(); let key = "test_key_early_hit_1234567890ab"; @@ -1066,7 +1135,15 @@ fn test_read_stops_at_first_hit_not_parallel() { let l2_accesses = l2_log.lock().await; assert_eq!(l0_accesses.len(), 1, "L0 should be checked first"); - assert_eq!(l1_accesses.len(), 2, "L1: put (setup) + get (check)"); + assert_eq!( + l1_accesses.len(), + 1, + "L1: put (setup); lookup uses raw access" + ); + assert_eq!( + l1_raw_log.lock().await.as_slice(), + &[format!("get_raw:{}", key)] + ); assert_eq!( l2_accesses.len(), 0, diff --git a/src/cache/readonly.rs b/src/cache/readonly.rs index 40f9873f5..231f99168 100644 --- a/src/cache/readonly.rs +++ b/src/cache/readonly.rs @@ -29,6 +29,10 @@ impl Storage for ReadOnlyStorage { self.0.get(key).await } + async fn get_with_raw(&self, key: &str) -> Result<(Cache, Option)> { + self.0.get_with_raw(key).await + } + /// Put `entry` in the cache under `key`. /// /// Returns a `Future` that will provide the result or error when the put is