Explicitly set memory affinity during replica allocation

usr/init/src/dynrep/allocator.rs

@@ -1,10 +1,11 @@
+use alloc::sync::Arc;
 use core::slice::from_raw_parts_mut;
-use x86::bits64::paging::LARGE_PAGE_SIZE;
 use lazy_static::lazy_static;
-use core::sync::atomic::AtomicU8;
-use core::sync::atomic::Ordering;
 
-use log::info;
+use kpi::system::NodeId;
+use spin::Mutex;
+use x86::bits64::paging::LARGE_PAGE_SIZE;
+
 use core::{
     alloc::{AllocError, Allocator, Layout},
     ptr::NonNull,
@@ -13,50 +14,70 @@ use core::{
 pub const BASE: u64 = 0x0510_0000_0000;
 pub const MAX_FRAMES: u64 = 600;
 
-static NODE_ID: AtomicU8 = AtomicU8::new(1);
+lazy_static! {
+    pub(crate) static ref ALLOC_AFFINITY: Arc<Mutex<NodeId>> = Arc::new(Mutex::new(0));
+}
 
 #[derive(Clone, Copy)]
 pub struct MyAllocator;
 
 impl MyAllocator {
-    fn allocate_pages(node_id: u8) {
+    fn allocate_pages(node_id: NodeId) {
         let mut allocated = 0;
         let node_offset = (node_id - 1) as u64 * LARGE_PAGE_SIZE as u64 * MAX_FRAMES;
         while allocated < MAX_FRAMES {
             // Allocate a large page of physical memory
             // Note that even if you allocate a base page, behind the scenes a large page is allocated
             // because DCM (and thus DiNOS) only allocates at large page granularity
             // 1 is the client machine id we want to allocate from
-            let (frame_id, paddr) = vibrio::syscalls::PhysicalMemory::allocate_large_page(node_id as usize)
-                .expect("Failed to get physical memory large page");
-            info!("large frame id={:?}, paddr={:?}", frame_id, paddr);
+            let (frame_id, paddr) =
+                vibrio::syscalls::PhysicalMemory::allocate_large_page(node_id as usize)
+                    .expect("Failed to get physical memory large page");
+            log::trace!(
+                "large frame id={:?}, paddr={:?}, mid={:?}",
+                frame_id,
+                paddr,
+                node_id
+            );
 
             // Map allocated physical memory into user space so we can actually access it.
             unsafe {
-                vibrio::syscalls::VSpace::map_frame(frame_id, BASE + node_offset + (allocated * LARGE_PAGE_SIZE as u64)).expect("Failed to map base page");
+                vibrio::syscalls::VSpace::map_frame(
+                    frame_id,
+                    BASE + node_offset + (allocated * LARGE_PAGE_SIZE as u64),
+                )
+                .expect("Failed to map base page");
             }
             allocated += 1;
         }
-        info!("# Allocated {} frames on {}", allocated, node_id);
     }
 }
 
 unsafe impl Allocator for MyAllocator {
     fn allocate(&self, layout: Layout) -> Result<NonNull<[u8]>, AllocError> {
-        let node_id = NODE_ID.fetch_add(1, Ordering::SeqCst);
-        let node_offset = (node_id - 1) as u64 * LARGE_PAGE_SIZE as u64 * MAX_FRAMES;
-        MyAllocator::allocate_pages(node_id);
-        info!("# Allocating {:?}", layout);
+        // Check argument
+        log::info!("{:?}", layout);
         if layout.size() > LARGE_PAGE_SIZE * MAX_FRAMES as usize {
             return Err(AllocError);
         }
 
-        let slice = unsafe {from_raw_parts_mut((BASE + node_offset) as *mut u8, layout.size()) };
+        // Grab affinity lock
+        let affinity = (*ALLOC_AFFINITY).lock();
+        log::info!("Affinity for alloc is: {:?}", *affinity + 1);
+
+        // Allocate and map pages
+        MyAllocator::allocate_pages(*affinity + 1);
+
+        // Get ptr to mapped memory
+        let node_offset = *affinity as u64 * LARGE_PAGE_SIZE as u64 * MAX_FRAMES;
+        let slice = unsafe { from_raw_parts_mut((BASE + node_offset) as *mut u8, layout.size()) };
+
+        log::info!("Finished allocating on {:?}", *affinity + 1);
         Ok(NonNull::from(slice))
     }
 
     unsafe fn deallocate(&self, ptr: NonNull<u8>, layout: Layout) {
-        info!("# Deallocating {:?}", layout);
+        log::info!("Deallocating {:?}", layout);
         /*for i in 0..MAX_FRAMES {
             vibrio::syscalls::VSpace::unmap((BASE + (i * LARGE_PAGE_SIZE as u64)) as u64, LARGE_PAGE_SIZE as u64)
                 .expect("Failed to unmap base page");

usr/init/src/dynrep/mod.rs

@@ -12,13 +12,13 @@ use core::sync::atomic::AtomicUsize;
 use core::sync::atomic::Ordering;
 
 use lineup::tls2::{Environment, SchedulerControlBlock};
-use nr2::nr::{Dispatch, NodeReplicated};
+use nr2::nr::{AffinityChange, Dispatch, NodeReplicated};
+use rawtime::Instant;
 use x86::bits64::paging::VAddr;
 use x86::random::rdrand64;
-use rawtime::Instant;
 
 mod allocator;
-use allocator::MyAllocator;
+use allocator::{MyAllocator, ALLOC_AFFINITY};
 
 pub const NUM_ENTRIES: u64 = 50_000_000;
 
@@ -31,7 +31,7 @@ struct HashTable {
 
 impl Default for HashTable {
     fn default() -> Self {
-        let allocator = MyAllocator{};
+        let allocator = MyAllocator {};
         let map = HashMap::<u64, u64, DefaultHashBuilder, MyAllocator>::with_capacity_in(
             NUM_ENTRIES as usize,
             allocator,
@@ -73,9 +73,9 @@ impl Dispatch for HashTable {
     }
 }
 
-fn run_bench(mid: usize, core_id : usize, replica: Arc<NodeReplicated<HashTable>>) {
+fn run_bench(mid: usize, core_id: usize, replica: Arc<NodeReplicated<HashTable>>) {
     let ttkn = replica.register(mid - 1).unwrap();
-    let mut random_key :u64 = 0;
+    let mut random_key: u64 = 0;
     let batch_size = 64;
     let duration = 5;
 
@@ -84,24 +84,23 @@ fn run_bench(mid: usize, core_id : usize, replica: Arc<NodeReplicated<HashTable>
         let mut ops = 0;
         let start = Instant::now();
         while start.elapsed().as_secs() < 1 {
-        for i in 0..batch_size {
+            for i in 0..batch_size {
                 unsafe { rdrand64(&mut random_key) };
                 random_key = random_key % NUM_ENTRIES;
                 let _ = replica.execute(OpRd::Get(random_key), ttkn).unwrap();
                 ops += 1;
+            }
         }
-        }
-        info!(
-            "dynhash,{},{},{},{}",mid,core_id, iterations, ops
-        );
+        info!("dynhash,{},{},{},{}", mid, core_id, iterations, ops);
         iterations += 1;
     }
 }
 
 unsafe extern "C" fn bencher_trampoline(args: *mut u8) -> *mut u8 {
     let current_gtid = vibrio::syscalls::System::core_id().expect("Can't get core id");
     let mid = kpi::system::mid_from_gtid(current_gtid);
-    let replica: Arc<NodeReplicated<HashTable>> = Arc::from_raw(args as *const NodeReplicated<HashTable>);
+    let replica: Arc<NodeReplicated<HashTable>> =
+        Arc::from_raw(args as *const NodeReplicated<HashTable>);
 
     // Synchronize with all cores
     POOR_MANS_BARRIER.fetch_sub(1, Ordering::Release);
@@ -131,9 +130,29 @@ pub fn userspace_dynrep_test() {
     log::info!("Found {:?} client machines", nnodes);
 
     // Create data structure, with as many replicas as there are clients (assuming 1 numa node per client)
-    // TODO: change # of replicas to nnodes
     let num_replicas = NonZeroUsize::new(nnodes).unwrap();
-    let replicas = Arc::new(NodeReplicated::<HashTable>::new(num_replicas, |_| 0).unwrap());
+    let replicas = Arc::new(
+        NodeReplicated::<HashTable>::new(num_replicas, |afc: AffinityChange| {
+            log::trace!("Got AffinityChange: {:?}", afc);
+            match afc {
+                AffinityChange::Replica(r) => {
+                    let mut affinity = (*ALLOC_AFFINITY).lock();
+                    let old_affinity = *affinity;
+                    *affinity = r;
+                    log::info!("Set alloc affinity to {:?}", r);
+                    return old_affinity;
+                }
+                AffinityChange::Revert(orig) => {
+                    //pcm.set_mem_affinity(orig).expect("Can't set affinity");
+                    let mut affinity = (*ALLOC_AFFINITY).lock();
+                    *affinity = orig;
+                    log::info!("Restored alloc affinity to {:?}", orig);
+                    return 0;
+                }
+            }
+        })
+        .unwrap(),
+    );
 
     let s = &vibrio::upcalls::PROCESS_SCHEDULER;
     let mut gtids = Vec::with_capacity(ncores);