From db28db85678a96f6ac5cca981315a1508d2577e8 Mon Sep 17 00:00:00 2001
From: Sebastien Boeuf <sebastien.boeuf@intel.com>
Date: Fri, 28 Aug 2020 19:18:33 +0200
Subject: [PATCH] vmm: Update NUMA nodes based on NumaConfig

Relying on the list of CPUs defined through the NumaConfig, this patch
will update the internal list of CPUs attached to each NUMA node.

Signed-off-by: Sebastien Boeuf <sebastien.boeuf@intel.com>
---
 vmm/src/memory_manager.rs | 14 ++++++++++++++
 vmm/src/vm.rs             | 27 +++++++++++++++++++++++++--
 2 files changed, 39 insertions(+), 2 deletions(-)

diff --git a/vmm/src/memory_manager.rs b/vmm/src/memory_manager.rs
index e1de6bfeb..906960fde 100644
--- a/vmm/src/memory_manager.rs
+++ b/vmm/src/memory_manager.rs
@@ -66,12 +66,21 @@ struct HotPlugState {
 #[derive(Clone)]
 pub struct NumaNode {
     memory_regions: Vec<Arc<GuestRegionMmap>>,
+    cpus: Vec<u8>,
 }
 
 impl NumaNode {
     pub fn memory_regions(&self) -> &Vec<Arc<GuestRegionMmap>> {
         &self.memory_regions
     }
+
+    pub fn cpus(&self) -> &Vec<u8> {
+        &self.cpus
+    }
+
+    pub fn cpus_mut(&mut self) -> &mut Vec<u8> {
+        &mut self.cpus
+    }
 }
 
 pub type NumaNodes = BTreeMap<u32, NumaNode>;
@@ -356,6 +365,7 @@ impl MemoryManager {
                             node_id,
                             NumaNode {
                                 memory_regions: vec![region.clone()],
+                                cpus: Vec::new(),
                             },
                         );
                     }
@@ -1233,6 +1243,10 @@ impl MemoryManager {
     pub fn numa_nodes(&self) -> &NumaNodes {
         &self.numa_nodes
     }
+
+    pub fn numa_nodes_mut(&mut self) -> &mut NumaNodes {
+        &mut self.numa_nodes
+    }
 }
 
 #[cfg(feature = "acpi")]
diff --git a/vmm/src/vm.rs b/vmm/src/vm.rs
index 16fb30b22..bf5d8dcc5 100644
--- a/vmm/src/vm.rs
+++ b/vmm/src/vm.rs
@@ -24,8 +24,8 @@ extern crate vm_allocator;
 extern crate vm_memory;
 
 use crate::config::{
-    DeviceConfig, DiskConfig, FsConfig, HotplugMethod, NetConfig, PmemConfig, ValidationError,
-    VmConfig, VsockConfig,
+    DeviceConfig, DiskConfig, FsConfig, HotplugMethod, NetConfig, NumaConfig, PmemConfig,
+    ValidationError, VmConfig, VsockConfig,
 };
 use crate::cpu;
 use crate::device_manager::{self, get_win_size, Console, DeviceManager, DeviceManagerError};
@@ -311,6 +311,11 @@ impl Vm {
             .transpose()
             .map_err(Error::InitramfsFile)?;
 
+        // Update NUMA based on NumaConfig.
+        if let Some(numa_cfg) = config.lock().unwrap().numa.clone() {
+            Self::update_numa(numa_cfg, &memory_manager)?;
+        }
+
         Ok(Vm {
             kernel,
             initramfs,
@@ -328,6 +333,24 @@ impl Vm {
         })
     }
 
+    fn update_numa(
+        configs: Vec<NumaConfig>,
+        memory_manager: &Arc<Mutex<MemoryManager>>,
+    ) -> Result<()> {
+        let mut mm = memory_manager.lock().unwrap();
+        let numa_nodes = mm.numa_nodes_mut();
+
+        for config in configs.iter() {
+            if let Some(cpus) = &config.cpus {
+                if let Some(node) = numa_nodes.get_mut(&config.id) {
+                    node.cpus_mut().extend(cpus);
+                }
+            }
+        }
+
+        Ok(())
+    }
+
     pub fn new(
         config: Arc<Mutex<VmConfig>>,
         exit_evt: EventFd,