From d8dc153be976263314a21b15a7d89ea0c6f9f697 Mon Sep 17 00:00:00 2001
From: Qingsong Chen <changxian.cqs@antgroup.com>
Date: Fri, 10 Jan 2025 02:44:40 +0000
Subject: [PATCH] Fix deadlock of `CurrentTx` in mlsdisk

---
 kernel/comps/mlsdisk/src/tx/current.rs | 43 +++++++++++++++++++-------
 kernel/comps/mlsdisk/src/tx/mod.rs     |  2 +-
 2 files changed, 32 insertions(+), 13 deletions(-)

diff --git a/kernel/comps/mlsdisk/src/tx/current.rs b/kernel/comps/mlsdisk/src/tx/current.rs
index af8db39fa..856c6bc48 100644
--- a/kernel/comps/mlsdisk/src/tx/current.rs
+++ b/kernel/comps/mlsdisk/src/tx/current.rs
@@ -46,33 +46,52 @@ impl<'a> CurrentTx<'a> {
     /// If the returned value is `Ok`, then the TX is committed successfully.
     /// Otherwise, the TX is aborted.
     pub fn commit(&self) -> Result<()> {
-        let mut tx_table = self.provider.tx_table.lock();
-        let Some(mut tx) = tx_table.remove(&CurrentThread::id()) else {
-            panic!("there should be one Tx exited on the current thread");
-        };
-        debug_assert!(tx.status() == TxStatus::Ongoing);
+        let mut tx_status = self
+            .provider
+            .tx_table
+            .lock()
+            .get(&CurrentThread::id())
+            .expect("there should be one Tx exited on the current thread")
+            .status();
+        debug_assert!(tx_status == TxStatus::Ongoing);
 
         let res = self.provider.call_precommit_handlers();
         if res.is_ok() {
             self.provider.call_commit_handlers();
-            tx.set_status(TxStatus::Committed);
+            tx_status = TxStatus::Committed;
         } else {
             self.provider.call_abort_handlers();
-            tx.set_status(TxStatus::Aborted);
+            tx_status = TxStatus::Aborted;
         }
 
+        let mut tx = self
+            .provider
+            .tx_table
+            .lock()
+            .remove(&CurrentThread::id())
+            .unwrap();
+        tx.set_status(tx_status);
         res
     }
 
     /// Aborts the current TX.
     pub fn abort(&self) {
-        let mut tx_table = self.provider.tx_table.lock();
-        let Some(mut tx) = tx_table.remove(&CurrentThread::id()) else {
-            panic!("there should be one Tx exited on the current thread");
-        };
-        debug_assert!(tx.status() == TxStatus::Ongoing);
+        let tx_status = self
+            .provider
+            .tx_table
+            .lock()
+            .get(&CurrentThread::id())
+            .expect("there should be one Tx exited on the current thread")
+            .status();
+        debug_assert!(tx_status == TxStatus::Ongoing);
 
         self.provider.call_abort_handlers();
+        let mut tx = self
+            .provider
+            .tx_table
+            .lock()
+            .remove(&CurrentThread::id())
+            .unwrap();
         tx.set_status(TxStatus::Aborted);
     }
 
diff --git a/kernel/comps/mlsdisk/src/tx/mod.rs b/kernel/comps/mlsdisk/src/tx/mod.rs
index bbcafae88..8cc98e73a 100644
--- a/kernel/comps/mlsdisk/src/tx/mod.rs
+++ b/kernel/comps/mlsdisk/src/tx/mod.rs
@@ -90,7 +90,7 @@ impl TxProvider {
         let init_fn = initializer_map
             .get(&TypeId::of::<T>())
             .unwrap()
-            .downcast_ref::<Box<dyn Fn() -> T>>()
+            .downcast_ref::<Box<dyn Fn() -> T + Send + Sync>>()
             .unwrap();
         init_fn()
     }