First step towards rackscale shared memory KernArcs

kernel/src/memory/mod.rs

@@ -102,7 +102,15 @@ impl KernelAllocator {
                 {
                     let affinity = { pcm.physical_memory.borrow().affinity };
                     if is_shmem_affinity(affinity) {
-                        panic!("MapBig not yet supported for shmem allocation");
+                        // Treating as type MemManager
+                        log::warn!("MapBig not supported for Rackscale.");
+                        let f = {
+                            let mut pmanager = pcm.try_mem_manager()?;
+                            pmanager.allocate_large_page()?
+                        };
+                        return unsafe {
+                            Ok(ptr::NonNull::new_unchecked(f.kernel_vaddr().as_mut_ptr()))
+                        };
                     }
                 }
                 // Big objects are mapped into the kernel address space

kernel/src/nrproc.rs

@@ -24,6 +24,9 @@ use crate::process::{
     Eid, Executor, Pid, Process, SliceAccess, UserSlice, MAX_FRAMES_PER_PROCESS, MAX_PROCESSES,
 };
 
+#[cfg(feature = "rackscale")]
+use crate::{arch::kcb::per_core_mem, memory::shmem_affinity::mid_to_shmem_affinity};
+
 /// The tokens per core to access the process replicas.
 #[thread_local]
 pub(crate) static PROCESS_TOKEN: Once<ArrayVec<ReplicaToken, { MAX_PROCESSES }>> = Once::new();
@@ -426,6 +429,7 @@ impl<P: Process> NrProcess<P> {
         }
     }
 
+    // For rackscale, kbuf must be allocated in shared memory
     pub(crate) fn write_to_userspace(to: &mut UserSlice, kbuf: &[u8]) -> Result<(), KError> {
         let node = *crate::environment::NODE_ID;
         let pid = to.pid;
@@ -441,6 +445,7 @@ impl<P: Process> NrProcess<P> {
         }
     }
 
+    // For rackscale, everything in the box must be allocated in shared memory
     #[cfg(feature = "rackscale")]
     pub(crate) fn userspace_exec_slice_mut(
         on: UserSlice,
@@ -459,6 +464,7 @@ impl<P: Process> NrProcess<P> {
         }
     }
 
