refactor: merge run_on_thread_local into signal handler

Merge run_on_thread_local into signal handler since it is
only used there.
The error returned from run_on_thread_local was ignored, so
instead replace it with logic to only use vcpu ptr if
TLS is initialized, without returning any errors.
The reason for not asserting on TLS being initialized here
is that during Firecracker shutdown, vcpus will be destroyed
and TLS will be reset. If signal will be send to Firecracker
during that time, the TLS accessed from a signal handler will
be empty. But this is expected, so no assertions/panics are needed.

Because Rust is a good language, it does not allow to reference
TLS_VCPU_PTR definded inside impl block inside the signal_handler
function. So move the TLS_VCPU_PTR definition outside the Vcpu impl
block.

Signed-off-by: Egor Lazarchuk <[email protected]>

Author: ShadowCurse

src/vmm/src/vstate/vcpu.rs

@@ -88,6 +88,8 @@ pub enum CopyKvmFdError {
     CreateVcpuError(#[from] kvm_ioctls::Error),
 }
 
+thread_local!(static TLS_VCPU_PTR: VcpuCell = const { Cell::new(None) });
+
 /// A wrapper around creating and using a vcpu.
 #[derive(Debug)]
 pub struct Vcpu {
@@ -110,15 +112,13 @@ pub struct Vcpu {
 }
 
 impl Vcpu {
-    thread_local!(static TLS_VCPU_PTR: VcpuCell = const { Cell::new(None) });
-
     /// Associates `self` with the current thread.
     ///
     /// It is a prerequisite to successfully run `init_thread_local_data()` before using
     /// `run_on_thread_local()` on the current thread.
     /// This function will panic if there already is a `Vcpu` present in the TLS.
     fn init_thread_local_data(&mut self) {
-        Self::TLS_VCPU_PTR.with(|cell: &VcpuCell| {
+        TLS_VCPU_PTR.with(|cell: &VcpuCell| {
             assert!(cell.get().is_none());
             cell.set(Some(self as *mut Vcpu));
         })
@@ -130,39 +130,12 @@ impl Vcpu {
     fn reset_thread_local_data(&mut self) {
         // Best-effort to clean up TLS. If the `Vcpu` was moved to another thread
         // _before_ running this, then there is nothing we can do.
-        Self::TLS_VCPU_PTR.with(|cell| {
+        TLS_VCPU_PTR.with(|cell| {
             let vcpu_ptr = cell.get().unwrap();
             assert!(std::ptr::eq(vcpu_ptr, self));
             // Have to do this trick because `update` method is
             // not stable on Cell.
-            Self::TLS_VCPU_PTR.with(|cell| cell.take());
-        })
-    }
-
-    /// Runs `func` for the `Vcpu` associated with the current thread.
-    ///
-    /// It requires that `init_thread_local_data()` was run on this thread.
-    ///
-    /// Fails if there is no `Vcpu` associated with the current thread.
-    ///
-    /// # Safety
-    ///
-    /// This is marked unsafe as it allows temporary aliasing through
-    /// dereferencing from pointer an already borrowed `Vcpu`.
-    unsafe fn run_on_thread_local<F>(func: F) -> Result<(), VcpuError>
-    where
-        F: FnOnce(&mut Vcpu),
-    {
-        Self::TLS_VCPU_PTR.with(|cell: &VcpuCell| {
-            if let Some(vcpu_ptr) = cell.get() {
-                // SAFETY: Dereferencing here is safe since `TLS_VCPU_PTR` is populated/non-empty,
-                // and it is being cleared on `Vcpu::drop` so there is no dangling pointer.
-                let vcpu_ref = unsafe { &mut *vcpu_ptr };
-                func(vcpu_ref);
-                Ok(())
-            } else {
-                Err(VcpuError::VcpuTlsNotPresent)
-            }
+            TLS_VCPU_PTR.with(|cell| cell.take());
         })
     }
 
@@ -172,14 +145,16 @@ impl Vcpu {
         self.init_thread_local_data();
 
         extern "C" fn handle_signal(_: c_int, _: *mut siginfo_t, _: *mut c_void) {
-            // SAFETY: This is safe because it's temporarily aliasing the `Vcpu` object, but we are
-            // only reading `vcpu.fd` which does not change for the lifetime of the `Vcpu`.
-            unsafe {
-                let _ = Vcpu::run_on_thread_local(|vcpu| {
+            TLS_VCPU_PTR.with(|cell: &VcpuCell| {
+                if let Some(vcpu_ptr) = cell.get() {
+                    // SAFETY: Dereferencing here is safe since `TLS_VCPU_PTR` is populated/non-empty,
+                    // and it is being cleared on `Vcpu::drop` so there is no dangling pointer.
+                    let vcpu = unsafe { &mut *vcpu_ptr };
+
                     vcpu.kvm_vcpu.fd.set_kvm_immediate_exit(1);
                     fence(Ordering::Release);
-                });
-            }
+                }
+            })
         }
 
         register_signal_handler(sigrtmin() + VCPU_RTSIG_OFFSET, handle_signal)
@@ -1014,34 +989,6 @@ pub(crate) mod tests {
         assert!(vcpu.kvm_vcpu.peripherals.mmio_bus.is_some());
     }
 
-    #[test]
-    fn test_vcpu_tls() {
-        let (_, _, mut vcpu) = setup_vcpu(0x1000);
-
-        // Running on the TLS vcpu should fail before we actually initialize it.
-        unsafe {
-            Vcpu::run_on_thread_local(|_| ()).unwrap_err();
-        }
-
-        // Initialize vcpu TLS.
-        vcpu.init_thread_local_data();
-
-        // Validate TLS vcpu is the local vcpu by changing the `id` then validating against
-        // the one in TLS.
-        vcpu.kvm_vcpu.index = 12;
-        unsafe {
-            Vcpu::run_on_thread_local(|v| assert_eq!(v.kvm_vcpu.index, 12)).unwrap();
-        }
-
-        // Reset vcpu TLS.
-        vcpu.reset_thread_local_data();
-
-        // Running on the TLS vcpu after TLS reset should fail.
-        unsafe {
-            Vcpu::run_on_thread_local(|_| ()).unwrap_err();
-        }
-    }
-
     #[test]
     #[should_panic]
     fn test_tls_double_init() {