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    /// Return all free block IDs.
63    pub fn free_block_ids(&self) -> Vec<u64> {
64        let state = self.state.lock().unwrap();
65        state.free_set.iter().copied().collect()
66    }
67}
68
69impl SlotAllocator for BitmapAllocator {
70    fn allocate(&self) -> FsResult<u64> {
71        let mut state = self.state.lock().map_err(|e| FsError::Internal(e.to_string()))?;
72        let block_id = *state.free_set.iter().next().ok_or(FsError::NoFreeBlocks)?;
73        state.free_set.remove(&block_id);
74        state.allocated.insert(block_id);
75        Ok(block_id)
76    }
77
78    fn free(&self, block_id: u64) -> FsResult<()> {
79        if block_id >= self.total_blocks || block_id < FIRST_DATA_BLOCK {
80            return Err(FsError::BlockOutOfRange(block_id));
81        }
82        let mut state = self.state.lock().map_err(|e| FsError::Internal(e.to_string()))?;
83        state.allocated.remove(&block_id);
84        state.free_set.insert(block_id);
85        Ok(())
86    }
87
88    fn is_allocated(&self, block_id: u64) -> bool {
89        let state = self.state.lock().unwrap();
90        state.allocated.contains(&block_id)
91    }
92}
93
94#[cfg(test)]
95mod tests {
96    use super::*;
97
98    #[test]
99    fn test_allocate_and_free() {
100        let alloc = BitmapAllocator::new(10);
101        let id = alloc.allocate().unwrap();
102        assert!(id >= FIRST_DATA_BLOCK);
103        assert!(alloc.is_allocated(id));
104        alloc.free(id).unwrap();
105        assert!(!alloc.is_allocated(id));
106    }
107
108    #[test]
109    fn test_exhaustion() {
110        // Only FIRST_DATA_BLOCK..4 are allocatable => 1 block (block 3).
111        let alloc = BitmapAllocator::new(4);
112        let _id = alloc.allocate().unwrap();
113        assert!(alloc.allocate().is_err());
114    }
115}
doublecrypt_core/allocator.rs

doublecrypt_core/
allocator.rs
