Fix entity map drop behavior

Conrad Irwin created 2 years ago

The entity map needs to be able to distinguish between the case when
the entity_id is waiting to be dropped, and when it is completely gone.

Before 8bc207141, it assumed that entity_ids in dropped_entity_ids could
be re-used. This caused `take_dropped` to error because the slot had
been overwritten. The fix there caused weak handles to allow upgrading
a reference count from 0, which could resurrect items in
`dropped_entity_ids` which caused them to be dropped twice.

We could allow weak items to upgrade from 0, and delete from
dropped_entity_ids, but that seemed more complicated than necessary.

Change summary

crates/gpui2/src/app/entity_map.rs | 78 +++++++++++++++++++++++++++++--
1 file changed, 72 insertions(+), 6 deletions(-)

Detailed changes

crates/gpui2/src/app/entity_map.rs

@@ -106,7 +106,12 @@ impl EntityMap {
         dropped_entity_ids
             .into_iter()
             .map(|entity_id| {
-                ref_counts.counts.remove(entity_id);
+                let count = ref_counts.counts.remove(entity_id).unwrap();
+                debug_assert_eq!(
+                    count.load(SeqCst),
+                    0,
+                    "dropped an entity that was referenced"
+                );
                 (entity_id, self.entities.remove(entity_id).unwrap())
             })
             .collect()
@@ -356,11 +361,15 @@ impl AnyWeakHandle {
 
     pub fn upgrade(&self) -> Option<AnyHandle> {
         let entity_map = self.entity_ref_counts.upgrade()?;
-        entity_map
-            .read()
-            .counts
-            .get(self.entity_id)?
-            .fetch_add(1, SeqCst);
+        let entity_map = entity_map.read();
+        let ref_count = entity_map.counts.get(self.entity_id)?;
+
+        // entity_id is in dropped_entity_ids
+        if ref_count.load(SeqCst) == 0 {
+            return None;
+        }
+        ref_count.fetch_add(1, SeqCst);
+
         Some(AnyHandle {
             entity_id: self.entity_id,
             entity_type: self.entity_type,
@@ -460,3 +469,60 @@ impl<T> PartialEq<Handle<T>> for WeakHandle<T> {
         self.entity_id() == other.entity_id()
     }
 }
+
+#[cfg(test)]
+mod test {
+    use crate::EntityMap;
+
+    struct TestEntity {
+        pub i: i32,
+    }
+
+    #[test]
+    fn test_entity_map_slot_assignment_before_cleanup() {
+        // Tests that slots are not re-used before take_dropped.
+        let mut entity_map = EntityMap::new();
+
+        let slot = entity_map.reserve::<TestEntity>();
+        entity_map.insert(slot, TestEntity { i: 1 });
+
+        let slot = entity_map.reserve::<TestEntity>();
+        entity_map.insert(slot, TestEntity { i: 2 });
+
+        let dropped = entity_map.take_dropped();
+        assert_eq!(dropped.len(), 2);
+
+        assert_eq!(
+            dropped
+                .into_iter()
+                .map(|(_, entity)| entity.downcast::<TestEntity>().unwrap().i)
+                .collect::<Vec<i32>>(),
+            vec![1, 2],
+        );
+    }
+
+    #[test]
+    fn test_entity_map_weak_upgrade_before_cleanup() {
+        // Tests that weak handles are not upgraded before take_dropped
+        let mut entity_map = EntityMap::new();
+
+        let slot = entity_map.reserve::<TestEntity>();
+        let handle = entity_map.insert(slot, TestEntity { i: 1 });
+        let weak = handle.downgrade();
+        drop(handle);
+
+        let strong = weak.upgrade();
+        assert_eq!(strong, None);
+
+        let dropped = entity_map.take_dropped();
+        assert_eq!(dropped.len(), 1);
+
+        assert_eq!(
+            dropped
+                .into_iter()
+                .map(|(_, entity)| entity.downcast::<TestEntity>().unwrap().i)
+                .collect::<Vec<i32>>(),
+            vec![1],
+        );
+    }
+}