skippy_cache/payload/
blob_store.rs

1use std::{collections::HashMap, sync::Arc, time::Instant};
2
3use super::CacheBytes;
4use super::bytes::{CacheBlockRef, CacheBytesRepr};
5
6const DEFAULT_BLOCK_SIZE_BYTES: usize = 1024 * 1024;
7
8#[derive(Debug)]
9pub struct CacheBlobStore {
10    block_size: usize,
11    physical_bytes: u64,
12    blocks: HashMap<String, CacheBlob>,
13}
14
15impl Default for CacheBlobStore {
16    fn default() -> Self {
17        Self::new(DEFAULT_BLOCK_SIZE_BYTES)
18    }
19}
20
21#[derive(Debug)]
22struct CacheBlob {
23    bytes: Arc<Vec<u8>>,
24    ref_count: u64,
25}
26
27impl CacheBlobStore {
28    pub fn new(block_size: usize) -> Self {
29        Self {
30            block_size: block_size.max(1),
31            physical_bytes: 0,
32            blocks: HashMap::new(),
33        }
34    }
35
36    pub fn store_bytes(&mut self, bytes: CacheBytes) -> (CacheBytes, CacheDedupeStats) {
37        let len = bytes.len;
38        let bytes = match bytes.repr {
39            CacheBytesRepr::Inline(bytes) => bytes,
40            repr @ CacheBytesRepr::Blocks(_) => {
41                return (CacheBytes { len, repr }, CacheDedupeStats::default());
42            }
43        };
44        let mut blocks = Vec::new();
45        let started = Instant::now();
46        let mut stats = CacheDedupeStats {
47            hash_bytes: bytes.len() as u64,
48            ..CacheDedupeStats::default()
49        };
50        for chunk in bytes.chunks(self.block_size) {
51            stats.block_count = stats.block_count.saturating_add(1);
52            let hash = blake3::hash(chunk).to_hex().to_string();
53            let entry = self.blocks.entry(hash.clone()).or_insert_with(|| {
54                self.physical_bytes = self.physical_bytes.saturating_add(chunk.len() as u64);
55                stats.new_block_count = stats.new_block_count.saturating_add(1);
56                CacheBlob {
57                    bytes: Arc::new(chunk.to_vec()),
58                    ref_count: 0,
59                }
60            });
61            if entry.ref_count > 0 {
62                stats.reused_block_count = stats.reused_block_count.saturating_add(1);
63            }
64            entry.ref_count = entry.ref_count.saturating_add(1);
65            blocks.push(CacheBlockRef::new(hash, entry.bytes.clone()));
66        }
67        stats.hash_ms = started.elapsed().as_secs_f64() * 1000.0;
68        (CacheBytes::blocks(bytes.len() as u64, blocks), stats)
69    }
70
71    pub fn release_bytes(&mut self, bytes: &CacheBytes) {
72        let hashes = bytes.block_hashes().map(str::to_string).collect::<Vec<_>>();
73        for hash in hashes {
74            let mut remove = false;
75            if let Some(entry) = self.blocks.get_mut(&hash) {
76                entry.ref_count = entry.ref_count.saturating_sub(1);
77                if entry.ref_count == 0 {
78                    self.physical_bytes =
79                        self.physical_bytes.saturating_sub(entry.bytes.len() as u64);
80                    remove = true;
81                }
82            }
83            if remove {
84                self.blocks.remove(&hash);
85            }
86        }
87    }
88
89    pub fn physical_bytes(&self) -> u64 {
90        self.physical_bytes
91    }
92
93    pub fn block_count(&self) -> usize {
94        self.blocks.len()
95    }
96}
97
98#[derive(Debug, Clone, Copy, Default)]
99pub struct CacheDedupeStats {
100    pub hash_ms: f64,
101    pub hash_bytes: u64,
102    pub block_count: usize,
103    pub new_block_count: usize,
104    pub reused_block_count: usize,
105}
106
107impl CacheDedupeStats {
108    pub fn saturating_add(self, other: Self) -> Self {
109        Self {
110            hash_ms: self.hash_ms + other.hash_ms,
111            hash_bytes: self.hash_bytes.saturating_add(other.hash_bytes),
112            block_count: self.block_count.saturating_add(other.block_count),
113            new_block_count: self.new_block_count.saturating_add(other.new_block_count),
114            reused_block_count: self
115                .reused_block_count
116                .saturating_add(other.reused_block_count),
117        }
118    }
119}
120
121#[cfg(test)]
122mod tests {
123    use crate::payload::{CacheBlobStore, ExactStatePayload};
124
125    #[test]
126    fn block_store_dedupes_repeated_payload_blocks() {
127        let mut blobs = CacheBlobStore::new(4);
128        let first = ExactStatePayload::full_state(b"aaaabbbb".to_vec());
129        let second = ExactStatePayload::full_state(b"aaaacccc".to_vec());
130
131        let (_, first_stats) = first.dedupe_into(&mut blobs);
132        let (second, second_stats) = second.dedupe_into(&mut blobs);
133
134        assert_eq!(first_stats.new_block_count, 2);
135        assert_eq!(second_stats.new_block_count, 1);
136        assert_eq!(second_stats.reused_block_count, 1);
137        assert_eq!(blobs.physical_bytes(), 12);
138        assert_eq!(second.byte_len(), 8);
139    }
140}
skippy_cache/payload/blob_store.rs

skippy_cache/payload/
blob_store.rs
