Use Bytes in get_raw/put_raw for zero-copy

Author: Felixoid

src/cache/cache.rs

@@ -49,6 +49,7 @@ use crate::compiler::PreprocessorCacheEntry;
 use crate::config::Config;
 use crate::config::{self, CacheType, PreprocessorCacheModeConfig};
 use async_trait::async_trait;
+use bytes::Bytes;
 
 use std::io;
 use std::sync::Arc;
@@ -77,14 +78,14 @@ pub trait Storage: Send + Sync {
     /// Get raw serialized cache entry bytes by `key` (for multi-level backfill).
     /// Returns `None` if the entry is not found, or if the implementation doesn't support raw access.
     /// This is used by multi-level caches to backfill faster levels.
-    async fn get_raw(&self, _key: &str) -> Result<Option<Vec<u8>>> {
+    async fn get_raw(&self, _key: &str) -> Result<Option<Bytes>> {
         Ok(None)
     }
 
     /// Put raw serialized cache entry bytes under `key` (for multi-level backfill).
     /// Returns an error if the implementation doesn't support raw access.
     /// This is used by multi-level caches to backfill faster levels.
-    async fn put_raw(&self, _key: &str, _data: Vec<u8>) -> Result<Duration> {
+    async fn put_raw(&self, _key: &str, _data: Bytes) -> Result<Duration> {
         Err(anyhow!("put_raw not implemented for this storage backend"))
     }
 
@@ -220,7 +221,7 @@ impl Storage for RemoteStorage {
         trace!("RemoteStorage::put({})", key);
         // Delegate to put_raw after serializing the entry
         let data = entry.finish()?;
-        self.put_raw(key, data).await
+        self.put_raw(key, data.into()).await
     }
 
     async fn check(&self) -> Result<CacheMode> {
@@ -297,18 +298,17 @@ impl Storage for RemoteStorage {
         &self.basedirs
     }
 
-    /// Get raw bytes from remote storage without any transformations.
+    /// Get raw bytes from remote storage for multi-level backfill.
     ///
-    /// Uses `to_vec()` instead of `to_bytes()` to preserve raw data unchanged.
-    /// This is critical for multi-level caching: when backfilling from remote to local,
-    /// we need the exact bytes without OpenDAL layer transformations (e.g., decompression).
-    /// If compression layers are configured, `to_bytes()` would decompress the data,
-    /// which would corrupt the cache entry when written to another level.
-    async fn get_raw(&self, key: &str) -> Result<Option<Vec<u8>>> {
+    /// Unlike `get()` which parses bytes into `CacheRead` (a `ZipArchive<Box<dyn ReadSeek>>`),
+    /// this returns the raw bytes without parsing. `CacheRead` is a one-way transformation —
+    /// there is no way to extract the original bytes back from the parsed ZIP archive.
+    /// For backfill we need the raw bytes to write directly to another cache level.
+    async fn get_raw(&self, key: &str) -> Result<Option<Bytes>> {
         trace!("opendal::Operator::get_raw({})", key);
         match self.operator.read(&normalize_key(key)).await {
             Ok(res) => {
-                let data = res.to_vec();
+                let data = res.to_bytes();
                 trace!(
                     "opendal::Operator::get_raw({}): Found {} bytes",
                     key,
@@ -328,12 +328,12 @@ impl Storage for RemoteStorage {
         }
     }
 
-    /// Write raw bytes to remote storage.
+    /// Write raw bytes to remote storage for multi-level backfill.
     ///
-    /// This is the primitive write operation used by both `put()` and multi-level backfill.
-    /// For backfill operations, raw bytes are passed directly from one cache level to another
-    /// to preserve the exact data format (including any compression applied by OpenDAL layers).
-    async fn put_raw(&self, key: &str, data: Vec<u8>) -> Result<Duration> {
+    /// Unlike `put()` which takes a `CacheWrite` and serializes it, this writes
+    /// pre-serialized bytes directly. Paired with `get_raw()` for efficient
+    /// level-to-level data transfer without a deserialize/reserialize round-trip.
+    async fn put_raw(&self, key: &str, data: Bytes) -> Result<Duration> {
         trace!("opendal::Operator::put_raw({}, {} bytes)", key, data.len());
         let start = std::time::Instant::now();
 

src/cache/disk.rs

@@ -16,6 +16,7 @@ use crate::cache::{Cache, CacheMode, CacheRead, CacheWrite, Storage};
 use crate::compiler::PreprocessorCacheEntry;
 use crate::lru_disk_cache::{Error as LruError, ReadSeek};
 use async_trait::async_trait;
+use bytes::Bytes;
 use std::ffi::OsStr;
 use std::io::{BufWriter, Read, Write};
 use std::path::{Path, PathBuf};
@@ -102,7 +103,7 @@ impl Storage for DiskCache {
             .await?
     }
 
-    async fn get_raw(&self, key: &str) -> Result<Option<Vec<u8>>> {
+    async fn get_raw(&self, key: &str) -> Result<Option<Bytes>> {
         trace!("DiskCache::get_raw({})", key);
         let path = make_key_path(key);
         let lru = self.lru.clone();
@@ -115,7 +116,7 @@ impl Storage for DiskCache {
                         let mut data = Vec::new();
                         io.read_to_end(&mut data)?;
                         trace!("DiskCache::get_raw({}): Found {} bytes", key, data.len());
-                        Ok(Some(data))
+                        Ok(Some(Bytes::from(data)))
                     }
                     Err(LruError::FileNotInCache) => {
                         trace!("DiskCache::get_raw({}): FileNotInCache", key);
@@ -135,10 +136,10 @@ impl Storage for DiskCache {
         trace!("DiskCache::put({})", key);
         // Delegate to put_raw after serializing the entry
         let data = entry.finish()?;
-        self.put_raw(key, data).await
+        self.put_raw(key, data.into()).await
     }
 
-    async fn put_raw(&self, key: &str, data: Vec<u8>) -> Result<Duration> {
+    async fn put_raw(&self, key: &str, data: Bytes) -> Result<Duration> {
         trace!("DiskCache::put_raw({}, {} bytes)", key, data.len());
 
         if self.rw_mode == CacheMode::ReadOnly {

src/cache/multilevel.rs

@@ -12,6 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use bytes::Bytes;
 use std::sync::Arc;
 use std::sync::atomic::{AtomicU64, Ordering};
 use std::time::{Duration, Instant};
@@ -519,23 +520,23 @@ impl MultiLevelStorage {
     async fn write_entry_from_bytes(
         level: &Arc<dyn Storage>,
         key: &str,
-        data: &Arc<Vec<u8>>,
+        data: &Bytes,
     ) -> Result<()> {
-        // Try to use put_raw for direct bytes write (most efficient)
-        level.put_raw(key, (**data).clone()).await?;
+        // Bytes::clone() is a cheap ref-count bump, no data copy
+        level.put_raw(key, data.clone()).await?;
         Ok(())
     }
 
     /// Write to levels starting from `start_idx` asynchronously
-    async fn write_remaining_levels_async(&self, key: &str, data: &Arc<Vec<u8>>, start_idx: usize) {
+    async fn write_remaining_levels_async(&self, key: &str, data: &Bytes, start_idx: usize) {
         for (idx, level) in self.levels.iter().enumerate().skip(start_idx) {
             // Check if level is read-only before spawning task
             if matches!(level.check().await, Ok(CacheMode::ReadOnly)) {
                 debug!("Level {} is read-only, skipping write", idx);
                 continue;
             }
 
-            let data = Arc::clone(data);
+            let data = data.clone();
             let key = key.to_string();
             let level = Arc::clone(level);
             let stats_arc = self.atomic_stats.get(idx).map(Arc::clone);
@@ -597,8 +598,6 @@ impl Storage for MultiLevelStorage {
                         // Try to get raw bytes for backfilling
                         match level.get_raw(key).await {
                             Ok(Some(raw_bytes)) => {
-                                let raw_bytes = Arc::new(raw_bytes);
-
                                 // Update backfill stats
                                 if let Some(stats) = self.atomic_stats.get(hit_level) {
                                     stats
@@ -610,7 +609,7 @@ impl Storage for MultiLevelStorage {
                                 // Iterate slice directly instead of creating Vec
                                 for backfill_idx in 0..idx {
                                     let key_bf = key_str.clone();
-                                    let bytes_bf = Arc::clone(&raw_bytes);
+                                    let bytes_bf = raw_bytes.clone();
                                     let level_bf = Arc::clone(&self.levels[backfill_idx]);
                                     let stats_arc =
                                         self.atomic_stats.get(backfill_idx).map(Arc::clone);
@@ -694,7 +693,7 @@ impl Storage for MultiLevelStorage {
         }
 
         // Serialize cache entry once
-        let data = Arc::new(entry.finish()?);
+        let data: Bytes = entry.finish()?.into();
         let key_str = key.to_string();
 
         match self.write_policy {
@@ -738,7 +737,7 @@ impl Storage for MultiLevelStorage {
                 let (tx, mut rx) = mpsc::channel(self.levels.len());
 
                 for (idx, level) in self.levels.iter().enumerate() {
-                    let data = Arc::clone(&data);
+                    let data = data.clone();
                     let key_str = key_str.clone();
                     let level = Arc::clone(level);
                     let tx = tx.clone();

src/cache/multilevel_test.rs

@@ -18,6 +18,7 @@ use crate::cache::disk::DiskCache;
 use crate::cache::readonly::ReadOnlyStorage;
 use crate::config::Config;
 use crate::config::PreprocessorCacheModeConfig;
+use bytes::Bytes;
 use std::collections::HashMap;
 use std::env;
 use std::fs;
@@ -256,13 +257,16 @@ impl Storage for InMemoryStorage {
     /// Implement get_raw() to enable backfill testing with remote-like backends.
     /// 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<Option<Vec<u8>>> {
-        Ok(self.data.lock().await.get(key).cloned())
+    async fn get_raw(&self, key: &str) -> Result<Option<Bytes>> {
+        Ok(self.data.lock().await.get(key).cloned().map(Bytes::from))
     }
 
     /// Implement put_raw() to enable backfill writes during testing.
-    async fn put_raw(&self, key: &str, data: Vec<u8>) -> Result<Duration> {
-        self.data.lock().await.insert(key.to_string(), data);
+    async fn put_raw(&self, key: &str, data: Bytes) -> Result<Duration> {
+        self.data
+            .lock()
+            .await
+            .insert(key.to_string(), data.to_vec());
         Ok(Duration::ZERO)
     }
 }
@@ -1087,7 +1091,7 @@ impl Storage for FailingStorage {
         Err(anyhow!("Intentional failure for testing"))
     }
 
-    async fn put_raw(&self, _key: &str, _entry: Vec<u8>) -> Result<Duration> {
+    async fn put_raw(&self, _key: &str, _entry: Bytes) -> Result<Duration> {
         Err(anyhow!("Intentional failure for testing"))
     }
 

src/cache/readonly.rs

@@ -19,6 +19,7 @@ use crate::cache::{Cache, CacheMode, CacheWrite, Storage};
 use crate::compiler::PreprocessorCacheEntry;
 use crate::config::PreprocessorCacheModeConfig;
 use crate::errors::*;
+use bytes::Bytes;
 
 pub struct ReadOnlyStorage(pub Arc<dyn Storage>);
 
@@ -95,7 +96,7 @@ impl Storage for ReadOnlyStorage {
     }
 
     /// Get raw serialized cache entry bytes (forwarded to inner storage)
-    async fn get_raw(&self, key: &str) -> Result<Option<Vec<u8>>> {
+    async fn get_raw(&self, key: &str) -> Result<Option<Bytes>> {
         self.0.get_raw(key).await
     }
 }