gpui: Fix race condition when upgrading a weak reference (#30952)

Release Notes:

- N/A

Author: laizy

crates/gpui/src/app/entity_map.rs

@@ -20,11 +20,11 @@ use std::{
     thread::panicking,
 };
 
+use super::Context;
+use crate::util::atomic_incr_if_not_zero;
 #[cfg(any(test, feature = "leak-detection"))]
 use collections::HashMap;
 
-use super::Context;
-
 slotmap::new_key_type! {
     /// A unique identifier for a entity across the application.
     pub struct EntityId;
@@ -529,11 +529,10 @@ impl AnyWeakEntity {
         let ref_counts = ref_counts.read();
         let ref_count = ref_counts.counts.get(self.entity_id)?;
 
-        // entity_id is in dropped_entity_ids
-        if ref_count.load(SeqCst) == 0 {
+        if atomic_incr_if_not_zero(ref_count) == 0 {
+            // entity_id is in dropped_entity_ids
             return None;
         }
-        ref_count.fetch_add(1, SeqCst);
         drop(ref_counts);
 
         Some(AnyEntity {

crates/gpui/src/util.rs

@@ -1,3 +1,5 @@
+use std::sync::atomic::AtomicUsize;
+use std::sync::atomic::Ordering::SeqCst;
 #[cfg(any(test, feature = "test-support"))]
 use std::time::Duration;
 
@@ -108,3 +110,18 @@ impl std::fmt::Debug for CwdBacktrace<'_> {
         fmt.finish()
     }
 }
+
+/// Increment the given atomic counter if it is not zero.
+/// Return the new value of the counter.
+pub(crate) fn atomic_incr_if_not_zero(counter: &AtomicUsize) -> usize {
+    let mut loaded = counter.load(SeqCst);
+    loop {
+        if loaded == 0 {
+            return 0;
+        }
+        match counter.compare_exchange_weak(loaded, loaded + 1, SeqCst, SeqCst) {
+            Ok(x) => return x + 1,
+            Err(actual) => loaded = actual,
+        }
+    }
+}

crates/gpui/src/window.rs

@@ -52,6 +52,7 @@ use uuid::Uuid;
 
 mod prompts;
 
+use crate::util::atomic_incr_if_not_zero;
 pub use prompts::*;
 
 pub(crate) const DEFAULT_WINDOW_SIZE: Size<Pixels> = size(px(1024.), px(700.));
@@ -263,15 +264,13 @@ impl FocusHandle {
     pub(crate) fn for_id(id: FocusId, handles: &Arc<FocusMap>) -> Option<Self> {
         let lock = handles.read();
         let ref_count = lock.get(id)?;
-        if ref_count.load(SeqCst) == 0 {
-            None
-        } else {
-            ref_count.fetch_add(1, SeqCst);
-            Some(Self {
-                id,
-                handles: handles.clone(),
-            })
+        if atomic_incr_if_not_zero(ref_count) == 0 {
+            return None;
         }
+        Some(Self {
+            id,
+            handles: handles.clone(),
+        })
     }
 
     /// Converts this focus handle into a weak variant, which does not prevent it from being released.