doublecrypt_core/
allocator.rs

1use crate::error::{FsError, FsResult};
2use crate::model::FIRST_DATA_BLOCK;
3use std::collections::BTreeSet;
4use std::sync::Mutex;
5
6/// Trait for allocating and freeing block slots.
7pub trait SlotAllocator: Send + Sync {
8    /// Allocate a free block and return its ID.
9    fn allocate(&self) -> FsResult<u64>;
10
11    /// Return a block to the free pool.
12    fn free(&self, block_id: u64) -> FsResult<()>;
13
14    /// Check if a block is currently allocated.
15    fn is_allocated(&self, block_id: u64) -> bool;
16}
17
18/// Simple bitmap-style allocator backed by a BTreeSet of free block IDs.
19pub struct BitmapAllocator {
20    total_blocks: u64,
21    state: Mutex<AllocatorState>,
22}
23
24struct AllocatorState {
25    /// Set of free block IDs.
26    free_set: BTreeSet<u64>,
27    /// Set of allocated block IDs.
28    allocated: BTreeSet<u64>,
29}
30
31impl BitmapAllocator {
32    /// Create a new allocator for `total_blocks` blocks.
33    /// Blocks 0..FIRST_DATA_BLOCK are reserved and never allocatable.
34    pub fn new(total_blocks: u64) -> Self {
35        let mut free_set = BTreeSet::new();
36        for id in FIRST_DATA_BLOCK..total_blocks {
37            free_set.insert(id);
38        }
39        Self {
40            total_blocks,
41            state: Mutex::new(AllocatorState {
42                free_set,
43                allocated: BTreeSet::new(),
44            }),
45        }
46    }
47
48    /// Mark a block as already allocated (used during mount/recovery).
49    pub fn mark_allocated(&self, block_id: u64) -> FsResult<()> {
50        let mut state = self.state.lock().map_err(|e| FsError::Internal(e.to_string()))?;
51        state.free_set.remove(&block_id);
52        state.allocated.insert(block_id);
53        Ok(())
54    }
55
56    /// Return the number of free blocks.
57    pub fn free_count(&self) -> u64 {
58        let state = self.state.lock().unwrap();
59        state.free_set.len() as u64
60    }
61}
62
63impl SlotAllocator for BitmapAllocator {
64    fn allocate(&self) -> FsResult<u64> {
65        let mut state = self.state.lock().map_err(|e| FsError::Internal(e.to_string()))?;
66        let block_id = *state.free_set.iter().next().ok_or(FsError::NoFreeBlocks)?;
67        state.free_set.remove(&block_id);
68        state.allocated.insert(block_id);
69        Ok(block_id)
70    }
71
72    fn free(&self, block_id: u64) -> FsResult<()> {
73        if block_id >= self.total_blocks || block_id < FIRST_DATA_BLOCK {
74            return Err(FsError::BlockOutOfRange(block_id));
75        }
76        let mut state = self.state.lock().map_err(|e| FsError::Internal(e.to_string()))?;
77        state.allocated.remove(&block_id);
78        state.free_set.insert(block_id);
79        Ok(())
80    }
81
82    fn is_allocated(&self, block_id: u64) -> bool {
83        let state = self.state.lock().unwrap();
84        state.allocated.contains(&block_id)
85    }
86}
87
88#[cfg(test)]
89mod tests {
90    use super::*;
91
92    #[test]
93    fn test_allocate_and_free() {
94        let alloc = BitmapAllocator::new(10);
95        let id = alloc.allocate().unwrap();
96        assert!(id >= FIRST_DATA_BLOCK);
97        assert!(alloc.is_allocated(id));
98        alloc.free(id).unwrap();
99        assert!(!alloc.is_allocated(id));
100    }
101
102    #[test]
103    fn test_exhaustion() {
104        // Only FIRST_DATA_BLOCK..4 are allocatable => 1 block (block 3).
105        let alloc = BitmapAllocator::new(4);
106        let _id = alloc.allocate().unwrap();
107        assert!(alloc.allocate().is_err());
108    }
109}
doublecrypt_core/allocator.rs

doublecrypt_core/
allocator.rs
