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 <yegorlz@amazon.co.uk>

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,61 +112,35 @@ 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));
         })
     }
 
-    /// 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)
-            }
-        })
-    }
-
     /// Registers a signal handler which makes use of TLS and kvm immediate exit to
     /// kick the vcpu running on the current thread, if there is one.
     fn register_kick_signal_handler(&mut self) {
         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)
@@ -595,7 +571,7 @@ fn handle_kvm_exit(
 
 impl Drop for Vcpu {
     fn drop(&mut self) {
-        Self::TLS_VCPU_PTR.with(|cell| {
+        TLS_VCPU_PTR.with(|cell| {
             // The reason for not asserting TLS being set here is that
             // it can happen that Vcpu::Drop is called on vcpus which never were
             // put on their threads. This can happen if some error occurs during Vmm
@@ -611,7 +587,7 @@ impl Drop for Vcpu {
 
                 // Have to do this trick because `update` method is
                 // not stable on Cell.
-                Self::TLS_VCPU_PTR.with(|cell| cell.take());
+                TLS_VCPU_PTR.with(|cell| cell.take());
             }
         })
     }
@@ -1017,34 +993,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::TLS_VCPU_PTR.with(|cell| cell.take());
-
-        // 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() {