refactor: assert on TLS reset

The reset must be called only once on vcpu drop. Replace
errors with asserts in this case.

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

Author: ShadowCurse

src/vmm/src/vstate/vcpu.rs

@@ -126,21 +126,16 @@ 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> {
+    /// Must be called after `init_thread_local_data`.
+    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: &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)
+        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());
         })
     }
 
@@ -615,7 +610,7 @@ fn handle_kvm_exit(
 
 impl Drop for Vcpu {
     fn drop(&mut self) {
-        let _ = self.reset_thread_local_data();
+        self.reset_thread_local_data();
     }
 }
 
@@ -1039,15 +1034,24 @@ 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();
+        vcpu.reset_thread_local_data();
+    }
+
+    #[test]
+    #[should_panic]
+    fn test_vcpu_tls_double_reset() {
+        let (_, _, mut vcpu) = setup_vcpu(0x1000);
+        vcpu.init_thread_local_data();
+        vcpu.reset_thread_local_data();
+        vcpu.reset_thread_local_data();
     }
 
     #[test]