From 9548e7e8570c37b31a431291e4302d5af1174595 Mon Sep 17 00:00:00 2001
From: Sebastien Boeuf <sebastien.boeuf@intel.com>
Date: Mon, 31 Aug 2020 15:05:49 +0200
Subject: [PATCH] vmm: Update NUMA node distances internally

Based on the NumaConfig which now provides distance information, we can
internally update the list of NUMA nodes with the exact distances they
should be located from other nodes.

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

diff --git a/vmm/src/memory_manager.rs b/vmm/src/memory_manager.rs
index 906960fde..e09425078 100644
--- a/vmm/src/memory_manager.rs
+++ b/vmm/src/memory_manager.rs
@@ -67,6 +67,7 @@ struct HotPlugState {
 pub struct NumaNode {
     memory_regions: Vec<Arc<GuestRegionMmap>>,
     cpus: Vec<u8>,
+    distances: BTreeMap<u32, u8>,
 }
 
 impl NumaNode {
@@ -81,6 +82,14 @@ impl NumaNode {
     pub fn cpus_mut(&mut self) -> &mut Vec<u8> {
         &mut self.cpus
     }
+
+    pub fn distances(&self) -> &BTreeMap<u32, u8> {
+        &self.distances
+    }
+
+    pub fn distances_mut(&mut self) -> &mut BTreeMap<u32, u8> {
+        &mut self.distances
+    }
 }
 
 pub type NumaNodes = BTreeMap<u32, NumaNode>;
@@ -366,6 +375,7 @@ impl MemoryManager {
                             NumaNode {
                                 memory_regions: vec![region.clone()],
                                 cpus: Vec::new(),
+                                distances: BTreeMap::new(),
                             },
                         );
                     }
diff --git a/vmm/src/vm.rs b/vmm/src/vm.rs
index bf5d8dcc5..9cd7dc9e4 100644
--- a/vmm/src/vm.rs
+++ b/vmm/src/vm.rs
@@ -195,6 +195,9 @@ pub enum Error {
 
     /// Failed serializing into JSON
     SerializeJson(serde_json::Error),
+
+    /// Invalid configuration for NUMA.
+    InvalidNumaConfig,
 }
 pub type Result<T> = result::Result<T, Error>;
 
@@ -339,12 +342,30 @@ impl Vm {
     ) -> Result<()> {
         let mut mm = memory_manager.lock().unwrap();
         let numa_nodes = mm.numa_nodes_mut();
+        let existing_nodes: Vec<u32> = numa_nodes.keys().cloned().collect();
 
         for config in configs.iter() {
-            if let Some(cpus) = &config.cpus {
-                if let Some(node) = numa_nodes.get_mut(&config.id) {
+            if let Some(node) = numa_nodes.get_mut(&config.id) {
+                if let Some(cpus) = &config.cpus {
                     node.cpus_mut().extend(cpus);
                 }
+
+                if let Some(distances) = &config.distances {
+                    for distance in distances.iter() {
+                        let dest = distance.destination;
+                        let dist = distance.distance;
+
+                        if !existing_nodes.contains(&dest) {
+                            error!("Unknown destination NUMA node {}", dest);
+                            return Err(Error::InvalidNumaConfig);
+                        }
+
+                        node.distances_mut().insert(dest, dist);
+                    }
+                }
+            } else {
+                error!("Unknown NUMA node {}", config.id);
+                return Err(Error::InvalidNumaConfig);
             }
         }