prevent mixing shmem affinities in realloc and dealloc

Author: hunhoffe

kernel/src/memory/mod.rs

@@ -631,7 +631,7 @@ unsafe impl GlobalAlloc for KernelAllocator {
             |pcm| {
                 #[cfg(feature = "rackscale")]
                 {
-                    use crate::transport::shmem::{is_shmem_addr, SHMEM_INITIALIZED};
+                    use crate::transport::shmem::{is_shmem_addr, is_shmem_addr_with_affinity, SHMEM_INITIALIZED};
                     use core::sync::atomic::Ordering;
 
                     // If shmem is not initialized, do not force it.
@@ -640,10 +640,10 @@ unsafe impl GlobalAlloc for KernelAllocator {
                         let affinity = { pcm.physical_memory.borrow().affinity };
                         if is_shmem_addr(ptr as u64, false, false) {
                             panic!("Should not be trying to dealloc non-kernel mapped shmem in kernel dealloc");
-                        } else if is_shmem_affinity(affinity) && !is_shmem_addr(ptr as u64, false, true) {
-                            log::error!("Trying to deallocate non-shmem into shmem allocator - losing this memory. Oh well.");
+                        } else if is_shmem_affinity(affinity) && !is_shmem_addr_with_affinity(ptr as u64, affinity, true) {
+                            log::error!("Trying to deallocate memory not in shmem affinity into shmem allocator - losing this memory. Oh well.");
                             return;
-                        } else if !is_shmem_affinity(affinity) && (is_shmem_addr(ptr as u64, false, true)){
+                        } else if !is_shmem_affinity(affinity) && is_shmem_addr(ptr as u64, false, true) {
                             log::error!("Trying to deallocate shmem into non-shmem allocator - losing this memory. Oh well.");
                             return;
                         }
@@ -743,18 +743,18 @@ unsafe impl GlobalAlloc for KernelAllocator {
                 } else {
                     #[cfg(feature = "rackscale")]
                     {
-                        use crate::transport::shmem::{is_shmem_addr, SHMEM_INITIALIZED};
+                        use crate::transport::shmem::{is_shmem_addr, is_shmem_addr_with_affinity, SHMEM_INITIALIZED};
                         use core::sync::atomic::Ordering;
 
                         // If shmem is not initialized, do not force it.
                         if 0 != SHMEM_INITIALIZED.load(Ordering::SeqCst) {
                             let affinity = { pcm.physical_memory.borrow().affinity };
                             if is_shmem_addr(ptr as u64, false, false) {
                                 panic!("Should not be trying to realloc non-kernel mapped shmem in kernel dealloc");
-                            } else if is_shmem_affinity(affinity) && !is_shmem_addr(ptr as u64, false, true) {
+                            } else if is_shmem_affinity(affinity) && !is_shmem_addr_with_affinity(ptr as u64, affinity, true) {
                                 // TODO(rackscale): should switch to non-shmem affinity for alloc below.
                                 // TODO(rackscale): check if shmem is a match for id?
-                                panic!("Trying to realloc non-shmem using shmem allocator");
+                                panic!("Trying to realloc shmem to wrong or non- shmem allocator");
                             } else if !is_shmem_affinity(affinity) && is_shmem_addr(ptr as u64, false, true) {
                                 // TODO(rackscale): should switch to use shmem affinity for alloc below.
                                 // TODO(rackscale): check if shmem is a match for id?

kernel/src/transport/shmem.rs

@@ -22,6 +22,7 @@ use crate::pci::claim_devices;
 use {
     crate::cmdline::{Mode, Transport},
     crate::error::{KError, KResult},
+    crate::memory::shmem_affinity::shmem_affinity_to_mid,
     alloc::boxed::Box,
     kpi::system::MachineId,
     rpc::rpc::MAX_BUFF_LEN,
@@ -444,6 +445,15 @@ pub(crate) fn is_shmem_addr(addr: u64, is_affinity: bool, is_kaddr: bool) -> boo
     }
 }
 
+#[cfg(feature = "rackscale")]
+#[inline(always)]
+pub(crate) fn is_shmem_addr_with_affinity(addr: u64, affinity: NodeId, is_kaddr: bool) -> bool {
+    let offset = if is_kaddr { KERNEL_BASE } else { 0 };
+    let mid = shmem_affinity_to_mid(affinity);
+    let frame = get_affinity_shmem_by_mid(mid);
+    addr >= frame.base.as_u64() + offset && addr < frame.base.as_u64() + offset + frame.size as u64
+}
+
 #[cfg(feature = "rackscale")]
 pub(crate) fn get_shmem_index_by_addr(addr: u64, is_kaddr: bool) -> Option<usize> {
     let offset = if is_kaddr { KERNEL_BASE } else { 0 };