Use new ChecksummedBlock in DataCache

mountpoint-s3/src/checksums.rs

@@ -0,0 +1,39 @@
+mod block;
+mod bytes;
+
+use mountpoint_s3_crt::checksums::crc32c::Crc32c;
+
+use thiserror::Error;
+
+pub use block::ChecksummedBlock;
+pub use bytes::ChecksummedBytes;
+
+/// Calculates the combined checksum for `AB` where `prefix_crc` is the checksum for `A`,
+/// `suffix_crc` is the checksum for `B`, and `suffic_len` is the length of `B`.
+pub fn combine_checksums(prefix_crc: Crc32c, suffix_crc: Crc32c, suffix_len: usize) -> Crc32c {
+    let combined = ::crc32c::crc32c_combine(prefix_crc.value(), suffix_crc.value(), suffix_len);
+    Crc32c::new(combined)
+}
+
+#[derive(Debug, Error)]
+pub enum IntegrityError {
+    #[error("Checksum mismatch. expected: {0:?}, actual: {1:?}")]
+    ChecksumMismatch(Crc32c, Crc32c),
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use mountpoint_s3_crt::checksums::crc32c;
+
+    #[test]
+    fn test_combine_checksums() {
+        let buf: &[u8] = b"123456789";
+        let (buf1, buf2) = buf.split_at(4);
+        let crc = crc32c::checksum(buf);
+        let crc1 = crc32c::checksum(buf1);
+        let crc2 = crc32c::checksum(buf2);
+        let combined = combine_checksums(crc1, crc2, buf2.len());
+        assert_eq!(combined, crc);
+    }
+}

mountpoint-s3/src/checksums/block.rs