+    // For rackscale, everything in the box must be allocated in shared memory
     pub(crate) fn userspace_exec_slice<'a>(
         on: &'a UserSlice,
         f: Box<dyn Fn(&'a [u8]) -> KResult<()>>,
@@ -495,14 +501,36 @@ where
                 Ok(ProcessResult::Resolved(paddr, rights))
             }
             ProcessOp::ReadSlice(uslice) => {
+                #[cfg(feature = "rackscale")]
+                let affinity = {
+                    // We want to allocate the kernel arc in shared memory
+                    let pcm = per_core_mem();
+                    let affinity = pcm.physical_memory.borrow().affinity;
+                    // Use local shared affinity for now
+                    pcm.set_mem_affinity(mid_to_shmem_affinity(*crate::environment::MACHINE_ID))
+                        .expect("Can't change affinity");
+                    affinity
+                };
+
                 // We're going to copy what we read into this thing
                 // TODO(panic+oom): need `try_new_uninit_slice` https://github.com/rust-lang/rust/issues/63291
                 let mut buffer = Arc::<[u8]>::new_uninit_slice(uslice.len());
                 let data = Arc::get_mut(&mut buffer).unwrap();
-                uslice.with_slice(&*self.process, |ubuf| {
+                let ret = uslice.with_slice(&*self.process, |ubuf| {
                     MaybeUninit::write_slice(data, ubuf);
                     Ok(())
-                })?;
+                });
+
+                #[cfg(feature = "rackscale")]
+                {
+                    // Restore affinity
+                    let pcm = per_core_mem();
+                    pcm.set_mem_affinity(affinity)
+                        .expect("Can't change affinity");
+                };
+                // restore memory affinity before exiting
+                ret?;
+
                 let buffer = unsafe {
                     // Safety: `assume_init`
                     // - Plain-old-data, that we just copied into `buffer` above
@@ -514,12 +542,42 @@ where
                 Ok(ProcessResult::ReadSlice(buffer))
             }
             ProcessOp::ReadString(uslice) => {
+                #[cfg(feature = "rackscale")]
+                let affinity = {
+                    // We want to allocate the kernel arc in shared memory
+                    let pcm = per_core_mem();
+                    let affinity = pcm.physical_memory.borrow().affinity;
+                    // Use local shared affinity for now
+                    pcm.set_mem_affinity(mid_to_shmem_affinity(*crate::environment::MACHINE_ID))
+                        .expect("Can't change affinity");
+                    affinity
+                };
+
                 let mut kbuf = Vec::try_with_capacity(uslice.len())?;
-                uslice.with_slice(&*self.process, |ubuf| {
+                let ret1 = uslice.with_slice(&*self.process, |ubuf| {
                     kbuf.extend_from_slice(ubuf);
                     Ok(())
-                })?;
-                Ok(ProcessResult::ReadString(String::from_utf8(kbuf)?))
+                });
+
+                // This deny is needed so it compiles without rackscale.
+                #[deny(clippy::let_and_return)]
+                let ret = match ret1 {
+                    Ok(()) => match String::from_utf8(kbuf) {
+                        Ok(str) => Ok(ProcessResult::ReadString(str)),
+                        Err(e) => Err(KError::from(e)),
+                    },
+                    Err(e) => Err(e),
+                };
+
+                #[cfg(feature = "rackscale")]
+                {
+                    // Restore affinity before exiting even in error case
+                    let pcm = per_core_mem();
+                    pcm.set_mem_affinity(affinity)
+                        .expect("Can't change affinity");
+                };
+
+                ret
             }
             ProcessOp::WriteSlice(uslice, kbuf) => {
                 if uslice.len() != kbuf.len() {

kernel/src/process.rs

@@ -38,7 +38,9 @@ use crate::fs::cnrfs;
 #[cfg(all(feature = "rackscale", target_arch = "x86_64"))]
 use {
     crate::arch::rackscale::get_shmem_frames::rpc_get_shmem_frames,
-    crate::memory::shmem_affinity::{is_shmem_affinity, local_shmem_affinity},
+    crate::memory::shmem_affinity::{
+        is_shmem_affinity, local_shmem_affinity, mid_to_shmem_affinity,
+    },
 };
 
 /// Process ID.
@@ -754,7 +756,8 @@ impl SliceAccess for &mut [u8] {
 /// while we copy the data in the file (and hence end up with inconsistent
 /// replicas).
 ///
-/// e.g., Any buffer that goes in the NR/CNR logs should be [`KernArcBuffer`].
+/// e.g., Any buffer that goes in the NR/CNR logs should be [`KernArcBuffer`],
+/// and for rackscale, [`KernArcBuffer`] must be allocated in shared memory.
 #[derive(PartialEq, Clone, Debug)]
 pub(crate) struct KernArcBuffer {
     pub buffer: Arc<[u8]>,
@@ -775,11 +778,30 @@ impl TryFrom<&[u8]> for KernArcBuffer {
 
     /// Converts a user-slice to a kernel slice.
     fn try_from(slice: &[u8]) -> KResult<Self> {
+        #[cfg(feature = "rackscale")]
+        let affinity = {
+            // We want to allocate the kernel arc in shared memory
+            let pcm = per_core_mem();
+            let affinity = pcm.physical_memory.borrow().affinity;
+            // Use local shared affinity for now
+            pcm.set_mem_affinity(mid_to_shmem_affinity(*crate::environment::MACHINE_ID))
+                .expect("Can't change affinity");
+            affinity
+        };
+
         // TODO: Panics on OOM, need a `try_new_uninit_slice()` https://github.com/rust-lang/rust/issues/63291
         let mut buffer = Arc::<[u8]>::new_uninit_slice(slice.len());
         let data = Arc::get_mut(&mut buffer).unwrap(); // not shared yet, no panic!
         MaybeUninit::write_slice(data, slice);
 
+        #[cfg(feature = "rackscale")]
+        {
+            // Restore affinity
+            let pcm = per_core_mem();
+            pcm.set_mem_affinity(affinity)
+                .expect("Can't change affinity");
+        };
+
         let buffer = unsafe {
             // Safety:
             // - Length == slice.len(): see above
@@ -971,15 +993,18 @@ impl UserSlice {
 }
 
 impl SliceAccess for UserSlice {
+    // For rackscale, everything in the box must be in shared memory
     fn read_slice<'a>(&'a self, f: Box<dyn Fn(&'a [u8]) -> KResult<()>>) -> KResult<()> {
         nrproc::NrProcess::<ArchProcess>::userspace_exec_slice(self, f)
     }
 
+    // For rackscale, buffer must be in shared memory
     fn write_slice(&mut self, buffer: &[u8]) -> KResult<()> {
         nrproc::NrProcess::<ArchProcess>::write_to_userspace(self, buffer)?;
         Ok(())
     }
 
+    // For rackscale, the buffer must be in shared memory
     fn write_subslice(&mut self, buffer: &[u8], offset: usize) -> KResult<()> {
         if self.len() < (offset + buffer.len()) {
             return Err(KError::InvalidOffset);