fix: atomic cas

bytedance · Aug 24, 2024 · c78d199 · c78d199
1 parent e7a6dc1
commit c78d199
Showing 1 changed file with 13 additions and 6 deletions.
diff --git a/monoio/src/task/state.rs b/monoio/src/task/state.rs
@@ -80,6 +80,7 @@ impl State {
         Snapshot(self.0.load(Acquire))
     }
 
+    #[allow(dead_code)]
     pub(crate) fn store(&self, val: Snapshot) {
         self.0.store(val.0, Release);
     }
@@ -156,12 +157,18 @@ impl State {
     /// Optimistically tries to swap the state assuming the join handle is
     /// __immediately__ dropped on spawn
     pub(super) fn drop_join_handle_fast(&self) -> Result<(), ()> {
-        if *self.load() == INITIAL_STATE {
-            self.store(Snapshot((INITIAL_STATE - REF_ONE) & !JOIN_INTEREST));
-            trace!("MONOIO DEBUG[State]: drop_join_handle_fast");
-            Ok(())
-        } else {
-            Err(())
+        match self.fetch_update(|curr| {
+            if *curr == INITIAL_STATE {
+                Some(Snapshot((INITIAL_STATE - REF_ONE) & !JOIN_INTEREST))
+            } else {
+                None
+            }
+        }) {
+            Ok(_) => {
+                trace!("MONOIO DEBUG[State]: drop_join_handle_fast");
+                Ok(())
+            }
+            Err(_) => Err(()),
         }
     }