refactor: move TLS reset into Vcpu::Drop

The reset must be called only once on vcpu drop, so directly
move it into the Drop impl. Tests can use a simplified version.
Replace errors with asserts since there is an assumption that TLS
will always hold correct vcpu pointer for a given thread.

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

Author: ShadowCurse

src/vmm/src/vstate/vcpu.rs

@@ -124,26 +124,6 @@ impl Vcpu {
         })
     }
 
-    /// Deassociates `self` from the current thread.
-    ///
-    /// Should be called if the current `self` had called `init_thread_local_data()` and
-    /// now needs to move to a different thread.
-    ///
-    /// Fails if `self` was not previously associated with the current thread.
-    fn reset_thread_local_data(&mut self) -> Result<(), VcpuError> {
-        // 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: &VcpuCell| {
-            if let Some(vcpu_ptr) = cell.get() {
-                if std::ptr::eq(vcpu_ptr, self) {
-                    Self::TLS_VCPU_PTR.with(|cell: &VcpuCell| cell.take());
-                    return Ok(());
-                }
-            }
-            Err(VcpuError::VcpuTlsNotPresent)
-        })
-    }
-
     /// Runs `func` for the `Vcpu` associated with the current thread.
     ///
     /// It requires that `init_thread_local_data()` was run on this thread.
@@ -615,7 +595,13 @@ fn handle_kvm_exit(
 
 impl Drop for Vcpu {
     fn drop(&mut self) {
-        let _ = self.reset_thread_local_data();
+        Self::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());
+        })
     }
 }
 
@@ -1011,6 +997,12 @@ pub(crate) mod tests {
         (vm, vcpu_handle, vcpu_exit_evt)
     }
 
+    impl Vcpu {
+        fn reset_thread_local_data(&mut self) {
+            Self::TLS_VCPU_PTR.with(|cell| cell.take());
+        }
+    }
+
     #[test]
     fn test_set_mmio_bus() {
         let (_, _, mut vcpu) = setup_vcpu(0x1000);
@@ -1039,15 +1031,12 @@ pub(crate) mod tests {
         }
 
         // Reset vcpu TLS.
-        vcpu.reset_thread_local_data().unwrap();
+        vcpu.reset_thread_local_data();
 
         // Running on the TLS vcpu after TLS reset should fail.
         unsafe {
             Vcpu::run_on_thread_local(|_| ()).unwrap_err();
         }
-
-        // Second reset should return error.
-        vcpu.reset_thread_local_data().unwrap_err();
     }
 
     #[test]
@@ -1058,6 +1047,17 @@ pub(crate) mod tests {
         vcpu.init_thread_local_data();
     }
 
+    #[test]
+    #[should_panic]
+    fn test_vcpu_tls_double_reset() {
+        let (_, _, mut vcpu) = setup_vcpu(0x1000);
+        vcpu.init_thread_local_data();
+        // This reset will succeed as the TLS
+        // was initialized
+        vcpu.reset_thread_local_data();
+        // The Drop in vcpu will panic now
+    }
+
     #[test]
     fn test_vcpu_kick() {
         let (_, vm, mut vcpu) = setup_vcpu(0x1000);