@@ -0,0 +1,211 @@
+use bytes::{Bytes, BytesMut};
+use mountpoint_s3_crt::checksums::crc32c::{self, Crc32c};
+
+use crate::checksums::{bytes::ChecksummedBytes, combine_checksums, IntegrityError};
+
+/// A `ChecksummedBlock` is a bytes buffer that carries its checksum.
+/// The implementation guarantees that its integrity will be validated when data is accessed.
+#[derive(Debug, Clone)]
+pub struct ChecksummedBlock {
+    bytes: Bytes,
+    checksum: Crc32c,
+}
+
+impl ChecksummedBlock {
+    pub fn new(bytes: Bytes, checksum: Crc32c) -> Self {
+        Self { bytes, checksum }
+    }
+
+    /// Create `ChecksummedBlock` from `Bytes`, calculating its checksum.
+    pub fn from_bytes(bytes: Bytes) -> Self {
+        let checksum = crc32c::checksum(&bytes);
+        Self::new(bytes, checksum)
+    }
+
+    /// Convert the `ChecksummedBlock` into `Bytes`, data integrity will be validated before converting.
+    ///
+    /// Return `IntegrityError` on data corruption.
+    pub fn into_bytes(self) -> Result<Bytes, IntegrityError> {
+        self.validate()?;
+
+        Ok(self.bytes)
+    }
+
+    /// Convert into a `ChecksummedBytes`.
+    pub fn into_checksummed_bytes(self) -> ChecksummedBytes {
+        ChecksummedBytes::new(self.bytes, self.checksum)
+    }
+
+    /// Returns the number of bytes contained in this `ChecksummedBlock`.
+    pub fn len(&self) -> usize {
+        self.bytes.len()
+    }
+
+    /// Returns true if the `ChecksummedBlock` has a length of 0.
+    pub fn is_empty(&self) -> bool {
+        self.bytes.is_empty()
+    }
+
+    /// Append the given bytes to current `ChecksummedBlock`.
+    pub fn extend(&mut self, extend: ChecksummedBlock) {
+        if self.is_empty() {
+            *self = extend;
+            return;
+        }
+        if extend.is_empty() {
+            return;
+        }
+
+        let total_len = self.bytes.len() + extend.len();
+        let mut bytes_mut = BytesMut::with_capacity(total_len);
+        bytes_mut.extend_from_slice(&self.bytes);
+        bytes_mut.extend_from_slice(&extend.bytes);
+        let new_bytes = bytes_mut.freeze();
+        let new_checksum = combine_checksums(self.checksum, extend.checksum, extend.len());
+        *self = ChecksummedBlock {
+            bytes: new_bytes,
+            checksum: new_checksum,
+        };
+    }
+
+    /// Validate data integrity in this `ChecksummedBlock`.
+    ///
+    /// Return `IntegrityError` on data corruption.
+    pub fn validate(&self) -> Result<(), IntegrityError> {
+        let checksum = crc32c::checksum(&self.bytes);
+        if self.checksum != checksum {
+            return Err(IntegrityError::ChecksumMismatch(self.checksum, checksum));
+        }
+        Ok(())
+    }
+}
+
+impl Default for ChecksummedBlock {
+    fn default() -> Self {
+        let bytes = Bytes::new();
+        let checksum = Crc32c::new(0);
+        Self { bytes, checksum }
+    }
+}
+
+impl From<ChecksummedBlock> for ChecksummedBytes {
+    fn from(value: ChecksummedBlock) -> Self {
+        value.into_checksummed_bytes()
+    }
+}
+
+// Implement equality for tests only. We implement equality, and will panic if the data is corrupted.
+#[cfg(test)]
+impl PartialEq for ChecksummedBlock {
+    fn eq(&self, other: &Self) -> bool {
+        if self.bytes != other.bytes {
+            return false;
+        }
+
+        if self.checksum == other.checksum {
+            return true;
+        }
+
+        self.validate().expect("should be valid");
+        other.validate().expect("should be valid");
+
+        true
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use bytes::Bytes;
+    use mountpoint_s3_crt::checksums::crc32c;
+
+    use super::*;
+
+    #[test]
+    fn test_into_bytes() {
+        let bytes = Bytes::from_static(b"some bytes");
+        let expected = bytes.clone();
+        let checksum = crc32c::checksum(&bytes);
+        let checksummed_block = ChecksummedBlock::new(bytes, checksum);
+
+        let actual = checksummed_block.into_bytes().unwrap();
+        assert_eq!(expected, actual);
+    }
+
+    #[test]
+    fn test_into_bytes_integrity_error() {
+        let bytes = Bytes::from_static(b"some bytes");
+        let checksum = crc32c::checksum(&bytes);
+        let mut checksummed_block = ChecksummedBlock::new(bytes, checksum);
+        checksummed_block.bytes = Bytes::from_static(b"new bytes");
+
+        let actual = checksummed_block.into_bytes();
+        assert!(matches!(actual, Err(IntegrityError::ChecksumMismatch(_, _))));
+    }
+
+    #[test]
+    fn test_into_checksummed_bytes() {
+        let bytes = Bytes::from_static(b"some bytes");
+        let checksum = crc32c::checksum(&bytes);
+        let checksummed_block = ChecksummedBlock::new(bytes, checksum);
+        let checksummed_bytes = checksummed_block.clone().into_checksummed_bytes();
+
+        assert_eq!(
+            checksummed_block.into_bytes().unwrap(),
+            checksummed_bytes.into_bytes().unwrap()
+        );
+    }
+
+    #[test]
+    fn test_extend() {
+        let bytes = Bytes::from_static(b"some bytes");
+        let expected = Bytes::from_static(b"some bytes extended");
+        let checksum = crc32c::checksum(&bytes);
+        let mut checksummed_block = ChecksummedBlock::new(bytes, checksum);
+
+        let extend = Bytes::from_static(b" extended");
+        let extend_checksum = crc32c::checksum(&extend);
+        let extend = ChecksummedBlock::new(extend, extend_checksum);
+        checksummed_block.extend(extend);
+        let actual = checksummed_block.bytes;
+        assert_eq!(expected, actual);
+    }
+
+    #[test]
+    fn test_extend_self_corrupted() {
+        let bytes = Bytes::from_static(b"some bytes");
+        let checksum = crc32c::checksum(&bytes);
+        let mut checksummed_block = ChecksummedBlock::new(bytes, checksum);
+
+        let currupted_bytes = Bytes::from_static(b"corrupted data");
+        checksummed_block.bytes = currupted_bytes.clone();
+
+        let extend = Bytes::from_static(b" extended");
+        let extend_checksum = crc32c::checksum(&extend);
+        let extend = ChecksummedBlock::new(extend, extend_checksum);
+        checksummed_block.extend(extend);
+        assert!(matches!(
+            checksummed_block.validate(),
+            Err(IntegrityError::ChecksumMismatch(_, _))
+        ));
+    }
+
+    #[test]
+    fn test_extend_other_corrupted() {
+        let bytes = Bytes::from_static(b"some bytes");
+        let checksum = crc32c::checksum(&bytes);
+        let mut checksummed_block = ChecksummedBlock::new(bytes, checksum);
+
+        let extend = Bytes::from_static(b" extended");
+        let extend_checksum = crc32c::checksum(&extend);
+        let mut extend = ChecksummedBlock::new(extend, extend_checksum);
+
+        let currupted_bytes = Bytes::from_static(b"corrupted data");
+        extend.bytes = currupted_bytes.clone();
+
+        checksummed_block.extend(extend);
+        assert!(matches!(
+            checksummed_block.validate(),
+            Err(IntegrityError::ChecksumMismatch(_, _))
+        ));
+    }
+}

mountpoint-s3/src/prefetch/checksummed_bytes.rs → mountpoint-s3/src/checksums/bytes.rs

@@ -1,6 +1,9 @@
+use std::ops::RangeBounds;
+
 use bytes::{Bytes, BytesMut};
 use mountpoint_s3_crt::checksums::crc32c::{self, Crc32c};
-use thiserror::Error;
+
+use crate::checksums::IntegrityError;
 
 /// A `ChecksummedBytes` is a bytes buffer that carries its checksum.
 /// The implementation guarantees that its integrity will be validated when data transformation occurs.
@@ -56,6 +59,18 @@ impl ChecksummedBytes {
         }
     }
 
