armdb 0.2.0

sharded bitcask key-value storage optimized for NVMe
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114

use std::sync::Arc;

use quick_cache::Weighter;
use quick_cache::sync::Cache;

use crate::config::CacheConfig;
use crate::io::aligned_buf::AlignedBuf;

/// Cache key: identifies a 4096-byte aligned block on disk.
/// Includes `shard_id` because file_ids are assigned independently per shard
/// (each shard starts at file_id=1), so (file_id, block_offset) alone is ambiguous.
#[derive(Clone, Copy, PartialEq, Eq, Hash)]
pub struct BlockKey {
    pub shard_id: u8,
    pub file_id: u32,
    pub block_offset: u64,
}

#[derive(Clone)]
struct BlockWeighter;

impl Weighter<BlockKey, Arc<AlignedBuf>> for BlockWeighter {
    fn weight(&self, _key: &BlockKey, _value: &Arc<AlignedBuf>) -> u64 {
        4096
    }
}

/// Block cache for `VarTree`, backed by `quick_cache` with S3-FIFO eviction.
/// Caches 4096-byte aligned disk blocks — the natural unit of O_DIRECT reads.
pub struct BlockCache {
    inner: Option<Cache<BlockKey, Arc<AlignedBuf>, BlockWeighter>>,
}

impl BlockCache {
    pub fn new(config: &CacheConfig) -> Self {
        if config.max_size == 0 {
            return Self { inner: None };
        }
        let cache = Cache::with_weighter(config.estimated_items, config.max_size, BlockWeighter);
        Self { inner: Some(cache) }
    }

    pub fn get(&self, key: &BlockKey) -> Option<Arc<AlignedBuf>> {
        self.inner.as_ref()?.get(key)
    }

    pub fn insert(&self, key: BlockKey, block: Arc<AlignedBuf>) {
        if let Some(cache) = &self.inner {
            cache.insert(key, block);
        }
    }

    /// Invalidate all cached blocks for a file. Called after compaction
    /// replaces a file's contents (file_id reuse).
    pub fn invalidate_file(&self, shard_id: u8, file_id: u32, total_bytes: u64) {
        if let Some(cache) = &self.inner {
            let num_blocks = total_bytes.div_ceil(4096);
            for i in 0..num_blocks {
                cache.remove(&BlockKey {
                    shard_id,
                    file_id,
                    block_offset: i * 4096,
                });
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::config::CacheConfig;
    use crate::io::aligned_buf::AlignedBuf;

    fn make_marked_buf(marker: u8) -> Arc<AlignedBuf> {
        let mut buf = AlignedBuf::zeroed(4096);
        buf[0] = marker;
        Arc::new(buf)
    }

    #[test]
    fn block_cache_disambiguates_file_ids_across_u16_boundary() {
        let cfg = CacheConfig {
            max_size: 1 << 20,
            estimated_items: 32,
        };
        let cache = BlockCache::new(&cfg);

        // Build two distinguishable 4 KiB buffers.
        let buf_a = make_marked_buf(0xAA);
        let buf_b = make_marked_buf(0xBB);

        let key_low = BlockKey {
            shard_id: 0,
            file_id: 65_535,
            block_offset: 0,
        };
        let key_high = BlockKey {
            shard_id: 0,
            file_id: 65_536,
            block_offset: 0,
        };

        cache.insert(key_low, buf_a.clone());
        cache.insert(key_high, buf_b.clone());

        let got_low = cache.get(&key_low).expect("low present");
        let got_high = cache.get(&key_high).expect("high present");

        assert_eq!(got_low[0], 0xAA);
        assert_eq!(got_high[0], 0xBB);
        assert!(!Arc::ptr_eq(&got_low, &got_high));
    }
}