Skip to main content

sklears_datasets/
memory_pool.rs

1//! Memory pool allocation for high-performance dataset generation
2//!
3//! This module provides sophisticated memory pool management for efficient
4//! allocation and reuse of memory blocks in dataset generation operations.
5//! It reduces allocation overhead and improves cache locality.
6
7use scirs2_core::ndarray::{Array1, Array2, ArrayViewMut1, ArrayViewMut2};
8use std::alloc::{alloc, dealloc, Layout};
9use std::collections::{HashMap, VecDeque};
10use std::ptr::NonNull;
11use std::sync::{
12    atomic::{AtomicUsize, Ordering},
13    Arc, Mutex, RwLock,
14};
15use thiserror::Error;
16
17/// Global monotonic counter used to assign unique pool IDs.
18static NEXT_POOL_ID: AtomicUsize = AtomicUsize::new(1);
19
20/// Memory pool errors
21#[derive(Error, Debug)]
22pub enum MemoryPoolError {
23    #[error("Pool exhausted: no available blocks of size {0}")]
24    PoolExhausted(usize),
25    #[error("Invalid block size: {0}")]
26    InvalidBlockSize(usize),
27    #[error("Allocation failed: {0}")]
28    AllocationFailed(String),
29    #[error("Block not found: {0:?}")]
30    BlockNotFound(usize),
31    #[error("Double free detected: {0:?}")]
32    DoubleFree(usize),
33    #[error("Pool corruption detected: {0}")]
34    PoolCorruption(String),
35    #[error("Alignment error: required {required}, got {actual}")]
36    AlignmentError { required: usize, actual: usize },
37}
38
39pub type MemoryPoolResult<T> = Result<T, MemoryPoolError>;
40
41/// Memory block metadata
42#[derive(Debug, Clone)]
43struct MemoryBlock {
44    ptr: NonNull<u8>,
45    size: usize,
46    layout: Layout,
47    allocated_at: std::time::Instant,
48    #[allow(dead_code)] // pool_id reserved for multi-pool debugging and diagnostics
49    pool_id: usize,
50    magic: u64, // For corruption detection
51}
52
53const MAGIC_VALUE: u64 = 0xDEADBEEFCAFEBABE;
54
55impl MemoryBlock {
56    fn new(size: usize, align: usize, pool_id: usize) -> MemoryPoolResult<Self> {
57        let layout = Layout::from_size_align(size, align)
58            .map_err(|e| MemoryPoolError::AllocationFailed(e.to_string()))?;
59
60        let ptr = unsafe { alloc(layout) };
61        if ptr.is_null() {
62            return Err(MemoryPoolError::AllocationFailed(
63                "System allocation failed".to_string(),
64            ));
65        }
66
67        Ok(Self {
68            ptr: NonNull::new(ptr).ok_or_else(|| {
69                MemoryPoolError::AllocationFailed("NonNull creation failed".to_string())
70            })?,
71            size,
72            layout,
73            allocated_at: std::time::Instant::now(),
74            pool_id,
75            magic: MAGIC_VALUE,
76        })
77    }
78
79    fn is_valid(&self) -> bool {
80        self.magic == MAGIC_VALUE
81    }
82
83    fn invalidate(&mut self) {
84        self.magic = 0;
85    }
86
87    fn as_slice_mut(&mut self) -> &mut [u8] {
88        unsafe { std::slice::from_raw_parts_mut(self.ptr.as_ptr(), self.size) }
89    }
90
91    #[allow(dead_code)] // available to callers but not yet used in current test paths
92    fn as_f64_slice_mut(&mut self) -> MemoryPoolResult<&mut [f64]> {
93        if !self.size.is_multiple_of(std::mem::size_of::<f64>()) {
94            return Err(MemoryPoolError::InvalidBlockSize(self.size));
95        }
96
97        let len = self.size / std::mem::size_of::<f64>();
98        Ok(unsafe { std::slice::from_raw_parts_mut(self.ptr.as_ptr() as *mut f64, len) })
99    }
100}
101
102impl Drop for MemoryBlock {
103    fn drop(&mut self) {
104        if self.is_valid() {
105            unsafe {
106                dealloc(self.ptr.as_ptr(), self.layout);
107            }
108            self.invalidate();
109        }
110    }
111}
112
113// Safety: MemoryBlock owns the memory pointed to by ptr (like Box<[u8]>).
114// The pointer represents uniquely owned memory, and all access is protected
115// by RwLock synchronization in the pool. There is no concurrent unsynchronized
116// access to the underlying memory.
117unsafe impl Send for MemoryBlock {}
118unsafe impl Sync for MemoryBlock {}
119
120/// Memory pool configuration
121#[derive(Debug, Clone)]
122pub struct MemoryPoolConfig {
123    /// Initial number of blocks per size bucket
124    pub initial_blocks_per_size: usize,
125    /// Maximum number of blocks per size bucket
126    pub max_blocks_per_size: usize,
127    /// Growth factor when expanding pool
128    pub growth_factor: f64,
129    /// Alignment requirement for blocks
130    pub alignment: usize,
131    /// Enable block reuse
132    pub enable_reuse: bool,
133    /// Maximum idle time before releasing blocks
134    pub max_idle_time: std::time::Duration,
135    /// Enable statistics collection
136    pub enable_stats: bool,
137}
138
139impl Default for MemoryPoolConfig {
140    fn default() -> Self {
141        Self {
142            initial_blocks_per_size: 4,
143            max_blocks_per_size: 64,
144            growth_factor: 1.5,
145            alignment: 64, // Cache line alignment
146            enable_reuse: true,
147            max_idle_time: std::time::Duration::from_secs(300), // 5 minutes
148            enable_stats: true,
149        }
150    }
151}
152
153/// Memory pool statistics
154#[derive(Debug, Clone, Default)]
155pub struct MemoryPoolStats {
156    pub total_allocations: usize,
157    pub total_deallocations: usize,
158    pub current_allocations: usize,
159    pub peak_allocations: usize,
160    pub total_bytes_allocated: usize,
161    pub current_bytes_allocated: usize,
162    pub peak_bytes_allocated: usize,
163    pub pool_hits: usize,
164    pub pool_misses: usize,
165    pub blocks_created: usize,
166    pub blocks_reused: usize,
167    pub blocks_released: usize,
168}
169
170impl MemoryPoolStats {
171    pub fn hit_rate(&self) -> f64 {
172        if self.pool_hits + self.pool_misses == 0 {
173            0.0
174        } else {
175            self.pool_hits as f64 / (self.pool_hits + self.pool_misses) as f64
176        }
177    }
178
179    pub fn reuse_rate(&self) -> f64 {
180        if self.blocks_created == 0 {
181            0.0
182        } else {
183            self.blocks_reused as f64 / self.blocks_created as f64
184        }
185    }
186}
187
188/// Size bucket for grouping similar-sized allocations
189struct SizeBucket {
190    size: usize,
191    available: VecDeque<MemoryBlock>,
192    allocated: HashMap<usize, MemoryBlock>,
193    total_created: usize,
194}
195
196impl SizeBucket {
197    fn new(size: usize) -> Self {
198        Self {
199            size,
200            available: VecDeque::new(),
201            allocated: HashMap::new(),
202            total_created: 0,
203        }
204    }
205
206    fn allocate(
207        &mut self,
208        pool_id: usize,
209        config: &MemoryPoolConfig,
210    ) -> MemoryPoolResult<(usize, *mut u8, bool)> {
211        let (block, was_reused) = if let Some(block) = self.available.pop_front() {
212            (block, true)
213        } else {
214            // Create new block
215            let block = MemoryBlock::new(self.size, config.alignment, pool_id)?;
216            self.total_created += 1;
217            (block, false)
218        };
219
220        let ptr = block.ptr.as_ptr();
221        let block_id = ptr as usize;
222        self.allocated.insert(block_id, block);
223
224        Ok((block_id, ptr, was_reused))
225    }
226
227    fn deallocate(&mut self, block_id: usize, enable_reuse: bool) -> MemoryPoolResult<()> {
228        if let Some(mut block) = self.allocated.remove(&block_id) {
229            if !block.is_valid() {
230                return Err(MemoryPoolError::PoolCorruption(
231                    "Block magic value corrupted".to_string(),
232                ));
233            }
234
235            if enable_reuse {
236                // Zero out the memory for security
237                let slice = block.as_slice_mut();
238                slice.fill(0);
239
240                self.available.push_back(block);
241            } else {
242                // Block will be dropped here and its memory freed by
243                // `MemoryBlock`'s `Drop` impl. Do NOT call `block.invalidate()`
244                // beforehand: `Drop::drop` only calls `dealloc()` when
245                // `is_valid()` is true, so pre-invalidating here would make
246                // the drop silently skip deallocation, leaking the block on
247                // every non-reuse deallocation (confirmed via a Miri leak
248                // report from `test_pool_exhaustion`).
249            }
250
251            Ok(())
252        } else {
253            Err(MemoryPoolError::BlockNotFound(block_id))
254        }
255    }
256
257    fn cleanup_idle(&mut self, max_idle_time: std::time::Duration) -> usize {
258        let now = std::time::Instant::now();
259        let mut removed = 0;
260
261        self.available.retain(|block| {
262            if now.duration_since(block.allocated_at) > max_idle_time {
263                removed += 1;
264                false
265            } else {
266                true
267            }
268        });
269
270        removed
271    }
272}
273
274/// High-performance memory pool for dataset generation
275pub struct MemoryPool {
276    config: MemoryPoolConfig,
277    buckets: RwLock<HashMap<usize, SizeBucket>>,
278    stats: RwLock<MemoryPoolStats>,
279    pool_id: usize,
280    next_cleanup: Mutex<std::time::Instant>,
281}
282
283impl MemoryPool {
284    /// Create a new memory pool
285    pub fn new(config: MemoryPoolConfig) -> Self {
286        let max_idle_time = config.max_idle_time;
287        Self {
288            config,
289            buckets: RwLock::new(HashMap::new()),
290            stats: RwLock::new(MemoryPoolStats::default()),
291            pool_id: NEXT_POOL_ID.fetch_add(1, Ordering::Relaxed),
292            next_cleanup: Mutex::new(std::time::Instant::now() + max_idle_time),
293        }
294    }
295
296    /// Return this pool's unique identifier.
297    pub fn pool_id(&self) -> usize {
298        self.pool_id
299    }
300
301    /// Allocate a block of memory
302    pub fn allocate(&self, size: usize) -> MemoryPoolResult<PooledBlock> {
303        if size == 0 {
304            return Err(MemoryPoolError::InvalidBlockSize(size));
305        }
306
307        // Round up to nearest power of 2 for better bucketing
308        let bucket_size = size.next_power_of_two().max(64);
309
310        let (block_id, ptr, was_reused) = {
311            let mut buckets = self
312                .buckets
313                .write()
314                .unwrap_or_else(|poisoned| poisoned.into_inner());
315            let bucket = buckets
316                .entry(bucket_size)
317                .or_insert_with(|| SizeBucket::new(bucket_size));
318
319            // Check if we have room for more allocations
320            if bucket.allocated.len() >= self.config.max_blocks_per_size {
321                return Err(MemoryPoolError::PoolExhausted(bucket_size));
322            }
323
324            bucket.allocate(self.pool_id, &self.config)?
325        };
326
327        // Update statistics
328        if self.config.enable_stats {
329            let mut stats = self
330                .stats
331                .write()
332                .unwrap_or_else(|poisoned| poisoned.into_inner());
333            stats.total_allocations += 1;
334            stats.current_allocations += 1;
335            stats.peak_allocations = stats.peak_allocations.max(stats.current_allocations);
336            stats.total_bytes_allocated += bucket_size;
337            stats.current_bytes_allocated += bucket_size;
338            stats.peak_bytes_allocated = stats
339                .peak_bytes_allocated
340                .max(stats.current_bytes_allocated);
341
342            if was_reused {
343                stats.blocks_reused += 1;
344            } else {
345                stats.blocks_created += 1;
346            }
347
348            if bucket_size != size {
349                stats.pool_hits += 1;
350            } else {
351                stats.pool_misses += 1;
352            }
353        }
354
355        // Periodic cleanup
356        self.maybe_cleanup();
357
358        Ok(PooledBlock {
359            ptr,
360            size: bucket_size,
361            block_id,
362            pool: self as *const Self,
363        })
364    }
365
366    /// Deallocate a block of memory
367    pub fn deallocate(&self, block_id: usize, size: usize) -> MemoryPoolResult<()> {
368        let bucket_size = size.next_power_of_two().max(64);
369
370        {
371            let mut buckets = self
372                .buckets
373                .write()
374                .unwrap_or_else(|poisoned| poisoned.into_inner());
375            if let Some(bucket) = buckets.get_mut(&bucket_size) {
376                bucket.deallocate(block_id, self.config.enable_reuse)?;
377            } else {
378                return Err(MemoryPoolError::BlockNotFound(block_id));
379            }
380        }
381
382        // Update statistics
383        if self.config.enable_stats {
384            let mut stats = self
385                .stats
386                .write()
387                .unwrap_or_else(|poisoned| poisoned.into_inner());
388            stats.total_deallocations += 1;
389            stats.current_allocations = stats.current_allocations.saturating_sub(1);
390            stats.current_bytes_allocated =
391                stats.current_bytes_allocated.saturating_sub(bucket_size);
392        }
393
394        Ok(())
395    }
396
397    /// Get pool statistics
398    pub fn stats(&self) -> MemoryPoolStats {
399        if self.config.enable_stats {
400            self.stats
401                .read()
402                .unwrap_or_else(|poisoned| poisoned.into_inner())
403                .clone()
404        } else {
405            MemoryPoolStats::default()
406        }
407    }
408
409    /// Force cleanup of idle blocks
410    pub fn cleanup(&self) -> usize {
411        let mut total_removed = 0;
412        let mut buckets = self
413            .buckets
414            .write()
415            .unwrap_or_else(|poisoned| poisoned.into_inner());
416
417        for bucket in buckets.values_mut() {
418            total_removed += bucket.cleanup_idle(self.config.max_idle_time);
419        }
420
421        // Update next cleanup time
422        *self
423            .next_cleanup
424            .lock()
425            .unwrap_or_else(|poisoned| poisoned.into_inner()) =
426            std::time::Instant::now() + self.config.max_idle_time;
427
428        total_removed
429    }
430
431    /// Maybe perform cleanup if enough time has passed
432    fn maybe_cleanup(&self) {
433        let now = std::time::Instant::now();
434        let should_cleanup = {
435            let next_cleanup = self
436                .next_cleanup
437                .lock()
438                .unwrap_or_else(|poisoned| poisoned.into_inner());
439            now >= *next_cleanup
440        };
441
442        if should_cleanup {
443            self.cleanup();
444        }
445    }
446
447    /// Allocate memory for a 2D array
448    pub fn allocate_array2(&self, nrows: usize, ncols: usize) -> MemoryPoolResult<PooledArray2> {
449        let size = nrows * ncols * std::mem::size_of::<f64>();
450        let block = self.allocate(size)?;
451
452        Ok(PooledArray2 {
453            block,
454            nrows,
455            ncols,
456        })
457    }
458
459    /// Allocate memory for a 1D array
460    pub fn allocate_array1(&self, len: usize) -> MemoryPoolResult<PooledArray1> {
461        let size = len * std::mem::size_of::<f64>();
462        let block = self.allocate(size)?;
463
464        Ok(PooledArray1 { block, len })
465    }
466
467    /// Get bucket information for debugging
468    pub fn bucket_info(&self) -> Vec<(usize, usize, usize)> {
469        let buckets = self
470            .buckets
471            .read()
472            .unwrap_or_else(|poisoned| poisoned.into_inner());
473        buckets
474            .iter()
475            .map(|(size, bucket)| (*size, bucket.available.len(), bucket.allocated.len()))
476            .collect()
477    }
478}
479
480impl Default for MemoryPool {
481    fn default() -> Self {
482        Self::new(MemoryPoolConfig::default())
483    }
484}
485
486/// RAII wrapper for pooled memory blocks
487pub struct PooledBlock {
488    ptr: *mut u8,
489    size: usize,
490    block_id: usize,
491    pool: *const MemoryPool,
492}
493
494impl PooledBlock {
495    /// Get a mutable slice view of the block
496    pub fn as_slice_mut(&mut self) -> &mut [u8] {
497        unsafe { std::slice::from_raw_parts_mut(self.ptr, self.size) }
498    }
499
500    /// Get the block size
501    pub fn size(&self) -> usize {
502        self.size
503    }
504
505    /// Get a mutable slice as f64 values
506    pub fn as_f64_slice_mut(&mut self) -> MemoryPoolResult<&mut [f64]> {
507        if !self.size.is_multiple_of(std::mem::size_of::<f64>()) {
508            return Err(MemoryPoolError::InvalidBlockSize(self.size));
509        }
510
511        let len = self.size / std::mem::size_of::<f64>();
512        Ok(unsafe { std::slice::from_raw_parts_mut(self.ptr as *mut f64, len) })
513    }
514}
515
516impl Drop for PooledBlock {
517    fn drop(&mut self) {
518        if !self.pool.is_null() {
519            unsafe {
520                if let Err(e) = (*self.pool).deallocate(self.block_id, self.size) {
521                    eprintln!("Error deallocating pooled block: {}", e);
522                }
523            }
524        }
525    }
526}
527
528unsafe impl Send for PooledBlock {}
529unsafe impl Sync for PooledBlock {}
530
531/// Pooled 2D array
532pub struct PooledArray2 {
533    block: PooledBlock,
534    nrows: usize,
535    ncols: usize,
536}
537
538impl PooledArray2 {
539    /// Get dimensions
540    pub fn dim(&self) -> (usize, usize) {
541        (self.nrows, self.ncols)
542    }
543
544    /// Get mutable array view
545    pub fn view_mut(&mut self) -> MemoryPoolResult<ArrayViewMut2<'_, f64>> {
546        let slice = self.block.as_f64_slice_mut()?;
547        ArrayViewMut2::from_shape((self.nrows, self.ncols), slice)
548            .map_err(|e| MemoryPoolError::AllocationFailed(format!("Shape error: {}", e)))
549    }
550
551    /// Convert to owned Array2
552    pub fn to_array(mut self) -> MemoryPoolResult<Array2<f64>> {
553        let slice = self.block.as_f64_slice_mut()?;
554        let total_elements = self.nrows * self.ncols;
555        // Only use the requested number of elements, not the full allocated block
556        let data = slice[..total_elements].to_vec();
557        let array = Array2::from_shape_vec((self.nrows, self.ncols), data).map_err(|e| {
558            MemoryPoolError::AllocationFailed(format!("Array creation error: {}", e))
559        })?;
560        Ok(array)
561    }
562}
563
564/// Pooled 1D array
565pub struct PooledArray1 {
566    block: PooledBlock,
567    len: usize,
568}
569
570impl PooledArray1 {
571    /// Get length
572    pub fn len(&self) -> usize {
573        self.len
574    }
575
576    /// Check if empty
577    pub fn is_empty(&self) -> bool {
578        self.len == 0
579    }
580
581    /// Get mutable array view
582    pub fn view_mut(&mut self) -> MemoryPoolResult<ArrayViewMut1<'_, f64>> {
583        let slice = self.block.as_f64_slice_mut()?;
584        ArrayViewMut1::from_shape(self.len, slice)
585            .map_err(|e| MemoryPoolError::AllocationFailed(format!("Shape error: {}", e)))
586    }
587
588    /// Convert to owned Array1
589    pub fn to_array(mut self) -> MemoryPoolResult<Array1<f64>> {
590        let slice = self.block.as_f64_slice_mut()?;
591        // Only use the requested number of elements, not the full allocated block
592        let array = Array1::from_vec(slice[..self.len].to_vec());
593        Ok(array)
594    }
595}
596
597/// Thread-safe memory pool for concurrent access
598pub type SharedMemoryPool = Arc<MemoryPool>;
599
600/// Create a shared memory pool
601pub fn create_shared_pool(config: MemoryPoolConfig) -> SharedMemoryPool {
602    Arc::new(MemoryPool::new(config))
603}
604
605lazy_static::lazy_static! {
606    /// Global memory pool instance
607    static ref GLOBAL_POOL: SharedMemoryPool = create_shared_pool(MemoryPoolConfig::default());
608}
609
610/// Get the global memory pool
611pub fn global_pool() -> &'static SharedMemoryPool {
612    &GLOBAL_POOL
613}
614
615/// Convenience function to allocate from global pool
616pub fn allocate_global(size: usize) -> MemoryPoolResult<PooledBlock> {
617    global_pool().allocate(size)
618}
619
620/// Convenience function to allocate 2D array from global pool
621pub fn allocate_array2_global(nrows: usize, ncols: usize) -> MemoryPoolResult<PooledArray2> {
622    global_pool().allocate_array2(nrows, ncols)
623}
624
625/// Convenience function to allocate 1D array from global pool
626pub fn allocate_array1_global(len: usize) -> MemoryPoolResult<PooledArray1> {
627    global_pool().allocate_array1(len)
628}
629
630#[allow(non_snake_case)]
631#[cfg(test)]
632mod tests {
633    use super::*;
634
635    #[test]
636    fn test_memory_pool_basic() -> MemoryPoolResult<()> {
637        let pool = MemoryPool::new(MemoryPoolConfig::default());
638
639        // Allocate a block
640        let mut block = pool.allocate(1024)?;
641        assert_eq!(block.size(), 1024);
642
643        // Write some data
644        let slice = block.as_slice_mut();
645        slice[0] = 42;
646        slice[1023] = 24;
647
648        // Check data
649        assert_eq!(slice[0], 42);
650        assert_eq!(slice[1023], 24);
651
652        Ok(())
653    }
654
655    #[test]
656    fn test_memory_pool_reuse() -> MemoryPoolResult<()> {
657        let config = MemoryPoolConfig {
658            enable_reuse: true,
659            ..Default::default()
660        };
661        let pool = MemoryPool::new(config);
662
663        // Allocate and deallocate
664        {
665            let _block1 = pool.allocate(512)?;
666        } // Block is returned to pool
667
668        // Allocate again - should reuse
669        let _block2 = pool.allocate(512)?;
670
671        let stats = pool.stats();
672        assert!(stats.blocks_reused > 0 || stats.pool_hits > 0);
673
674        Ok(())
675    }
676
677    #[test]
678    fn test_pooled_array2() -> MemoryPoolResult<()> {
679        let pool = MemoryPool::new(MemoryPoolConfig::default());
680
681        let mut array = pool.allocate_array2(100, 50)?;
682        assert_eq!(array.dim(), (100, 50));
683
684        // Get mutable view and fill with data
685        {
686            let mut view = array.view_mut()?;
687            view.fill(std::f64::consts::PI);
688        }
689
690        // Convert to owned array
691        let owned = array.to_array()?;
692        assert_eq!(owned.dim(), (100, 50));
693        assert!((owned[[0, 0]] - std::f64::consts::PI).abs() < f64::EPSILON);
694
695        Ok(())
696    }
697
698    #[test]
699    fn test_pooled_array1() -> MemoryPoolResult<()> {
700        let pool = MemoryPool::new(MemoryPoolConfig::default());
701
702        let mut array = pool.allocate_array1(1000)?;
703        assert_eq!(array.len(), 1000);
704
705        // Get mutable view and fill with data
706        {
707            let mut view = array.view_mut()?;
708            view.fill(2.71);
709        }
710
711        // Convert to owned array
712        let owned = array.to_array()?;
713        assert_eq!(owned.len(), 1000);
714        assert!((owned[0] - 2.71).abs() < f64::EPSILON);
715
716        Ok(())
717    }
718
719    #[test]
720    fn test_memory_pool_stats() -> MemoryPoolResult<()> {
721        let config = MemoryPoolConfig {
722            enable_stats: true,
723            ..Default::default()
724        };
725        let pool = MemoryPool::new(config);
726
727        // Initial stats
728        let initial_stats = pool.stats();
729        assert_eq!(initial_stats.total_allocations, 0);
730
731        // Allocate some blocks
732        let _block1 = pool.allocate(256)?;
733        let _block2 = pool.allocate(512)?;
734
735        let stats = pool.stats();
736        assert_eq!(stats.total_allocations, 2);
737        assert_eq!(stats.current_allocations, 2);
738        assert!(stats.current_bytes_allocated > 0);
739
740        Ok(())
741    }
742
743    #[test]
744    fn test_global_pool() -> MemoryPoolResult<()> {
745        let mut block = allocate_global(128)?;
746        assert_eq!(block.size(), 128);
747
748        let slice = block.as_slice_mut();
749        slice[0] = 255;
750        assert_eq!(slice[0], 255);
751
752        let array2 = allocate_array2_global(10, 20)?;
753        assert_eq!(array2.dim(), (10, 20));
754
755        let array1 = allocate_array1_global(100)?;
756        assert_eq!(array1.len(), 100);
757
758        Ok(())
759    }
760
761    #[test]
762    fn test_pool_exhaustion() {
763        let config = MemoryPoolConfig {
764            max_blocks_per_size: 2,
765            enable_reuse: false,
766            ..Default::default()
767        };
768        let pool = MemoryPool::new(config);
769
770        // Allocate up to the limit
771        let _block1 = pool.allocate(64).expect("operation should succeed");
772        let _block2 = pool.allocate(64).expect("operation should succeed");
773
774        // This should fail
775        assert!(pool.allocate(64).is_err());
776    }
777
778    #[test]
779    fn test_cleanup() -> MemoryPoolResult<()> {
780        let config = MemoryPoolConfig {
781            max_idle_time: std::time::Duration::from_millis(1),
782            enable_reuse: true,
783            ..Default::default()
784        };
785        let pool = MemoryPool::new(config);
786
787        // Allocate and immediately deallocate
788        {
789            let _block = pool.allocate(256)?;
790        }
791
792        // Wait for idle time to pass
793        std::thread::sleep(std::time::Duration::from_millis(10));
794
795        // Force cleanup
796        let removed = pool.cleanup();
797        assert!(removed > 0);
798
799        Ok(())
800    }
801
802    #[test]
803    fn test_pool_ids_are_unique_and_nonzero() {
804        // Each new pool should receive a distinct, nonzero identifier.
805        let p1 = MemoryPool::new(MemoryPoolConfig::default());
806        let p2 = MemoryPool::new(MemoryPoolConfig::default());
807        let p3 = MemoryPool::new(MemoryPoolConfig::default());
808
809        assert_ne!(p1.pool_id(), 0, "pool_id must not be zero");
810        assert_ne!(p1.pool_id(), p2.pool_id(), "pool IDs must be unique");
811        assert_ne!(p2.pool_id(), p3.pool_id(), "pool IDs must be unique");
812        assert_ne!(p1.pool_id(), p3.pool_id(), "pool IDs must be unique");
813    }
814
815    #[test]
816    fn test_bucket_alignment() -> MemoryPoolResult<()> {
817        let pool = MemoryPool::new(MemoryPoolConfig::default());
818
819        // Different sizes should be bucketed to powers of 2
820        let block1 = pool.allocate(100)?; // Should be rounded to 128
821        let block2 = pool.allocate(200)?; // Should be rounded to 256
822
823        assert!(block1.size() >= 100);
824        assert!(block2.size() >= 200);
825
826        // Size should be power of 2
827        assert!(block1.size().is_power_of_two());
828        assert!(block2.size().is_power_of_two());
829
830        Ok(())
831    }
832}