+    /// Returns a slice of self for the provided range.
+    ///
+    /// This operation just increases the reference count and sets a few indices,
+    /// so there will be no validation and the checksum will not be recomputed.
+    pub fn slice(&self, range: impl RangeBounds<usize>) -> Self {
+        Self {
+            orig_bytes: self.orig_bytes.clone(),
+            curr_slice: self.curr_slice.slice(range),
+            checksum: self.checksum,
+        }
+    }
+
     /// Append the given checksummed bytes to current `ChecksummedBytes`, ensure that data integrity will
     /// be validated.
     ///
@@ -118,13 +133,6 @@ impl Default for ChecksummedBytes {
         }
     }
 }
-
-#[derive(Debug, Error)]
-pub enum IntegrityError {
-    #[error("Checksum mismatch. expected: {0:?}, actual: {1:?}")]
-    ChecksumMismatch(Crc32c, Crc32c),
-}
-
 // Implement equality for tests only. We implement equality, and will panic if the data is corrupted.
 #[cfg(test)]
 impl PartialEq for ChecksummedBytes {
@@ -149,9 +157,7 @@ mod tests {
     use bytes::Bytes;
     use mountpoint_s3_crt::checksums::crc32c;
 
-    use crate::prefetch::checksummed_bytes::IntegrityError;
-
-    use super::ChecksummedBytes;
+    use super::*;
 
     #[test]
     fn test_into_bytes() {

mountpoint-s3/src/data_cache.rs

@@ -10,7 +10,7 @@ use std::ops::Range;
 
 use thiserror::Error;
 
-pub use crate::prefetch::checksummed_bytes::ChecksummedBytes;
+pub use crate::checksums::ChecksummedBlock;
 
 /// Indexes blocks within a given object.
 pub type BlockIndex = u64;
@@ -36,10 +36,10 @@ pub trait DataCache<Key> {
     /// Get block of data from the cache for the given [Key] and [BlockIndex], if available.
     ///
     /// Operation may fail due to errors, or return [None] if the block was not available in the cache.
-    fn get_block(&self, cache_key: &Key, block_idx: BlockIndex) -> DataCacheResult<Option<ChecksummedBytes>>;
+    fn get_block(&self, cache_key: &Key, block_idx: BlockIndex) -> DataCacheResult<Option<ChecksummedBlock>>;
 
     /// Put block of data to the cache for the given [Key] and [BlockIndex].
-    fn put_block(&self, cache_key: Key, block_idx: BlockIndex, bytes: ChecksummedBytes) -> DataCacheResult<()>;
+    fn put_block(&self, cache_key: Key, block_idx: BlockIndex, bytes: ChecksummedBlock) -> DataCacheResult<()>;
 
     /// Returns the block size for the data cache.
     fn block_size(&self) -> u64;

mountpoint-s3/src/data_cache/in_memory_data_cache.rs

@@ -5,12 +5,12 @@ use std::default::Default;
 use std::hash::Hash;
 use std::ops::Range;
 
-use super::{BlockIndex, ChecksummedBytes, DataCache, DataCacheResult};
+use super::{BlockIndex, ChecksummedBlock, DataCache, DataCacheResult};
 use crate::sync::RwLock;
 
 /// Simple in-memory (RAM) implementation of [DataCache]. Recommended for use in testing only.
 pub struct InMemoryDataCache<CacheKey> {
-    data: RwLock<HashMap<CacheKey, HashMap<BlockIndex, ChecksummedBytes>>>,
+    data: RwLock<HashMap<CacheKey, HashMap<BlockIndex, ChecksummedBlock>>>,
     block_size: u64,
 }
 
@@ -25,13 +25,13 @@ impl<Key> InMemoryDataCache<Key> {
 }
 
 impl<Key: Eq + Hash> DataCache<Key> for InMemoryDataCache<Key> {
-    fn get_block(&self, cache_key: &Key, block_idx: BlockIndex) -> DataCacheResult<Option<ChecksummedBytes>> {
+    fn get_block(&self, cache_key: &Key, block_idx: BlockIndex) -> DataCacheResult<Option<ChecksummedBlock>> {
         let data = self.data.read().unwrap();
         let block_data = data.get(cache_key).and_then(|blocks| blocks.get(&block_idx)).cloned();
         Ok(block_data)
     }
 
-    fn put_block(&self, cache_key: Key, block_idx: BlockIndex, bytes: ChecksummedBytes) -> DataCacheResult<()> {
+    fn put_block(&self, cache_key: Key, block_idx: BlockIndex, bytes: ChecksummedBlock) -> DataCacheResult<()> {
         let mut data = self.data.write().unwrap();
         let blocks = data.entry(cache_key).or_default();
         blocks.insert(block_idx, bytes);
@@ -59,20 +59,18 @@ mod tests {
 
     use bytes::Bytes;
 
-    use mountpoint_s3_crt::checksums::crc32c;
-
     type TestCacheKey = String;
 
     #[test]
     fn test_put_get() {
         let data_1 = Bytes::from_static(b"Hello world");
-        let data_1 = ChecksummedBytes::new(data_1.clone(), crc32c::checksum(&data_1));
+        let data_1 = ChecksummedBlock::from_bytes(data_1.clone());
         let data_2 = Bytes::from_static(b"Foo bar");
-        let data_2 = ChecksummedBytes::new(data_2.clone(), crc32c::checksum(&data_2));
+        let data_2 = ChecksummedBlock::from_bytes(data_2.clone());
         let data_3 = Bytes::from_static(b"Baz");
-        let data_3 = ChecksummedBytes::new(data_3.clone(), crc32c::checksum(&data_3));
+        let data_3 = ChecksummedBlock::from_bytes(data_3.clone());
 
-        let mut cache = InMemoryDataCache::new(8 * 1024 * 1024);
+        let cache = InMemoryDataCache::new(8 * 1024 * 1024);
         let cache_key_1: TestCacheKey = String::from("a");
         let cache_key_2: TestCacheKey = String::from("b");
 
@@ -136,11 +134,11 @@ mod tests {
     #[test]
     fn test_cached_indices() {
         let data_1 = Bytes::from_static(b"Hello world");
-        let data_1 = ChecksummedBytes::new(data_1.clone(), crc32c::checksum(&data_1));
+        let data_1 = ChecksummedBlock::from_bytes(data_1.clone());
         let data_2 = Bytes::from_static(b"Foo bar");
-        let data_2 = ChecksummedBytes::new(data_2.clone(), crc32c::checksum(&data_2));
+        let data_2 = ChecksummedBlock::from_bytes(data_2.clone());
 
-        let mut cache = InMemoryDataCache::new(8 * 1024 * 1024);
+        let cache = InMemoryDataCache::new(8 * 1024 * 1024);
         let cache_key_1: TestCacheKey = String::from("a");
         let cache_key_2: TestCacheKey = String::from("b");