From 4a56fcfd6bafe4b17f91c1f7ebb7d1dc00ea4274 Mon Sep 17 00:00:00 2001
From: danbugs <danilochiarlone@gmail.com>
Date: Wed, 6 Nov 2024 23:33:57 +0000
Subject: [PATCH] [performance] Move set_up_hypervisor_partition to inside the
 hv handler thread

As per KVM docs, vCPU ioctls should be issued from the same thread that was used to create the vCPU.
In accordance to that, this PR moves set_up_hypervisor_partition to inside the hv handler thread,
which manages vCPU interactions like init and dispatching guest funciton calls.

Signed-off-by: danbugs <danilochiarlone@gmail.com>
---
 .../src/hypervisor/hypervisor_handler.rs      | 50 ++++++++++++-------
 1 file changed, 31 insertions(+), 19 deletions(-)

diff --git a/src/hyperlight_host/src/hypervisor/hypervisor_handler.rs b/src/hyperlight_host/src/hypervisor/hypervisor_handler.rs
index bfb55fa2f..2724ab6ac 100644
--- a/src/hyperlight_host/src/hypervisor/hypervisor_handler.rs
+++ b/src/hyperlight_host/src/hypervisor/hypervisor_handler.rs
@@ -16,6 +16,7 @@ limitations under the License.
 
 #[cfg(target_os = "windows")]
 use core::ffi::c_void;
+use std::ops::DerefMut;
 use std::sync::atomic::{AtomicBool, Ordering};
 use std::sync::{Arc, Mutex};
 use std::thread;
@@ -125,7 +126,7 @@ impl HvHandlerExecVars {
             .thread_id
             .try_lock()
             .map_err(|_| new_error!("Failed to get_thread_id"))?)
-        .ok_or_else(|| new_error!("thread_id not set"))
+            .ok_or_else(|| new_error!("thread_id not set"))
     }
 
     #[cfg(target_os = "windows")]
@@ -228,15 +229,11 @@ impl HypervisorHandler {
     #[instrument(err(Debug), skip_all, parent = Span::current(), level = "Trace")]
     pub(crate) fn start_hypervisor_handler(
         &mut self,
-        mut sandbox_memory_manager: SandboxMemoryManager<GuestSharedMemory>,
+        sandbox_memory_manager: SandboxMemoryManager<GuestSharedMemory>,
     ) -> Result<()> {
         let configuration = self.configuration.clone();
-        let mut hv = set_up_hypervisor_partition(
-            &mut sandbox_memory_manager,
-            configuration.outb_handler.clone(),
-        )?;
         #[cfg(target_os = "windows")]
-        let in_process = sandbox_memory_manager.is_in_process();
+            let in_process = sandbox_memory_manager.is_in_process();
 
         *self.execution_variables.shm.try_lock().unwrap() = Some(sandbox_memory_manager);
 
@@ -267,18 +264,8 @@ impl HypervisorHandler {
         #[cfg(target_os = "linux")]
         self.execution_variables.run_cancelled.store(false);
 
-        #[cfg(target_os = "windows")]
-        if !in_process {
-            self.execution_variables
-                .set_partition_handle(hv.get_partition_handle())?;
-        }
-
         let to_handler_rx = self.communication_channels.to_handler_rx.clone();
-        #[cfg(target_os = "windows")]
-        let execution_variables = self.execution_variables.clone();
-        #[cfg(target_os = "linux")]
         let mut execution_variables = self.execution_variables.clone();
-        // ^^^ this needs to be mut on linux to set_thread_id
         let from_handler_tx = self.communication_channels.from_handler_tx.clone();
         let hv_handler_clone = self.clone();
 
@@ -295,9 +282,24 @@ impl HypervisorHandler {
             thread::Builder::new()
                 .name("Hypervisor Handler".to_string())
                 .spawn(move || -> Result<()> {
+                    let mut hv: Option<Box<dyn Hypervisor>> = None;
                     for action in to_handler_rx {
                         match action {
                             HypervisorHandlerAction::Initialise => {
+                                {
+                                    hv = Some(set_up_hypervisor_partition(
+                                        execution_variables.shm.try_lock().unwrap().deref_mut().as_mut().unwrap(),
+                                        configuration.outb_handler.clone(),
+                                    )?);
+                                }
+                                let hv = hv.as_mut().unwrap();
+
+                                #[cfg(target_os = "windows")]
+                                if !in_process {
+                                    execution_variables
+                                        .set_partition_handle(hv.get_partition_handle())?;
+                                }
+
                                 #[cfg(target_os = "linux")]
                                 {
                                     // We cannot use the Killable trait, so we get the `pthread_t` via a libc
@@ -328,6 +330,7 @@ impl HypervisorHandler {
                                     .shared_mem
                                     .lock
                                     .try_read();
+
                                 let res = hv.initialise(
                                     configuration.peb_addr.clone(),
                                     configuration.seed,
@@ -362,6 +365,8 @@ impl HypervisorHandler {
                                 }
                             }
                             HypervisorHandlerAction::DispatchCallFromHost(function_name) => {
+                                let hv = hv.as_mut().unwrap();
+
                                 // Lock to indicate an action is being performed in the hypervisor
                                 execution_variables.running.store(true, Ordering::SeqCst);
 
@@ -735,7 +740,7 @@ impl HypervisorHandler {
                         0,
                         0,
                     )
-                    .map_err(|e| new_error!("Failed to cancel guest execution {:?}", e))?;
+                        .map_err(|e| new_error!("Failed to cancel guest execution {:?}", e))?;
                 }
             }
             // if running in-process on windows, we currently have no way of cancelling the execution
@@ -911,7 +916,7 @@ mod tests {
             None,
             None,
         )
-        .unwrap();
+            .unwrap();
 
         usbox.evolve(Noop::default()).unwrap()
     }
@@ -944,6 +949,13 @@ mod tests {
         }
     }
 
+    #[test]
+    fn create_10_sandboxes() {
+        for _ in 0..10 {
+            create_multi_use_sandbox();
+        }
+    }
+
     #[test]
     fn hello_world() -> Result<()> {
         let mut sandbox = create_multi_use_sandbox();