Skip to main content

feoxdb/storage/
free_space.rs

1use crate::constants::*;
2use crate::error::{FeoxError, Result};
3use std::collections::BTreeMap;
4
5/// Represents a contiguous range of free sectors on disk
6#[derive(Debug, Clone, PartialEq, Eq)]
7pub struct FreeSpace {
8    pub start: u64, // Starting sector
9    pub size: u64,  // Size in sectors
10}
11
12/// Free space manager using dual RB-trees for efficient allocation and coalescing
13/// Matches the kernel implementation's design for correctness
14pub struct FreeSpaceManager {
15    /// Tree sorted by (size, start) for best-fit allocation
16    /// Using composite key ensures uniqueness
17    by_size: BTreeMap<(u64, u64), FreeSpace>,
18
19    /// Tree sorted by start address for efficient merging
20    by_start: BTreeMap<u64, FreeSpace>,
21
22    /// Total free space in bytes
23    total_free: u64,
24
25    /// Device size in bytes (for validation)
26    device_size: u64,
27
28    /// Fragmentation percentage (0-100)
29    fragmentation_percent: u32,
30}
31
32impl FreeSpaceManager {
33    /// Create a new free space manager
34    pub fn new() -> Self {
35        Self {
36            by_size: BTreeMap::new(),
37            by_start: BTreeMap::new(),
38            total_free: 0,
39            device_size: 0,
40            fragmentation_percent: 0,
41        }
42    }
43
44    /// Initialize with device size and initial free space
45    pub fn initialize(&mut self, device_size: u64) -> Result<()> {
46        self.device_size = device_size;
47
48        let metadata_sectors = FEOX_DATA_START_BLOCK;
49        let total_sectors = device_size / FEOX_BLOCK_SIZE as u64;
50
51        if total_sectors <= metadata_sectors {
52            return Err(FeoxError::InvalidDevice);
53        }
54
55        // Add all space after metadata as free
56        let free_sectors = total_sectors - metadata_sectors;
57        self.insert_free_space(FreeSpace {
58            start: metadata_sectors,
59            size: free_sectors,
60        })?;
61
62        Ok(())
63    }
64
65    /// Set device size for bounds checking (used when rebuilding from scan)
66    pub fn set_device_size(&mut self, device_size: u64) {
67        self.device_size = device_size;
68    }
69
70    /// Allocate sectors using best-fit algorithm
71    pub fn allocate_sectors(&mut self, sectors_needed: u64) -> Result<u64> {
72        if sectors_needed == 0 {
73            return Err(FeoxError::InvalidArgument);
74        }
75
76        let best_fit = self
77            .by_size
78            .range((sectors_needed, 0)..)
79            .next()
80            .map(|((size, start), space)| (*size, *start, space.clone()));
81
82        if let Some((size, start, space)) = best_fit {
83            // Validate the space
84            if !self.is_valid_free_space(&space) {
85                return Err(FeoxError::CorruptedData);
86            }
87
88            // Remove from both trees
89            self.by_size.remove(&(size, start));
90            self.by_start.remove(&space.start);
91            self.total_free -= space.size * FEOX_BLOCK_SIZE as u64;
92
93            let allocated_start = space.start;
94
95            // Handle remaining space if any
96            if space.size > sectors_needed {
97                let remaining = FreeSpace {
98                    start: space.start + sectors_needed,
99                    size: space.size - sectors_needed,
100                };
101
102                // Insert remaining space back
103                if let Err(e) = self.insert_free_space(remaining) {
104                    // Try to restore original space on error
105                    let _ = self.insert_free_space(space);
106                    return Err(e);
107                }
108            }
109
110            self.update_fragmentation();
111
112            Ok(allocated_start)
113        } else {
114            Err(FeoxError::OutOfSpace)
115        }
116    }
117
118    /// Release sectors back to free space pool with coalescing
119    pub fn release_sectors(&mut self, start: u64, count: u64) -> Result<()> {
120        if start < FEOX_DATA_START_BLOCK || count == 0 {
121            return Err(FeoxError::InvalidArgument);
122        }
123
124        // Validate bounds
125        if !self.is_valid_sector_range(start, count) {
126            return Err(FeoxError::InvalidArgument);
127        }
128
129        // Try to merge with adjacent spaces
130        let merged = self.try_merge_spaces(start, count)?;
131
132        // Insert the merged space (this will update total_free)
133        self.insert_free_space(merged)?;
134
135        self.update_fragmentation();
136
137        Ok(())
138    }
139
140    /// Try to merge with adjacent free spaces
141    fn try_merge_spaces(&mut self, start: u64, size: u64) -> Result<FreeSpace> {
142        let end = start.checked_add(size).ok_or(FeoxError::InvalidArgument)?;
143
144        // Overlap is rejected from these two probes, before either neighbour is
145        // removed or debited, so a rejected release leaves the free list untouched.
146        // Checking after the merge would already have destroyed the neighbours.
147        let preceding = self
148            .by_start
149            .range(..start)
150            .next_back()
151            .map(|(_, space)| space.clone());
152        if let Some(ref space) = preceding {
153            if space.start + space.size > start {
154                return Err(FeoxError::DuplicateKey);
155            }
156        }
157
158        let following = self
159            .by_start
160            .range(start..=end)
161            .next()
162            .map(|(_, space)| space.clone());
163        if let Some(ref space) = following {
164            if space.start < end {
165                return Err(FeoxError::DuplicateKey);
166            }
167        }
168
169        let mut merged_start = start;
170        let mut merged_size = size;
171
172        let prev = preceding.filter(|space| space.start + space.size == start);
173        let next = following.filter(|space| space.start == end);
174
175        // Merge with predecessor if found
176        if let Some(prev_space) = prev {
177            // Remove from both trees
178            self.by_size.remove(&(prev_space.size, prev_space.start));
179            self.by_start.remove(&prev_space.start);
180
181            // Subtract the removed space from total_free (will be re-added when inserting merged)
182            self.total_free -= prev_space.size * FEOX_BLOCK_SIZE as u64;
183
184            merged_start = prev_space.start;
185            merged_size += prev_space.size;
186        }
187
188        // Merge with successor if found
189        if let Some(next_space) = next {
190            // Remove from both trees
191            self.by_size.remove(&(next_space.size, next_space.start));
192            self.by_start.remove(&next_space.start);
193
194            // Subtract the removed space from total_free (will be re-added when inserting merged)
195            self.total_free -= next_space.size * FEOX_BLOCK_SIZE as u64;
196
197            merged_size += next_space.size;
198        }
199
200        Ok(FreeSpace {
201            start: merged_start,
202            size: merged_size,
203        })
204    }
205
206    /// Insert a free space into both trees
207    fn insert_free_space(&mut self, space: FreeSpace) -> Result<()> {
208        if space.size == 0 {
209            return Err(FeoxError::InvalidArgument);
210        }
211
212        // Validate the space
213        if !self.is_valid_free_space(&space) {
214            return Err(FeoxError::InvalidArgument);
215        }
216
217        // Check for duplicates
218        if self.by_start.contains_key(&space.start) {
219            return Err(FeoxError::DuplicateKey);
220        }
221
222        // Update total free space
223        self.total_free += space.size * FEOX_BLOCK_SIZE as u64;
224
225        // Insert into both trees
226        self.by_size
227            .insert((space.size, space.start), space.clone());
228        self.by_start.insert(space.start, space);
229
230        Ok(())
231    }
232
233    /// Check if a free space is valid
234    fn is_valid_free_space(&self, space: &FreeSpace) -> bool {
235        if space.start < FEOX_DATA_START_BLOCK {
236            return false;
237        }
238
239        // Check device bounds
240        if self.device_size > 0 {
241            let device_sectors = self.device_size / FEOX_BLOCK_SIZE as u64;
242            if space.start >= device_sectors {
243                return false;
244            }
245            if space.start.checked_add(space.size).is_none() {
246                return false;
247            }
248            if space.start + space.size > device_sectors {
249                return false;
250            }
251        }
252
253        true
254    }
255
256    /// Check if a sector range is valid
257    fn is_valid_sector_range(&self, start: u64, count: u64) -> bool {
258        if start < FEOX_DATA_START_BLOCK || count == 0 {
259            return false;
260        }
261
262        if self.device_size > 0 {
263            let device_sectors = self.device_size / FEOX_BLOCK_SIZE as u64;
264            if start >= device_sectors {
265                return false;
266            }
267            match start.checked_add(count) {
268                Some(end) if end <= device_sectors => {}
269                _ => return false,
270            }
271        }
272
273        true
274    }
275
276    /// Update fragmentation metric
277    fn update_fragmentation(&mut self) {
278        if self.total_free == 0 {
279            self.fragmentation_percent = 0;
280            return;
281        }
282
283        let num_chunks = self.by_start.len() as u64;
284        if num_chunks <= 1 {
285            self.fragmentation_percent = 0;
286            return;
287        }
288
289        // Find largest free chunk
290        let largest = self
291            .by_size
292            .iter()
293            .next_back()
294            .map(|(_, space)| space.size * FEOX_BLOCK_SIZE as u64)
295            .unwrap_or(0);
296
297        // Calculate fragmentation as percentage of free space not in largest chunk
298        if self.total_free > 0 {
299            let fragmented = self.total_free - largest;
300            self.fragmentation_percent = ((fragmented * 100) / self.total_free) as u32;
301        }
302    }
303
304    /// Get total free space in bytes
305    pub fn get_total_free(&self) -> u64 {
306        self.total_free
307    }
308
309    /// Get fragmentation percentage
310    pub fn get_fragmentation(&self) -> u32 {
311        self.fragmentation_percent
312    }
313
314    /// Get number of free chunks
315    pub fn get_free_chunks_count(&self) -> usize {
316        self.by_start.len()
317    }
318
319    /// Get largest free chunk in bytes
320    pub fn get_largest_free_chunk(&self) -> u64 {
321        self.by_size
322            .iter()
323            .next_back()
324            .map(|(_, space)| space.size * FEOX_BLOCK_SIZE as u64)
325            .unwrap_or(0)
326    }
327}
328
329impl Default for FreeSpaceManager {
330    fn default() -> Self {
331        Self::new()
332    }
333}