1use scirs2_core::ndarray::{Array2, ArrayViewMut1, ArrayViewMut2};
33use std::alloc::{GlobalAlloc, Layout, System};
34use std::collections::VecDeque;
35use std::ptr::NonNull;
36use std::sync::Mutex;
37
38#[cfg(any(target_os = "linux", target_os = "android"))]
40use libc;
41#[cfg(target_os = "linux")]
42use std::fs;
43
44use std::sync::atomic::Ordering;
46
47#[cfg(test)]
50use num_cpus;
51
52#[cfg(not(test))]
54mod num_cpus {
55 pub fn get() -> usize {
56 std::thread::available_parallelism()
57 .map(|n| n.get())
58 .unwrap_or(4)
59 }
60}
61
62#[derive(Debug, Clone)]
64pub struct MemoryPoolConfig {
65 pub max_pool_size: usize,
67 pub cache_line_size: usize,
69 pub numa_aware: bool,
71 pub prefetch_distance: usize,
73 pub arena_block_size: usize,
75 pub numa_node_hint: i32,
77 pub auto_numa_discovery: bool,
79 pub enable_thread_affinity: bool,
81 pub enable_memory_warming: bool,
83 pub large_object_threshold: usize,
85 pub max_memory_usage: usize,
87}
88
89impl Default for MemoryPoolConfig {
90 fn default() -> Self {
91 Self {
92 max_pool_size: 1000,
93 cache_line_size: 64,
94 numa_aware: true,
95 prefetch_distance: 8,
96 arena_block_size: 1024 * 1024, numa_node_hint: -1, auto_numa_discovery: true,
99 enable_thread_affinity: true,
100 enable_memory_warming: true,
101 large_object_threshold: 64 * 1024, max_memory_usage: 1024 * 1024 * 1024, }
104 }
105}
106
107pub struct DistancePool {
109 config: MemoryPoolConfig,
110 distance_buffers: Mutex<VecDeque<Box<[f64]>>>,
111 index_buffers: Mutex<VecDeque<Box<[usize]>>>,
112 matrix_buffers: Mutex<VecDeque<Array2<f64>>>,
113 large_buffers: Mutex<VecDeque<Box<[f64]>>>, stats: PoolStatistics,
115 memory_usage: std::sync::atomic::AtomicUsize, numa_node: std::sync::atomic::AtomicI32, }
118
119impl DistancePool {
120 pub fn new(capacity: usize) -> Self {
122 Self::with_config(capacity, MemoryPoolConfig::default())
123 }
124
125 pub fn with_config(capacity: usize, config: MemoryPoolConfig) -> Self {
127 let numa_node = if config.numa_aware && config.numa_node_hint >= 0 {
128 config.numa_node_hint
129 } else {
130 Self::detect_numa_node()
131 };
132
133 Self {
134 config,
135 distance_buffers: Mutex::new(VecDeque::with_capacity(capacity)),
136 index_buffers: Mutex::new(VecDeque::with_capacity(capacity)),
137 matrix_buffers: Mutex::new(VecDeque::with_capacity(capacity / 4)), large_buffers: Mutex::new(VecDeque::with_capacity(capacity / 10)), stats: PoolStatistics::new(),
140 memory_usage: std::sync::atomic::AtomicUsize::new(0),
141 numa_node: std::sync::atomic::AtomicI32::new(numa_node),
142 }
143 }
144
145 pub fn get_distance_buffer(&self, size: usize) -> DistanceBuffer {
147 let buffer_size_bytes = size * std::mem::size_of::<f64>();
149 let is_large = buffer_size_bytes > self.config.large_object_threshold;
150
151 let current_usage = self.memory_usage.load(std::sync::atomic::Ordering::Relaxed);
153 if current_usage + buffer_size_bytes > self.config.max_memory_usage {
154 self.cleanup_excess_memory();
155 }
156
157 let buffer = if is_large {
158 self.get_large_buffer(size)
159 } else {
160 let mut buffers = self.distance_buffers.lock().expect("Operation failed");
161
162 for i in 0..buffers.len() {
164 if buffers[i].len() >= size && buffers[i].len() <= size * 2 {
165 let buffer = buffers.remove(i).expect("Operation failed");
166 self.stats.record_hit();
167 return DistanceBuffer::new(buffer, self);
168 }
169 }
170
171 self.stats.record_miss();
173 self.create_aligned_buffer(size)
174 };
175
176 self.memory_usage
178 .fetch_add(buffer_size_bytes, std::sync::atomic::Ordering::Relaxed);
179
180 DistanceBuffer::new(buffer, self)
181 }
182
183 fn get_large_buffer(&self, size: usize) -> Box<[f64]> {
185 let mut buffers = self.large_buffers.lock().expect("Operation failed");
186
187 for i in 0..buffers.len() {
189 if buffers[i].len() == size {
190 let buffer = buffers.remove(i).expect("Operation failed");
191 self.stats.record_hit();
192 return buffer;
193 }
194 }
195
196 self.stats.record_miss();
198 if self.config.numa_aware {
199 self.create_numa_aligned_buffer(size)
200 } else {
201 self.create_aligned_buffer(size)
202 }
203 }
204
205 pub fn get_index_buffer(&self, size: usize) -> IndexBuffer {
207 let mut buffers = self.index_buffers.lock().expect("Operation failed");
208
209 for i in 0..buffers.len() {
211 if buffers[i].len() >= size && buffers[i].len() <= size * 2 {
212 let buffer = buffers.remove(i).expect("Operation failed");
213 self.stats.record_hit();
214 return IndexBuffer::new(buffer, self);
215 }
216 }
217
218 self.stats.record_miss();
220 let new_buffer = vec![0usize; size].into_boxed_slice();
221 IndexBuffer::new(new_buffer, self)
222 }
223
224 pub fn get_matrix_buffer(&self, rows: usize, cols: usize) -> MatrixBuffer {
226 let mut buffers = self.matrix_buffers.lock().expect("Operation failed");
227
228 for i in 0..buffers.len() {
230 let (r, c) = buffers[i].dim();
231 if r >= rows && c >= cols && r <= rows * 2 && c <= cols * 2 {
232 let mut matrix = buffers.remove(i).expect("Operation failed");
233 matrix = matrix.slice_mut(s![..rows, ..cols]).to_owned();
235 self.stats.record_hit();
236 return MatrixBuffer::new(matrix, self);
237 }
238 }
239
240 self.stats.record_miss();
242 let matrix = Array2::zeros((rows, cols));
243 MatrixBuffer::new(matrix, self)
244 }
245
246 fn create_aligned_buffer(&self, size: usize) -> Box<[f64]> {
264 vec![0.0_f64; size].into_boxed_slice()
265 }
266
267 fn create_numa_aligned_buffer(&self, size: usize) -> Box<[f64]> {
269 let numa_node = self.numa_node.load(Ordering::Relaxed);
270
271 #[cfg(target_os = "linux")]
272 {
273 if self.config.numa_aware && numa_node >= 0 {
274 match Self::allocate_on_numa_node_linux(size, numa_node as u32) {
275 Ok(buffer) => {
276 if self.config.enable_memory_warming {
277 Self::warm_memory(&buffer);
278 }
279 return buffer;
280 }
281 Err(_) => {
282 }
284 }
285 }
286 }
287
288 #[cfg(target_os = "windows")]
289 {
290 if self.config.numa_aware && numa_node >= 0 {
291 match Self::allocate_on_numa_node_windows(size, numa_node as u32) {
292 Ok(buffer) => {
293 if self.config.enable_memory_warming {
294 Self::warm_memory(&buffer);
295 }
296 return buffer;
297 }
298 Err(_) => {
299 }
301 }
302 }
303 }
304
305 let buffer = self.create_aligned_buffer(size);
307
308 if self.config.enable_memory_warming {
310 Self::warm_memory(&buffer);
311 }
312
313 buffer
314 }
315
316 #[cfg(target_os = "linux")]
318 fn allocate_on_numa_node_linux(
319 size: usize,
320 node: u32,
321 ) -> Result<Box<[f64]>, Box<dyn std::error::Error>> {
322 Ok(vec![0.0_f64; size].into_boxed_slice())
326 }
327
328 #[cfg(target_os = "windows")]
330 fn allocate_on_numa_node_windows(
331 size: usize,
332 _node: u32,
333 ) -> Result<Box<[f64]>, Box<dyn std::error::Error>> {
334 Ok(vec![0.0_f64; size].into_boxed_slice())
338 }
339
340 pub fn bind_thread_to_numa_node(node: u32) -> Result<(), Box<dyn std::error::Error>> {
342 #[cfg(target_os = "linux")]
343 {
344 Self::bind_thread_to_numa_node_linux(node)
345 }
346 #[cfg(target_os = "windows")]
347 {
348 Self::bind_thread_to_numa_node_windows(node)
349 }
350 #[cfg(not(any(target_os = "linux", target_os = "windows")))]
351 {
352 Ok(()) }
354 }
355
356 #[cfg(target_os = "linux")]
357 fn bind_thread_to_numa_node_linux(node: u32) -> Result<(), Box<dyn std::error::Error>> {
358 if let Some(_cpu_count) = Self::get_node_cpu_count(node) {
363 let mut cpu_set: libc::cpu_set_t = unsafe { std::mem::zeroed() };
364
365 let cpulist_path = format!("/sys/devices/system/node/node{}/cpulist", node);
367 if let Ok(cpulist) = fs::read_to_string(&cpulist_path) {
368 for range in cpulist.trim().split(',') {
369 if let Some((start, end)) = range.split_once('-') {
370 if let (Ok(s), Ok(e)) = (start.parse::<u32>(), end.parse::<u32>()) {
371 for cpu in s..=e {
372 unsafe { libc::CPU_SET(cpu as usize, &mut cpu_set) };
373 }
374 }
375 } else if let Ok(cpu) = range.parse::<u32>() {
376 unsafe { libc::CPU_SET(cpu as usize, &mut cpu_set) };
377 }
378 }
379
380 unsafe {
382 libc::sched_setaffinity(
383 0, std::mem::size_of::<libc::cpu_set_t>(),
385 &cpu_set,
386 );
387 }
388 }
389 }
390
391 Ok(())
392 }
393
394 #[cfg(target_os = "windows")]
395 fn bind_thread_to_numa_node_windows(node: u32) -> Result<(), Box<dyn std::error::Error>> {
396 Ok(())
398 }
399
400 fn warm_memory(buffer: &[f64]) {
402 if buffer.is_empty() {
403 return;
404 }
405
406 let page_size = 4096; let elements_per_page = page_size / std::mem::size_of::<f64>();
409
410 for i in (0..buffer.len()).step_by(elements_per_page) {
411 unsafe {
413 std::ptr::read_volatile(&buffer[i]);
414 }
415 }
416 }
417
418 fn detect_numa_node() -> i32 {
420 #[cfg(target_os = "linux")]
421 {
422 Self::detect_numa_node_linux().unwrap_or(0)
423 }
424 #[cfg(target_os = "windows")]
425 {
426 Self::detect_numa_node_windows().unwrap_or(0)
427 }
428 #[cfg(not(any(target_os = "linux", target_os = "windows")))]
429 {
430 0 }
432 }
433
434 #[cfg(target_os = "linux")]
436 fn detect_numa_node_linux() -> Option<i32> {
437 let _tid = unsafe { libc::gettid() };
439
440 match Self::get_current_numa_node_linux() {
442 Ok(node) => Some(node),
443 Err(_) => {
444 Self::detect_numa_from_cpu_linux()
446 }
447 }
448 }
449
450 #[cfg(target_os = "linux")]
451 fn get_current_numa_node_linux() -> Result<i32, Box<dyn std::error::Error>> {
452 let mut cpu: u32 = 0;
454 let mut node: u32 = 0;
455
456 let result = unsafe {
457 libc::syscall(
458 libc::SYS_getcpu,
459 &mut cpu as *mut u32,
460 &mut node as *mut u32,
461 std::ptr::null_mut::<libc::c_void>(),
462 )
463 };
464
465 if result == 0 {
466 Ok(node as i32)
467 } else {
468 Err("getcpu syscall failed".into())
469 }
470 }
471
472 #[cfg(target_os = "linux")]
473 fn detect_numa_from_cpu_linux() -> Option<i32> {
474 if let Ok(entries) = fs::read_dir("/sys/devices/system/node") {
476 for entry in entries.flatten() {
477 let name = entry.file_name();
478 if let Some(name_str) = name.to_str() {
479 if let Some(stripped) = name_str.strip_prefix("node") {
480 if let Ok(node_num) = stripped.parse::<i32>() {
481 return Some(node_num);
483 }
484 }
485 }
486 }
487 }
488 None
489 }
490
491 #[cfg(target_os = "windows")]
493 fn detect_numa_node_windows() -> Option<i32> {
494 Some(0)
498 }
499
500 pub fn get_numa_topology() -> NumaTopology {
502 #[cfg(target_os = "linux")]
503 {
504 Self::get_numa_topology_linux()
505 }
506 #[cfg(target_os = "windows")]
507 {
508 Self::get_numa_topology_windows()
509 }
510 #[cfg(not(any(target_os = "linux", target_os = "windows")))]
511 {
512 NumaTopology::default()
513 }
514 }
515
516 #[cfg(target_os = "linux")]
517 fn get_numa_topology_linux() -> NumaTopology {
518 let mut topology = NumaTopology::default();
519
520 if let Ok(entries) = fs::read_dir("/sys/devices/system/node") {
522 for entry in entries.flatten() {
523 let name = entry.file_name();
524 if let Some(name_str) = name.to_str() {
525 if let Some(stripped) = name_str.strip_prefix("node") {
526 if let Ok(_nodeid) = stripped.parse::<u32>() {
527 let meminfo_path =
529 format!("/sys/devices/system/node/{name_str}/meminfo");
530 if let Ok(meminfo) = fs::read_to_string(&meminfo_path) {
531 if let Some(total_kb) = Self::parse_meminfo_total(&meminfo) {
532 topology.nodes.push(NumaNode {
533 id: _nodeid,
534 total_memory_bytes: total_kb * 1024,
535 available_memory_bytes: total_kb * 1024, cpu_count: Self::get_node_cpu_count(_nodeid).unwrap_or(1),
537 });
538 }
539 }
540 }
541 }
542 }
543 }
544 }
545
546 if topology.nodes.is_empty() {
548 topology.nodes.push(NumaNode {
549 id: 0,
550 total_memory_bytes: Self::get_total_system_memory()
551 .unwrap_or(8 * 1024 * 1024 * 1024), available_memory_bytes: Self::get_available_system_memory()
553 .unwrap_or(4 * 1024 * 1024 * 1024), cpu_count: num_cpus::get() as u32,
555 });
556 }
557
558 topology
559 }
560
561 #[cfg(target_os = "linux")]
562 fn parse_meminfo_total(meminfo: &str) -> Option<u64> {
563 for line in meminfo.lines() {
564 if line.starts_with("Node") && line.contains("MemTotal:") {
565 let parts: Vec<&str> = line.split_whitespace().collect();
566 if parts.len() >= 3 {
567 return parts[2].parse().ok();
568 }
569 }
570 }
571 None
572 }
573
574 #[cfg(target_os = "linux")]
575 fn get_node_cpu_count(_nodeid: u32) -> Option<u32> {
576 let cpulist_path = format!("/sys/devices/system/node/node{}/cpulist", _nodeid);
577 if let Ok(cpulist) = fs::read_to_string(&cpulist_path) {
578 let mut count = 0;
580 for range in cpulist.trim().split(',') {
581 if let Some((start, end)) = range.split_once('-') {
582 if let (Ok(s), Ok(e)) = (start.parse::<u32>(), end.parse::<u32>()) {
583 count += e - s + 1;
584 }
585 } else if range.parse::<u32>().is_ok() {
586 count += 1;
587 }
588 }
589 Some(count)
590 } else {
591 None
592 }
593 }
594
595 #[cfg(target_os = "linux")]
596 fn get_total_system_memory() -> Option<u64> {
597 if let Ok(meminfo) = fs::read_to_string("/proc/meminfo") {
598 for line in meminfo.lines() {
599 if line.starts_with("MemTotal:") {
600 let parts: Vec<&str> = line.split_whitespace().collect();
601 if parts.len() >= 2 {
602 return parts[1].parse::<u64>().ok().map(|kb| kb * 1024);
603 }
604 }
605 }
606 }
607 None
608 }
609
610 #[cfg(target_os = "linux")]
611 fn get_available_system_memory() -> Option<u64> {
612 if let Ok(meminfo) = fs::read_to_string("/proc/meminfo") {
613 for line in meminfo.lines() {
614 if line.starts_with("MemAvailable:") {
615 let parts: Vec<&str> = line.split_whitespace().collect();
616 if parts.len() >= 2 {
617 return parts[1].parse::<u64>().ok().map(|kb| kb * 1024);
618 }
619 }
620 }
621 }
622 None
623 }
624
625 #[cfg(target_os = "windows")]
626 fn get_numa_topology_windows() -> NumaTopology {
627 NumaTopology::default()
630 }
631
632 fn cleanup_excess_memory(&self) {
634 let cleanup_ratio = 0.25; {
638 let mut buffers = self.distance_buffers.lock().expect("Operation failed");
639 let cleanup_count = (buffers.len() as f64 * cleanup_ratio) as usize;
640 for _ in 0..cleanup_count {
641 if let Some(buffer) = buffers.pop_back() {
642 let freed_bytes = buffer.len() * std::mem::size_of::<f64>();
643 self.memory_usage
644 .fetch_sub(freed_bytes, std::sync::atomic::Ordering::Relaxed);
645 }
646 }
647 }
648
649 {
650 let mut buffers = self.large_buffers.lock().expect("Operation failed");
651 let cleanup_count = (buffers.len() as f64 * cleanup_ratio) as usize;
652 for _ in 0..cleanup_count {
653 if let Some(buffer) = buffers.pop_back() {
654 let freed_bytes = buffer.len() * std::mem::size_of::<f64>();
655 self.memory_usage
656 .fetch_sub(freed_bytes, std::sync::atomic::Ordering::Relaxed);
657 }
658 }
659 }
660 }
661
662 fn return_distance_buffer(&self, buffer: Box<[f64]>) {
664 let buffer_size_bytes = buffer.len() * std::mem::size_of::<f64>();
665 let is_large = buffer_size_bytes > self.config.large_object_threshold;
666
667 self.memory_usage
669 .fetch_sub(buffer_size_bytes, std::sync::atomic::Ordering::Relaxed);
670
671 if is_large {
672 let mut buffers = self.large_buffers.lock().expect("Operation failed");
673 if buffers.len() < self.config.max_pool_size / 10 {
674 buffers.push_back(buffer);
675 }
676 } else {
678 let mut buffers = self.distance_buffers.lock().expect("Operation failed");
679 if buffers.len() < self.config.max_pool_size {
680 buffers.push_back(buffer);
681 }
682 }
684 }
685
686 fn return_index_buffer(&self, buffer: Box<[usize]>) {
688 let mut buffers = self.index_buffers.lock().expect("Operation failed");
689 if buffers.len() < self.config.max_pool_size {
690 buffers.push_back(buffer);
691 }
692 }
693
694 fn return_matrix_buffer(&self, matrix: Array2<f64>) {
696 let mut buffers = self.matrix_buffers.lock().expect("Operation failed");
697 if buffers.len() < self.config.max_pool_size / 4 {
698 buffers.push_back(matrix);
700 }
701 }
702
703 pub fn statistics(&self) -> PoolStatistics {
705 self.stats.clone()
706 }
707
708 pub fn memory_usage(&self) -> usize {
710 self.memory_usage.load(std::sync::atomic::Ordering::Relaxed)
711 }
712
713 pub fn current_numa_node(&self) -> i32 {
715 self.numa_node.load(std::sync::atomic::Ordering::Relaxed)
716 }
717
718 pub fn pool_info(&self) -> PoolInfo {
720 let distance_count = self
721 .distance_buffers
722 .lock()
723 .expect("Operation failed")
724 .len();
725 let index_count = self.index_buffers.lock().expect("Operation failed").len();
726 let matrix_count = self.matrix_buffers.lock().expect("Operation failed").len();
727 let large_count = self.large_buffers.lock().expect("Operation failed").len();
728
729 PoolInfo {
730 distance_buffer_count: distance_count,
731 index_buffer_count: index_count,
732 matrix_buffer_count: matrix_count,
733 large_buffer_count: large_count,
734 total_memory_usage: self.memory_usage(),
735 numa_node: self.current_numa_node(),
736 hit_rate: self.stats.hit_rate(),
737 }
738 }
739
740 pub fn clear(&self) {
742 self.distance_buffers
743 .lock()
744 .expect("Operation failed")
745 .clear();
746 self.index_buffers.lock().expect("Operation failed").clear();
747 self.matrix_buffers
748 .lock()
749 .expect("Operation failed")
750 .clear();
751 self.large_buffers.lock().expect("Operation failed").clear();
752 self.memory_usage
753 .store(0, std::sync::atomic::Ordering::Relaxed);
754 self.stats.reset();
755 }
756}
757
758use scirs2_core::ndarray::s;
760
761pub struct DistanceBuffer<'a> {
763 buffer: Option<Box<[f64]>>,
764 pool: &'a DistancePool,
765}
766
767impl<'a> DistanceBuffer<'a> {
768 fn new(buffer: Box<[f64]>, pool: &'a DistancePool) -> Self {
769 Self {
770 buffer: Some(buffer),
771 pool,
772 }
773 }
774
775 pub fn as_mut_slice(&mut self) -> &mut [f64] {
777 self.buffer.as_mut().expect("Operation failed").as_mut()
778 }
779
780 pub fn as_slice(&self) -> &[f64] {
782 self.buffer.as_ref().expect("Operation failed").as_ref()
783 }
784
785 pub fn len(&self) -> usize {
787 self.buffer.as_ref().expect("Operation failed").len()
788 }
789
790 pub fn is_empty(&self) -> bool {
792 self.len() == 0
793 }
794
795 pub fn as_array_mut(&mut self) -> ArrayViewMut1<f64> {
797 ArrayViewMut1::from(self.as_mut_slice())
798 }
799}
800
801impl Drop for DistanceBuffer<'_> {
802 fn drop(&mut self) {
803 if let Some(buffer) = self.buffer.take() {
804 self.pool.return_distance_buffer(buffer);
805 }
806 }
807}
808
809pub struct IndexBuffer<'a> {
811 buffer: Option<Box<[usize]>>,
812 pool: &'a DistancePool,
813}
814
815impl<'a> IndexBuffer<'a> {
816 fn new(buffer: Box<[usize]>, pool: &'a DistancePool) -> Self {
817 Self {
818 buffer: Some(buffer),
819 pool,
820 }
821 }
822
823 pub fn as_mut_slice(&mut self) -> &mut [usize] {
825 self.buffer.as_mut().expect("Operation failed").as_mut()
826 }
827
828 pub fn as_slice(&self) -> &[usize] {
830 self.buffer.as_ref().expect("Operation failed").as_ref()
831 }
832
833 pub fn len(&self) -> usize {
835 self.buffer.as_ref().expect("Operation failed").len()
836 }
837
838 pub fn is_empty(&self) -> bool {
840 self.len() == 0
841 }
842}
843
844impl Drop for IndexBuffer<'_> {
845 fn drop(&mut self) {
846 if let Some(buffer) = self.buffer.take() {
847 self.pool.return_index_buffer(buffer);
848 }
849 }
850}
851
852pub struct MatrixBuffer<'a> {
854 matrix: Option<Array2<f64>>,
855 pool: &'a DistancePool,
856}
857
858impl<'a> MatrixBuffer<'a> {
859 fn new(matrix: Array2<f64>, pool: &'a DistancePool) -> Self {
860 Self {
861 matrix: Some(matrix),
862 pool,
863 }
864 }
865
866 pub fn as_mut(&mut self) -> ArrayViewMut2<f64> {
868 self.matrix.as_mut().expect("Operation failed").view_mut()
869 }
870
871 pub fn dim(&mut self) -> (usize, usize) {
873 self.matrix.as_ref().expect("Operation failed").dim()
874 }
875
876 pub fn fill(&mut self, value: f64) {
878 self.matrix.as_mut().expect("Operation failed").fill(value);
879 }
880}
881
882impl Drop for MatrixBuffer<'_> {
883 fn drop(&mut self) {
884 if let Some(matrix) = self.matrix.take() {
885 self.pool.return_matrix_buffer(matrix);
886 }
887 }
888}
889
890pub struct ClusteringArena {
892 config: MemoryPoolConfig,
893 current_block: Mutex<Option<ArenaBlock>>,
894 full_blocks: Mutex<Vec<ArenaBlock>>,
895 stats: ArenaStatistics,
896}
897
898impl ClusteringArena {
899 pub fn new() -> Self {
901 Self::with_config(MemoryPoolConfig::default())
902 }
903
904 pub fn with_config(config: MemoryPoolConfig) -> Self {
906 Self {
907 config,
908 current_block: Mutex::new(None),
909 full_blocks: Mutex::new(Vec::new()),
910 stats: ArenaStatistics::new(),
911 }
912 }
913
914 pub fn alloc_temp_vec<T: Default + Clone>(&self, size: usize) -> ArenaVec<T> {
916 let layout = Layout::array::<T>(size).expect("Operation failed");
917 let ptr = self.allocate_raw(layout);
918
919 unsafe {
920 for i in 0..size {
922 std::ptr::write(ptr.as_ptr().add(i) as *mut T, T::default());
923 }
924
925 ArenaVec::new(ptr.as_ptr() as *mut T, size)
926 }
927 }
928
929 fn allocate_raw(&self, layout: Layout) -> NonNull<u8> {
931 let mut current = self.current_block.lock().expect("Operation failed");
932
933 if current.is_none()
934 || !current
935 .as_ref()
936 .expect("Operation failed")
937 .can_allocate(layout)
938 {
939 if let Some(old_block) = current.take() {
941 self.full_blocks
942 .lock()
943 .expect("Operation failed")
944 .push(old_block);
945 }
946 *current = Some(ArenaBlock::new(self.config.arena_block_size));
947 }
948
949 current.as_mut().expect("Operation failed").allocate(layout)
950 }
951
952 pub fn reset(&self) {
954 let mut current = self.current_block.lock().expect("Operation failed");
955 let mut full_blocks = self.full_blocks.lock().expect("Operation failed");
956
957 if let Some(block) = current.take() {
958 full_blocks.push(block);
959 }
960
961 for block in full_blocks.iter_mut() {
963 block.reset();
964 }
965
966 if let Some(block) = full_blocks.pop() {
968 *current = Some(block);
969 }
970
971 self.stats.reset();
972 }
973
974 pub fn statistics(&self) -> ArenaStatistics {
976 self.stats.clone()
977 }
978}
979
980impl Default for ClusteringArena {
981 fn default() -> Self {
982 Self::new()
983 }
984}
985
986struct ArenaBlock {
988 memory: NonNull<u8>,
989 size: usize,
990 offset: usize,
991}
992
993unsafe impl Send for ArenaBlock {}
995unsafe impl Sync for ArenaBlock {}
996
997impl ArenaBlock {
998 fn new(size: usize) -> Self {
999 let layout = Layout::from_size_align(size, 64).expect("Operation failed"); let memory =
1001 unsafe { NonNull::new(System.alloc(layout)).expect("Failed to allocate arena block") };
1002
1003 Self {
1004 memory,
1005 size,
1006 offset: 0,
1007 }
1008 }
1009
1010 fn can_allocate(&self, layout: Layout) -> bool {
1011 let aligned_offset = (self.offset + layout.align() - 1) & !(layout.align() - 1);
1012 aligned_offset + layout.size() <= self.size
1013 }
1014
1015 fn allocate(&mut self, layout: Layout) -> NonNull<u8> {
1016 assert!(self.can_allocate(layout));
1017
1018 self.offset = (self.offset + layout.align() - 1) & !(layout.align() - 1);
1020
1021 let ptr = unsafe { NonNull::new_unchecked(self.memory.as_ptr().add(self.offset)) };
1022 self.offset += layout.size();
1023
1024 ptr
1025 }
1026
1027 fn reset(&mut self) {
1028 self.offset = 0;
1029 }
1030}
1031
1032impl Drop for ArenaBlock {
1033 fn drop(&mut self) {
1034 let layout = Layout::from_size_align(self.size, 64).expect("Operation failed");
1035 unsafe {
1036 System.dealloc(self.memory.as_ptr(), layout);
1037 }
1038 }
1039}
1040
1041pub struct ArenaVec<T> {
1043 ptr: *mut T,
1044 len: usize,
1045 phantom: std::marker::PhantomData<T>,
1046}
1047
1048impl<T> ArenaVec<T> {
1049 fn new(ptr: *mut T, len: usize) -> Self {
1050 Self {
1051 ptr,
1052 len,
1053 phantom: std::marker::PhantomData,
1054 }
1055 }
1056
1057 pub fn as_mut_slice(&mut self) -> &mut [T] {
1059 unsafe { std::slice::from_raw_parts_mut(self.ptr, self.len) }
1060 }
1061
1062 pub fn as_slice(&mut self) -> &[T] {
1064 unsafe { std::slice::from_raw_parts(self.ptr, self.len) }
1065 }
1066
1067 pub fn len(&mut self) -> usize {
1069 self.len
1070 }
1071
1072 pub fn is_empty(&self) -> bool {
1074 self.len == 0
1075 }
1076}
1077
1078#[derive(Debug, Clone)]
1082pub struct PoolInfo {
1083 pub distance_buffer_count: usize,
1085 pub index_buffer_count: usize,
1087 pub matrix_buffer_count: usize,
1089 pub large_buffer_count: usize,
1091 pub total_memory_usage: usize,
1093 pub numa_node: i32,
1095 pub hit_rate: f64,
1097}
1098
1099#[derive(Debug)]
1101pub struct PoolStatistics {
1102 hits: std::sync::atomic::AtomicUsize,
1103 misses: std::sync::atomic::AtomicUsize,
1104 total_allocations: std::sync::atomic::AtomicUsize,
1105}
1106
1107impl PoolStatistics {
1108 fn new() -> Self {
1109 Self {
1110 hits: std::sync::atomic::AtomicUsize::new(0),
1111 misses: std::sync::atomic::AtomicUsize::new(0),
1112 total_allocations: std::sync::atomic::AtomicUsize::new(0),
1113 }
1114 }
1115
1116 fn record_hit(&self) {
1117 self.hits.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
1118 }
1119
1120 fn record_miss(&self) {
1121 self.misses
1122 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
1123 self.total_allocations
1124 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
1125 }
1126
1127 fn reset(&self) {
1128 self.hits.store(0, std::sync::atomic::Ordering::Relaxed);
1129 self.misses.store(0, std::sync::atomic::Ordering::Relaxed);
1130 self.total_allocations
1131 .store(0, std::sync::atomic::Ordering::Relaxed);
1132 }
1133
1134 pub fn hit_rate(&self) -> f64 {
1136 let hits = self.hits.load(std::sync::atomic::Ordering::Relaxed);
1137 let total = hits + self.misses.load(std::sync::atomic::Ordering::Relaxed);
1138 if total == 0 {
1139 0.0
1140 } else {
1141 hits as f64 / total as f64 * 100.0
1142 }
1143 }
1144
1145 pub fn total_requests(&self) -> usize {
1147 self.hits.load(std::sync::atomic::Ordering::Relaxed)
1148 + self.misses.load(std::sync::atomic::Ordering::Relaxed)
1149 }
1150
1151 pub fn total_allocations(&self) -> usize {
1153 self.total_allocations
1154 .load(std::sync::atomic::Ordering::Relaxed)
1155 }
1156}
1157
1158impl Clone for PoolStatistics {
1159 fn clone(&self) -> Self {
1160 Self {
1161 hits: std::sync::atomic::AtomicUsize::new(
1162 self.hits.load(std::sync::atomic::Ordering::Relaxed),
1163 ),
1164 misses: std::sync::atomic::AtomicUsize::new(
1165 self.misses.load(std::sync::atomic::Ordering::Relaxed),
1166 ),
1167 total_allocations: std::sync::atomic::AtomicUsize::new(
1168 self.total_allocations
1169 .load(std::sync::atomic::Ordering::Relaxed),
1170 ),
1171 }
1172 }
1173}
1174
1175#[derive(Debug)]
1177pub struct ArenaStatistics {
1178 blocks_allocated: std::sync::atomic::AtomicUsize,
1179 total_memory: std::sync::atomic::AtomicUsize,
1180 active_objects: std::sync::atomic::AtomicUsize,
1181}
1182
1183impl ArenaStatistics {
1184 fn new() -> Self {
1185 Self {
1186 blocks_allocated: std::sync::atomic::AtomicUsize::new(0),
1187 total_memory: std::sync::atomic::AtomicUsize::new(0),
1188 active_objects: std::sync::atomic::AtomicUsize::new(0),
1189 }
1190 }
1191
1192 fn reset(&self) {
1193 self.blocks_allocated
1194 .store(0, std::sync::atomic::Ordering::Relaxed);
1195 self.total_memory
1196 .store(0, std::sync::atomic::Ordering::Relaxed);
1197 self.active_objects
1198 .store(0, std::sync::atomic::Ordering::Relaxed);
1199 }
1200
1201 pub fn blocks_allocated(&self) -> usize {
1203 self.blocks_allocated
1204 .load(std::sync::atomic::Ordering::Relaxed)
1205 }
1206
1207 pub fn total_memory(&self) -> usize {
1209 self.total_memory.load(std::sync::atomic::Ordering::Relaxed)
1210 }
1211
1212 pub fn active_objects(&self) -> usize {
1214 self.active_objects
1215 .load(std::sync::atomic::Ordering::Relaxed)
1216 }
1217}
1218
1219impl Clone for ArenaStatistics {
1220 fn clone(&self) -> Self {
1221 Self {
1222 blocks_allocated: std::sync::atomic::AtomicUsize::new(
1223 self.blocks_allocated
1224 .load(std::sync::atomic::Ordering::Relaxed),
1225 ),
1226 total_memory: std::sync::atomic::AtomicUsize::new(
1227 self.total_memory.load(std::sync::atomic::Ordering::Relaxed),
1228 ),
1229 active_objects: std::sync::atomic::AtomicUsize::new(
1230 self.active_objects
1231 .load(std::sync::atomic::Ordering::Relaxed),
1232 ),
1233 }
1234 }
1235}
1236
1237#[derive(Debug, Clone)]
1239pub struct NumaTopology {
1240 pub nodes: Vec<NumaNode>,
1242}
1243
1244#[derive(Debug, Clone)]
1246pub struct NumaNode {
1247 pub id: u32,
1249 pub total_memory_bytes: u64,
1251 pub available_memory_bytes: u64,
1253 pub cpu_count: u32,
1255}
1256
1257impl Default for NumaTopology {
1258 fn default() -> Self {
1259 Self {
1260 nodes: vec![NumaNode {
1261 id: 0,
1262 total_memory_bytes: 8 * 1024 * 1024 * 1024, available_memory_bytes: 4 * 1024 * 1024 * 1024, cpu_count: 4, }],
1266 }
1267 }
1268}
1269
1270impl NumaTopology {
1271 pub fn get_optimal_node(&self) -> u32 {
1273 if !self.nodes.is_empty() {
1276 self.nodes[0].id
1277 } else {
1278 0
1279 }
1280 }
1281
1282 pub fn get_node_with_most_memory(&self) -> Option<u32> {
1284 self.nodes
1285 .iter()
1286 .max_by_key(|node| node.available_memory_bytes)
1287 .map(|node| node.id)
1288 }
1289
1290 pub fn total_system_memory(&self) -> u64 {
1292 self.nodes.iter().map(|node| node.total_memory_bytes).sum()
1293 }
1294
1295 pub fn total_available_memory(&self) -> u64 {
1297 self.nodes
1298 .iter()
1299 .map(|node| node.available_memory_bytes)
1300 .sum()
1301 }
1302
1303 pub fn has_node(&self, _nodeid: u32) -> bool {
1305 self.nodes.iter().any(|node| node.id == _nodeid)
1306 }
1307
1308 pub fn get_node_info(&self, _nodeid: u32) -> Option<&NumaNode> {
1310 self.nodes.iter().find(|node| node.id == _nodeid)
1311 }
1312}
1313
1314static GLOBAL_DISTANCE_POOL: std::sync::OnceLock<DistancePool> = std::sync::OnceLock::new();
1316static GLOBAL_CLUSTERING_ARENA: std::sync::OnceLock<ClusteringArena> = std::sync::OnceLock::new();
1317
1318#[allow(dead_code)]
1320pub fn global_distance_pool() -> &'static DistancePool {
1321 GLOBAL_DISTANCE_POOL.get_or_init(|| DistancePool::new(1000))
1322}
1323
1324#[allow(dead_code)]
1326pub fn global_clustering_arena() -> &'static ClusteringArena {
1327 GLOBAL_CLUSTERING_ARENA.get_or_init(ClusteringArena::new)
1328}
1329
1330#[allow(dead_code)]
1332pub fn create_numa_optimized_pool(capacity: usize) -> DistancePool {
1333 let config = MemoryPoolConfig {
1334 numa_aware: true,
1335 auto_numa_discovery: true,
1336 enable_thread_affinity: true,
1337 ..Default::default()
1338 };
1339
1340 DistancePool::with_config(capacity, config)
1341}
1342
1343#[allow(dead_code)]
1345pub fn get_numa_topology() -> NumaTopology {
1346 DistancePool::get_numa_topology()
1347}
1348
1349#[allow(dead_code)]
1351pub fn test_numa_capabilities() -> NumaCapabilities {
1352 NumaCapabilities::detect()
1353}
1354
1355#[derive(Debug, Clone)]
1357pub struct NumaCapabilities {
1358 pub numa_available: bool,
1360 pub num_nodes: u32,
1362 pub memory_binding_supported: bool,
1364 pub thread_affinity_supported: bool,
1366 pub platform_details: String,
1368}
1369
1370impl NumaCapabilities {
1371 pub fn detect() -> Self {
1373 #[cfg(target_os = "linux")]
1374 {
1375 Self::detect_linux()
1376 }
1377 #[cfg(target_os = "windows")]
1378 {
1379 Self::detect_windows()
1380 }
1381 #[cfg(not(any(target_os = "linux", target_os = "windows")))]
1382 {
1383 Self {
1384 numa_available: false,
1385 num_nodes: 1,
1386 memory_binding_supported: false,
1387 thread_affinity_supported: false,
1388 platform_details: "Unsupported platform".to_string(),
1389 }
1390 }
1391 }
1392
1393 #[cfg(target_os = "linux")]
1394 fn detect_linux() -> Self {
1395 let numa_available = std::path::Path::new("/sys/devices/system/node").exists();
1396 let num_nodes = if numa_available {
1397 DistancePool::get_numa_topology().nodes.len() as u32
1398 } else {
1399 1
1400 };
1401
1402 Self {
1403 numa_available,
1404 num_nodes,
1405 memory_binding_supported: numa_available,
1406 thread_affinity_supported: true, platform_details: format!("Linux with {num_nodes} NUMA nodes"),
1408 }
1409 }
1410
1411 #[cfg(target_os = "windows")]
1412 fn detect_windows() -> Self {
1413 Self {
1414 numa_available: true, num_nodes: 1, memory_binding_supported: true,
1417 thread_affinity_supported: true,
1418 platform_details: "Windows NUMA support".to_string(),
1419 }
1420 }
1421
1422 pub fn should_enable_numa(&self) -> bool {
1424 self.numa_available && self.num_nodes > 1
1425 }
1426
1427 pub fn recommended_memory_strategy(&self) -> &'static str {
1429 if self.should_enable_numa() {
1430 "NUMA-aware"
1431 } else {
1432 "Standard"
1433 }
1434 }
1435}
1436
1437#[cfg(test)]
1438mod tests {
1439 use super::*;
1440
1441 #[test]
1442 fn test_distance_pool() {
1443 let pool = DistancePool::new(10);
1444
1445 let mut buffer1 = pool.get_distance_buffer(100);
1447 assert_eq!(buffer1.len(), 100);
1448
1449 buffer1.as_mut_slice()[0] = 42.0;
1451 assert_eq!(buffer1.as_slice()[0], 42.0);
1452
1453 let buffer2 = pool.get_distance_buffer(50);
1455 assert_eq!(buffer2.len(), 50);
1456
1457 drop(buffer1);
1459
1460 let buffer3 = pool.get_distance_buffer(100);
1462 assert_eq!(buffer3.len(), 100);
1463 }
1465
1466 #[test]
1467 fn test_arena_allocator() {
1468 let arena = ClusteringArena::new();
1469
1470 let mut vec1 = arena.alloc_temp_vec::<f64>(100);
1472 let mut vec2 = arena.alloc_temp_vec::<usize>(50);
1473
1474 vec1.as_mut_slice()[0] = std::f64::consts::PI;
1476 vec2.as_mut_slice()[0] = 42;
1477
1478 assert_eq!(vec1.as_slice()[0], std::f64::consts::PI);
1479 assert_eq!(vec2.as_slice()[0], 42);
1480
1481 arena.reset();
1483
1484 let mut vec3 = arena.alloc_temp_vec::<f64>(200);
1486 vec3.as_mut_slice()[0] = 2.71;
1487 assert_eq!(vec3.as_slice()[0], 2.71);
1488 }
1489
1490 #[test]
1491 fn test_pool_statistics() {
1492 let pool = DistancePool::new(2);
1493
1494 let stats = pool.statistics();
1496 assert_eq!(stats.total_requests(), 0);
1497 assert_eq!(stats.total_allocations(), 0);
1498
1499 let _buffer1 = pool.get_distance_buffer(100);
1501 let stats = pool.statistics();
1502 assert_eq!(stats.total_requests(), 1);
1503 assert_eq!(stats.total_allocations(), 1);
1504 assert!(stats.hit_rate() < 1.0);
1505
1506 drop(_buffer1);
1508 let _buffer2 = pool.get_distance_buffer(100);
1509 let stats = pool.statistics();
1510 assert_eq!(stats.total_requests(), 2);
1511 assert_eq!(stats.total_allocations(), 1); assert!(stats.hit_rate() > 0.0);
1513 }
1514
1515 #[test]
1516 fn test_matrix_buffer() {
1517 let pool = DistancePool::new(5);
1518
1519 let mut matrix = pool.get_matrix_buffer(10, 10);
1520 assert_eq!(matrix.dim(), (10, 10));
1521
1522 matrix.fill(42.0);
1523 drop(matrix);
1526
1527 let mut matrix2 = pool.get_matrix_buffer(8, 8);
1529 assert_eq!(matrix2.dim(), (8, 8));
1530 }
1531
1532 #[test]
1533 fn test_global_pools() {
1534 let pool = global_distance_pool();
1536 let arena = global_clustering_arena();
1537
1538 let buffer = pool.get_distance_buffer(10);
1539 let _vec = arena.alloc_temp_vec::<f64>(10);
1540
1541 }
1543
1544 #[cfg(target_os = "windows")]
1545 #[test]
1546 fn test_windows_numa_fallback_returns_ok() {
1547 let result = DistancePool::allocate_on_numa_node_windows(1024, 0);
1548 assert!(result.is_ok());
1549 let buf = result.expect("allocation should succeed");
1550 assert_eq!(buf.len(), 1024);
1551 }
1552}