1use crate::{App, AppContext, GpuiBorrow, VisualContext, Window, seal::Sealed};
2use anyhow::{Context as _, Result};
3use collections::FxHashSet;
4use derive_more::{Deref, DerefMut};
5use parking_lot::{RwLock, RwLockUpgradableReadGuard};
6use slotmap::{KeyData, SecondaryMap, SlotMap};
7use std::{
8 any::{Any, TypeId, type_name},
9 cell::{Cell, RefCell},
10 cmp::Ordering,
11 fmt::{self, Display},
12 hash::{Hash, Hasher},
13 marker::PhantomData,
14 num::NonZeroU64,
15 sync::{
16 Arc, Weak,
17 atomic::{AtomicU64, AtomicUsize, Ordering::SeqCst},
18 },
19 thread::panicking,
20};
21
22use super::Context;
23use crate::util::atomic_incr_if_not_zero;
24#[cfg(any(test, gpui_leak_detection))]
25use collections::HashMap;
26
27slotmap::new_key_type! {
28 pub struct EntityId;
33}
34
35impl From<u64> for EntityId {
36 fn from(value: u64) -> Self {
37 Self(KeyData::from_ffi(value))
38 }
39}
40
41impl EntityId {
42 pub fn as_non_zero_u64(self) -> NonZeroU64 {
49 NonZeroU64::new(self.0.as_ffi()).unwrap()
50 }
51
52 pub fn as_u64(self) -> u64 {
54 self.0.as_ffi()
55 }
56}
57
58impl Display for EntityId {
59 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
60 write!(f, "{}", self.as_u64())
61 }
62}
63
64pub(crate) struct EntityMap {
65 entities: SecondaryMap<EntityId, Box<dyn Any>>,
66 pub accessed_entities: RefCell<FxHashSet<EntityId>>,
67 pub(crate) access_log: RefCell<Vec<EntityId>>,
69 pub(crate) recordings: Cell<usize>,
70 ref_counts: Arc<RwLock<EntityRefCounts>>,
71}
72
73#[doc(hidden)]
74pub(crate) struct EntityRefCounts {
75 counts: SlotMap<EntityId, AtomicUsize>,
76 dropped_entity_ids: Vec<EntityId>,
77 #[cfg(any(test, gpui_leak_detection))]
78 leak_detector: LeakDetector,
79}
80
81pub(super) struct LeaseInner {
82 pub(super) entity: Option<Box<dyn Any>>,
83}
84
85impl EntityMap {
86 pub fn new() -> Self {
87 Self {
88 entities: SecondaryMap::new(),
89 accessed_entities: RefCell::new(FxHashSet::default()),
90 access_log: RefCell::default(),
91 recordings: Cell::new(0),
92 ref_counts: Arc::new(RwLock::new(EntityRefCounts {
93 counts: SlotMap::with_key(),
94 dropped_entity_ids: Vec::new(),
95 #[cfg(any(test, gpui_leak_detection))]
96 leak_detector: LeakDetector {
97 next_handle_id: 0,
98 entity_handles: HashMap::default(),
99 },
100 })),
101 }
102 }
103
104 #[doc(hidden)]
105 pub fn ref_counts_drop_handle(&self) -> Arc<RwLock<EntityRefCounts>> {
106 self.ref_counts.clone()
107 }
108
109 #[cfg(any(test, gpui_leak_detection))]
115 pub fn leak_detector_snapshot(&self) -> LeakDetectorSnapshot {
116 self.ref_counts.read().leak_detector.snapshot()
117 }
118
119 #[cfg(any(test, gpui_leak_detection))]
123 pub fn assert_no_new_leaks(&self, snapshot: &LeakDetectorSnapshot) {
124 self.ref_counts
125 .read()
126 .leak_detector
127 .assert_no_new_leaks(snapshot)
128 }
129
130 pub fn reserve<T: 'static>(&self) -> Slot<T> {
132 let id = self.ref_counts.write().counts.insert(1.into());
133 Slot(Entity::new(id, Arc::downgrade(&self.ref_counts)))
134 }
135
136 pub fn insert<T>(&mut self, slot: Slot<T>, entity: T) -> Entity<T>
138 where
139 T: 'static,
140 {
141 self.note_access(slot.entity_id);
142
143 let handle = slot.0;
144 self.entities.insert(handle.entity_id, Box::new(entity));
145 handle
146 }
147
148 #[track_caller]
150 pub fn lease<T>(&mut self, pointer: &Entity<T>) -> Lease<T> {
151 Lease {
152 inner: self.lease_erased(pointer, type_name::<T>()),
153 id: pointer.entity_id,
154 entity_type: PhantomData,
155 }
156 }
157
158 pub fn end_lease<T>(&mut self, lease: Lease<T>) {
160 self.end_lease_erased(lease.id, lease.inner);
161 }
162
163 #[inline(always)]
164 pub fn read<T: 'static>(&self, entity: &Entity<T>) -> &T {
165 self.assert_valid_context(entity);
166 self.read_inner(entity.entity_id)
167 .and_then(|entity| entity.downcast_ref())
168 .unwrap_or_else(|| double_lease_panic("read", type_name::<T>()))
169 }
170
171 #[track_caller]
172 pub(super) fn lease_erased(&mut self, pointer: &AnyEntity, entity_type: &str) -> LeaseInner {
173 self.assert_valid_context(pointer);
174 let entity = Some(
175 self.lease_inner(pointer.entity_id)
176 .unwrap_or_else(|| double_lease_panic("update", entity_type)),
177 );
178 LeaseInner { entity }
179 }
180
181 pub(super) fn end_lease_erased(&mut self, entity_id: EntityId, mut lease: LeaseInner) {
182 self.end_lease_inner(entity_id, lease.entity.take().unwrap());
183 }
184
185 fn assert_valid_context(&self, entity: &AnyEntity) {
186 debug_assert!(
187 Weak::ptr_eq(&entity.entity_map, &Arc::downgrade(&self.ref_counts)),
188 "used a entity with the wrong context"
189 );
190 }
191
192 pub fn clear_accessed(&mut self) {
193 self.accessed_entities.get_mut().clear();
194 }
195
196 pub fn take_dropped(&mut self) -> Vec<(EntityId, Box<dyn Any>)> {
197 let mut ref_counts = &mut *self.ref_counts.write();
198 let dropped_entity_ids = ref_counts.dropped_entity_ids.drain(..);
199 let mut accessed_entities = self.accessed_entities.get_mut();
200
201 dropped_entity_ids
202 .filter_map(|entity_id| {
203 let count = ref_counts.counts.remove(entity_id).unwrap();
204 debug_assert_eq!(
205 count.load(SeqCst),
206 0,
207 "dropped an entity that was referenced"
208 );
209 accessed_entities.remove(&entity_id);
210 Some((entity_id, self.entities.remove(entity_id)?))
213 })
214 .collect()
215 }
216
217 #[inline(never)]
218 fn read_inner(&self, entity_id: EntityId) -> Option<&dyn Any> {
219 self.note_access(entity_id);
220 self.entities.get(entity_id).map(Box::as_ref)
221 }
222
223 pub(crate) fn note_access(&self, entity_id: EntityId) {
224 self.accessed_entities.borrow_mut().insert(entity_id);
225 if self.recordings.get() > 0 {
226 self.access_log.borrow_mut().push(entity_id);
227 }
228 }
229
230 #[inline(never)]
231 fn lease_inner(&mut self, entity_id: EntityId) -> Option<Box<dyn Any>> {
232 self.note_access(entity_id);
233 self.entities.remove(entity_id)
234 }
235
236 #[inline(never)]
237 fn end_lease_inner(&mut self, entity_id: EntityId, entity: Box<dyn Any>) {
238 self.entities.insert(entity_id, entity);
239 }
240}
241
242#[track_caller]
243fn double_lease_panic(operation: &str, entity_type: &str) -> ! {
244 panic!("cannot {operation} {entity_type} while it is already being updated")
245}
246
247pub(crate) struct Lease<T> {
248 pub id: EntityId,
249 inner: LeaseInner,
250 entity_type: PhantomData<T>,
251}
252
253impl<T: 'static> core::ops::Deref for Lease<T> {
254 type Target = T;
255
256 fn deref(&self) -> &Self::Target {
257 self.inner.entity.as_ref().unwrap().downcast_ref().unwrap()
258 }
259}
260
261impl<T: 'static> core::ops::DerefMut for Lease<T> {
262 fn deref_mut(&mut self) -> &mut Self::Target {
263 self.inner.entity.as_mut().unwrap().downcast_mut().unwrap()
264 }
265}
266
267impl Drop for LeaseInner {
268 fn drop(&mut self) {
269 if self.entity.is_some() && !panicking() {
270 panic!("Leases must be ended with EntityMap::end_lease")
271 }
272 }
273}
274
275#[derive(Deref, DerefMut)]
276pub(crate) struct Slot<T>(Entity<T>);
277
278pub struct AnyEntity {
299 pub(crate) entity_id: EntityId,
300 pub(crate) entity_type: TypeId,
301 entity_map: Weak<RwLock<EntityRefCounts>>,
302 #[cfg(any(test, gpui_leak_detection))]
303 handle_id: HandleId,
304}
305
306impl AnyEntity {
307 fn new(
308 id: EntityId,
309 entity_type: TypeId,
310 entity_map: Weak<RwLock<EntityRefCounts>>,
311 #[cfg(any(test, gpui_leak_detection))] type_name: &'static str,
312 ) -> Self {
313 Self {
314 entity_id: id,
315 entity_type,
316 #[cfg(any(test, gpui_leak_detection))]
317 handle_id: entity_map
318 .clone()
319 .upgrade()
320 .unwrap()
321 .write()
322 .leak_detector
323 .handle_created(id, Some(type_name)),
324 entity_map,
325 }
326 }
327
328 #[inline]
330 pub fn entity_id(&self) -> EntityId {
331 self.entity_id
332 }
333
334 #[inline]
336 pub fn entity_type(&self) -> TypeId {
337 self.entity_type
338 }
339
340 pub fn downgrade(&self) -> AnyWeakEntity {
342 AnyWeakEntity {
343 entity_id: self.entity_id,
344 entity_type: self.entity_type,
345 entity_ref_counts: self.entity_map.clone(),
346 }
347 }
348
349 pub fn downcast<T: 'static>(self) -> Result<Entity<T>, AnyEntity> {
353 if TypeId::of::<T>() == self.entity_type {
354 Ok(Entity {
355 any_entity: self,
356 entity_type: PhantomData,
357 })
358 } else {
359 Err(self)
360 }
361 }
362}
363
364impl Clone for AnyEntity {
365 fn clone(&self) -> Self {
366 if let Some(entity_map) = self.entity_map.upgrade() {
367 let entity_map = entity_map.read();
368 let count = entity_map
369 .counts
370 .get(self.entity_id)
371 .expect("detected over-release of a entity");
372 let prev_count = count.fetch_add(1, SeqCst);
373 assert_ne!(prev_count, 0, "Detected over-release of a entity.");
374 }
375
376 Self {
377 entity_id: self.entity_id,
378 entity_type: self.entity_type,
379 entity_map: self.entity_map.clone(),
380 #[cfg(any(test, gpui_leak_detection))]
381 handle_id: self
382 .entity_map
383 .upgrade()
384 .unwrap()
385 .write()
386 .leak_detector
387 .handle_created(self.entity_id, None),
388 }
389 }
390}
391
392impl Drop for AnyEntity {
393 fn drop(&mut self) {
394 if let Some(entity_map) = self.entity_map.upgrade() {
395 let entity_map = entity_map.upgradable_read();
396 let count = entity_map
397 .counts
398 .get(self.entity_id)
399 .expect("detected over-release of a handle.");
400 let prev_count = count.fetch_sub(1, SeqCst);
401 assert_ne!(prev_count, 0, "Detected over-release of a entity.");
402 if prev_count == 1 {
403 let mut entity_map = RwLockUpgradableReadGuard::upgrade(entity_map);
405 entity_map.dropped_entity_ids.push(self.entity_id);
406 }
407 }
408
409 #[cfg(any(test, gpui_leak_detection))]
410 if let Some(entity_map) = self.entity_map.upgrade() {
411 entity_map
412 .write()
413 .leak_detector
414 .handle_released(self.entity_id, self.handle_id)
415 }
416 }
417}
418
419impl<T> From<Entity<T>> for AnyEntity {
420 #[inline]
421 fn from(entity: Entity<T>) -> Self {
422 entity.any_entity
423 }
424}
425
426impl Hash for AnyEntity {
427 #[inline]
428 fn hash<H: Hasher>(&self, state: &mut H) {
429 self.entity_id.hash(state);
430 }
431}
432
433impl PartialEq for AnyEntity {
434 #[inline]
435 fn eq(&self, other: &Self) -> bool {
436 self.entity_id == other.entity_id
437 }
438}
439
440impl Eq for AnyEntity {}
441
442impl Ord for AnyEntity {
443 #[inline]
444 fn cmp(&self, other: &Self) -> Ordering {
445 self.entity_id.cmp(&other.entity_id)
446 }
447}
448
449impl PartialOrd for AnyEntity {
450 #[inline]
451 fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
452 Some(self.cmp(other))
453 }
454}
455
456impl std::fmt::Debug for AnyEntity {
457 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
458 f.debug_struct("AnyEntity")
459 .field("entity_id", &self.entity_id.as_u64())
460 .finish()
461 }
462}
463
464#[derive(Deref, DerefMut)]
494pub struct Entity<T> {
495 #[deref]
496 #[deref_mut]
497 pub(crate) any_entity: AnyEntity,
498 pub(crate) entity_type: PhantomData<fn(T) -> T>,
499}
500
501impl<T> Sealed for Entity<T> {}
502
503impl<T: 'static> Entity<T> {
504 #[inline]
505 fn new(id: EntityId, entity_map: Weak<RwLock<EntityRefCounts>>) -> Self
506 where
507 T: 'static,
508 {
509 Self {
510 any_entity: AnyEntity::new(
511 id,
512 TypeId::of::<T>(),
513 entity_map,
514 #[cfg(any(test, gpui_leak_detection))]
515 std::any::type_name::<T>(),
516 ),
517 entity_type: PhantomData,
518 }
519 }
520
521 #[inline]
525 pub fn entity_id(&self) -> EntityId {
526 self.any_entity.entity_id
527 }
528
529 #[inline]
534 pub fn downgrade(&self) -> WeakEntity<T> {
535 WeakEntity {
536 any_entity: self.any_entity.downgrade(),
537 entity_type: self.entity_type,
538 }
539 }
540
541 #[inline]
543 pub fn into_any(self) -> AnyEntity {
544 self.any_entity
545 }
546
547 #[inline]
573 pub fn read<'a>(&self, cx: &'a App) -> &'a T {
574 cx.entities.read(self)
575 }
576
577 #[inline]
582 pub fn read_with<R, C: AppContext>(&self, cx: &C, f: impl FnOnce(&T, &App) -> R) -> R {
583 cx.read_entity(self, f)
584 }
585
586 #[inline]
623 pub fn update<R, C: AppContext>(
624 &self,
625 cx: &mut C,
626 update: impl FnOnce(&mut T, &mut Context<T>) -> R,
627 ) -> R {
628 cx.update_entity(self, update)
629 }
630
631 #[inline]
637 pub fn as_mut<'a, C: AppContext>(&self, cx: &'a mut C) -> GpuiBorrow<'a, T> {
638 cx.as_mut(self)
639 }
640
641 pub fn write<C: AppContext>(&self, cx: &mut C, value: T) {
655 self.update(cx, |entity, cx| {
656 *entity = value;
657 cx.notify();
658 })
659 }
660
661 #[inline]
666 pub fn update_in<R, C: VisualContext>(
667 &self,
668 cx: &mut C,
669 update: impl FnOnce(&mut T, &mut Window, &mut Context<T>) -> R,
670 ) -> C::Result<R> {
671 cx.update_window_entity(self, update)
672 }
673}
674
675impl<T> Clone for Entity<T> {
676 #[inline]
677 fn clone(&self) -> Self {
678 Self {
679 any_entity: self.any_entity.clone(),
680 entity_type: self.entity_type,
681 }
682 }
683}
684
685impl<T> std::fmt::Debug for Entity<T> {
686 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
687 f.debug_struct("Entity")
688 .field("entity_id", &self.any_entity.entity_id)
689 .field("entity_type", &type_name::<T>())
690 .finish()
691 }
692}
693
694impl<T> Hash for Entity<T> {
695 #[inline]
696 fn hash<H: Hasher>(&self, state: &mut H) {
697 self.any_entity.hash(state);
698 }
699}
700
701impl<T> PartialEq for Entity<T> {
702 #[inline]
703 fn eq(&self, other: &Self) -> bool {
704 self.any_entity == other.any_entity
705 }
706}
707
708impl<T> Eq for Entity<T> {}
709
710impl<T> PartialEq<WeakEntity<T>> for Entity<T> {
711 #[inline]
712 fn eq(&self, other: &WeakEntity<T>) -> bool {
713 self.any_entity.entity_id() == other.entity_id()
714 }
715}
716
717impl<T: 'static> Ord for Entity<T> {
718 #[inline]
719 fn cmp(&self, other: &Self) -> std::cmp::Ordering {
720 self.entity_id().cmp(&other.entity_id())
721 }
722}
723
724impl<T: 'static> PartialOrd for Entity<T> {
725 #[inline]
726 fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
727 Some(self.cmp(other))
728 }
729}
730
731#[derive(Clone)]
735pub struct AnyWeakEntity {
736 pub(crate) entity_id: EntityId,
737 entity_type: TypeId,
738 entity_ref_counts: Weak<RwLock<EntityRefCounts>>,
739}
740
741impl AnyWeakEntity {
742 #[inline]
744 pub fn entity_id(&self) -> EntityId {
745 self.entity_id
746 }
747
748 pub fn is_upgradable(&self) -> bool {
753 let ref_count = self
754 .entity_ref_counts
755 .upgrade()
756 .and_then(|ref_counts| Some(ref_counts.read().counts.get(self.entity_id)?.load(SeqCst)))
757 .unwrap_or(0);
758 ref_count > 0
759 }
760
761 pub fn upgrade(&self) -> Option<AnyEntity> {
765 let ref_counts = &self.entity_ref_counts.upgrade()?;
766 let ref_counts = ref_counts.read();
767 let ref_count = ref_counts.counts.get(self.entity_id)?;
768
769 if atomic_incr_if_not_zero(ref_count) == 0 {
770 return None;
772 }
773 drop(ref_counts);
774
775 Some(AnyEntity {
776 entity_id: self.entity_id,
777 entity_type: self.entity_type,
778 entity_map: self.entity_ref_counts.clone(),
779 #[cfg(any(test, gpui_leak_detection))]
780 handle_id: self
781 .entity_ref_counts
782 .upgrade()
783 .unwrap()
784 .write()
785 .leak_detector
786 .handle_created(self.entity_id, None),
787 })
788 }
789
790 #[cfg(any(
832 test,
833 feature = "test-support",
834 feature = "leak-detection",
835 gpui_leak_detection
836 ))]
837 pub fn assert_released(&self) {
838 #[cfg(any(test, gpui_leak_detection))]
839 self.entity_ref_counts
840 .upgrade()
841 .unwrap()
842 .write()
843 .leak_detector
844 .assert_released(self.entity_id);
845
846 if self
847 .entity_ref_counts
848 .upgrade()
849 .and_then(|ref_counts| Some(ref_counts.read().counts.get(self.entity_id)?.load(SeqCst)))
850 .is_some()
851 {
852 panic!(
853 "entity is still alive, or was recently dropped and its resources are retained \
854 until the end of the effect cycle. Build with GPUI_LEAK_DETECTION=1 to list \
855 its live handles."
856 )
857 }
858 }
859
860 pub fn new_invalid() -> Self {
862 static UNIQUE_NON_CONFLICTING_ID_GENERATOR: AtomicU64 = AtomicU64::new(u64::MAX);
866 let entity_id = UNIQUE_NON_CONFLICTING_ID_GENERATOR.fetch_sub(1, SeqCst);
867
868 Self {
869 entity_id: entity_id.into(),
879 entity_type: TypeId::of::<()>(),
880 entity_ref_counts: Weak::new(),
881 }
882 }
883}
884
885impl std::fmt::Debug for AnyWeakEntity {
886 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
887 f.debug_struct(type_name::<Self>())
888 .field("entity_id", &self.entity_id)
889 .field("entity_type", &self.entity_type)
890 .finish()
891 }
892}
893
894impl<T> From<WeakEntity<T>> for AnyWeakEntity {
895 #[inline]
896 fn from(entity: WeakEntity<T>) -> Self {
897 entity.any_entity
898 }
899}
900
901impl Hash for AnyWeakEntity {
902 #[inline]
903 fn hash<H: Hasher>(&self, state: &mut H) {
904 self.entity_id.hash(state);
905 }
906}
907
908impl PartialEq for AnyWeakEntity {
909 #[inline]
910 fn eq(&self, other: &Self) -> bool {
911 self.entity_id == other.entity_id
912 }
913}
914
915impl Eq for AnyWeakEntity {}
916
917impl Ord for AnyWeakEntity {
918 #[inline]
919 fn cmp(&self, other: &Self) -> Ordering {
920 self.entity_id.cmp(&other.entity_id)
921 }
922}
923
924impl PartialOrd for AnyWeakEntity {
925 #[inline]
926 fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
927 Some(self.cmp(other))
928 }
929}
930
931#[derive(Deref, DerefMut)]
938pub struct WeakEntity<T> {
939 #[deref]
940 #[deref_mut]
941 any_entity: AnyWeakEntity,
942 entity_type: PhantomData<fn(T) -> T>,
943}
944
945impl<T> std::fmt::Debug for WeakEntity<T> {
946 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
947 f.debug_struct(type_name::<Self>())
948 .field("entity_id", &self.any_entity.entity_id)
949 .field("entity_type", &type_name::<T>())
950 .finish()
951 }
952}
953
954impl<T> Clone for WeakEntity<T> {
955 fn clone(&self) -> Self {
956 Self {
957 any_entity: self.any_entity.clone(),
958 entity_type: self.entity_type,
959 }
960 }
961}
962
963impl<T: 'static> WeakEntity<T> {
964 pub fn upgrade(&self) -> Option<Entity<T>> {
966 Some(Entity {
967 any_entity: self.any_entity.upgrade()?,
968 entity_type: self.entity_type,
969 })
970 }
971
972 #[inline(always)]
1006 pub fn update<C, R>(
1007 &self,
1008 cx: &mut C,
1009 update: impl FnOnce(&mut T, &mut Context<T>) -> R,
1010 ) -> Result<R>
1011 where
1012 C: AppContext,
1013 {
1014 let entity = self.upgrade().context("entity released")?;
1015 Ok(cx.update_entity(&entity, update))
1016 }
1017
1018 #[inline(always)]
1023 pub fn update_in<C, R>(
1024 &self,
1025 cx: &mut C,
1026 update: impl FnOnce(&mut T, &mut Window, &mut Context<T>) -> R,
1027 ) -> Result<R>
1028 where
1029 C: AppContext,
1030 {
1031 let entity = self.upgrade().context("entity released")?;
1032 cx.with_window(entity.entity_id(), |window, app| {
1033 entity.update(app, |entity, cx| update(entity, window, cx))
1034 })
1035 .context("entity has no current window")
1036 }
1037
1038 #[inline(always)]
1042 pub fn read_with<C, R>(&self, cx: &C, read: impl FnOnce(&T, &App) -> R) -> Result<R>
1043 where
1044 C: AppContext,
1045 {
1046 let entity = self.upgrade().context("entity released")?;
1047 Ok(cx.read_entity(&entity, read))
1048 }
1049
1050 #[inline]
1052 pub fn new_invalid() -> Self {
1053 Self {
1054 any_entity: AnyWeakEntity::new_invalid(),
1055 entity_type: PhantomData,
1056 }
1057 }
1058}
1059
1060impl<T> Hash for WeakEntity<T> {
1061 #[inline]
1062 fn hash<H: Hasher>(&self, state: &mut H) {
1063 self.any_entity.hash(state);
1064 }
1065}
1066
1067impl<T> PartialEq for WeakEntity<T> {
1068 #[inline]
1069 fn eq(&self, other: &Self) -> bool {
1070 self.any_entity == other.any_entity
1071 }
1072}
1073
1074impl<T> Eq for WeakEntity<T> {}
1075
1076impl<T> PartialEq<Entity<T>> for WeakEntity<T> {
1077 #[inline]
1078 fn eq(&self, other: &Entity<T>) -> bool {
1079 self.entity_id() == other.any_entity.entity_id()
1080 }
1081}
1082
1083impl<T: 'static> Ord for WeakEntity<T> {
1084 #[inline]
1085 fn cmp(&self, other: &Self) -> Ordering {
1086 self.entity_id().cmp(&other.entity_id())
1087 }
1088}
1089
1090impl<T: 'static> PartialOrd for WeakEntity<T> {
1091 #[inline]
1092 fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
1093 Some(self.cmp(other))
1094 }
1095}
1096
1097#[cfg(any(test, gpui_leak_detection))]
1102static LEAK_BACKTRACE: std::sync::LazyLock<bool> =
1103 std::sync::LazyLock::new(|| std::env::var("LEAK_BACKTRACE").is_ok_and(|b| !b.is_empty()));
1104
1105#[cfg(any(test, gpui_leak_detection))]
1110#[derive(Clone, Copy, Debug, Default, Hash, PartialEq, Eq)]
1111pub(crate) struct HandleId {
1112 id: u64,
1113}
1114
1115#[cfg(any(test, gpui_leak_detection))]
1157pub(crate) struct LeakDetector {
1158 next_handle_id: u64,
1159 entity_handles: HashMap<EntityId, EntityLeakData>,
1160}
1161
1162#[cfg(any(
1167 test,
1168 feature = "test-support",
1169 feature = "leak-detection",
1170 gpui_leak_detection
1171))]
1172#[derive(Default)]
1173pub struct LeakDetectorSnapshot {
1174 #[cfg(any(test, gpui_leak_detection))]
1175 entity_ids: collections::HashSet<EntityId>,
1176}
1177
1178#[cfg(any(test, gpui_leak_detection))]
1179struct EntityLeakData {
1180 handles: HashMap<HandleId, Option<backtrace::Backtrace>>,
1181 type_name: &'static str,
1182}
1183
1184#[cfg(any(test, gpui_leak_detection))]
1185impl LeakDetector {
1186 #[track_caller]
1192 pub fn handle_created(
1193 &mut self,
1194 entity_id: EntityId,
1195 type_name: Option<&'static str>,
1196 ) -> HandleId {
1197 let id = gpui_util::post_inc(&mut self.next_handle_id);
1198 let handle_id = HandleId { id };
1199 let handles = self
1200 .entity_handles
1201 .entry(entity_id)
1202 .or_insert_with(|| EntityLeakData {
1203 handles: HashMap::default(),
1204 type_name: type_name.unwrap_or("<unknown>"),
1205 });
1206 handles.handles.insert(
1207 handle_id,
1208 LEAK_BACKTRACE.then(backtrace::Backtrace::new_unresolved),
1209 );
1210 handle_id
1211 }
1212
1213 pub fn handle_released(&mut self, entity_id: EntityId, handle_id: HandleId) {
1218 if let std::collections::hash_map::Entry::Occupied(mut data) =
1219 self.entity_handles.entry(entity_id)
1220 {
1221 data.get_mut().handles.remove(&handle_id);
1222 if data.get().handles.is_empty() {
1223 data.remove();
1224 }
1225 }
1226 }
1227
1228 pub fn assert_released(&mut self, entity_id: EntityId) {
1236 use std::fmt::Write as _;
1237
1238 if let Some(data) = self.entity_handles.remove(&entity_id) {
1239 let mut out = String::new();
1240 for (_, backtrace) in data.handles {
1241 if let Some(mut backtrace) = backtrace {
1242 backtrace.resolve();
1243 let backtrace = BacktraceFormatter(backtrace);
1244 writeln!(out, "Leaked handle:\n{:?}", backtrace).unwrap();
1245 } else {
1246 writeln!(
1247 out,
1248 "Leaked handle: (export LEAK_BACKTRACE to find allocation site)"
1249 )
1250 .unwrap();
1251 }
1252 }
1253 panic!("Handles for {} leaked:\n{out}", data.type_name);
1254 }
1255 }
1256
1257 pub fn snapshot(&self) -> LeakDetectorSnapshot {
1263 LeakDetectorSnapshot {
1264 entity_ids: self.entity_handles.keys().copied().collect(),
1265 }
1266 }
1267
1268 pub fn assert_no_new_leaks(&self, snapshot: &LeakDetectorSnapshot) {
1280 use std::fmt::Write as _;
1281
1282 let mut out = String::new();
1283 for (entity_id, data) in &self.entity_handles {
1284 if snapshot.entity_ids.contains(entity_id) {
1285 continue;
1286 }
1287 for (_, backtrace) in &data.handles {
1288 if let Some(backtrace) = backtrace {
1289 let mut backtrace = backtrace.clone();
1290 backtrace.resolve();
1291 let backtrace = BacktraceFormatter(backtrace);
1292 writeln!(
1293 out,
1294 "Leaked handle for entity {} ({entity_id:?}):\n{:?}",
1295 data.type_name, backtrace
1296 )
1297 .unwrap();
1298 } else {
1299 writeln!(
1300 out,
1301 "Leaked handle for entity {} ({entity_id:?}): (export LEAK_BACKTRACE to find allocation site)",
1302 data.type_name
1303 )
1304 .unwrap();
1305 }
1306 }
1307 }
1308
1309 if !out.is_empty() {
1310 panic!("New entity leaks detected since snapshot:\n{out}");
1311 }
1312 }
1313}
1314
1315#[cfg(any(test, gpui_leak_detection))]
1316impl Drop for LeakDetector {
1317 fn drop(&mut self) {
1318 use std::fmt::Write;
1319
1320 if self.entity_handles.is_empty() || std::thread::panicking() {
1321 return;
1322 }
1323
1324 let mut out = String::new();
1325 for (entity_id, data) in self.entity_handles.drain() {
1326 for (_handle, backtrace) in data.handles {
1327 if let Some(mut backtrace) = backtrace {
1328 backtrace.resolve();
1329 let backtrace = BacktraceFormatter(backtrace);
1330 writeln!(
1331 out,
1332 "Leaked handle for entity {} ({entity_id:?}):\n{:?}",
1333 data.type_name, backtrace
1334 )
1335 .unwrap();
1336 } else {
1337 writeln!(
1338 out,
1339 "Leaked handle for entity {} ({entity_id:?}): (export LEAK_BACKTRACE to find allocation site)",
1340 data.type_name
1341 )
1342 .unwrap();
1343 }
1344 }
1345 }
1346 panic!("Exited with leaked handles:\n{out}");
1347 }
1348}
1349
1350#[cfg(any(test, gpui_leak_detection))]
1351struct BacktraceFormatter(backtrace::Backtrace);
1352
1353#[cfg(any(test, gpui_leak_detection))]
1354impl fmt::Debug for BacktraceFormatter {
1355 fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
1356 use backtrace::{BacktraceFmt, BytesOrWideString, PrintFmt};
1357
1358 let style = if fmt.alternate() {
1359 PrintFmt::Full
1360 } else {
1361 PrintFmt::Short
1362 };
1363
1364 let cwd = std::env::current_dir();
1369 let mut print_path = move |fmt: &mut fmt::Formatter<'_>, path: BytesOrWideString<'_>| {
1370 let path = path.into_path_buf();
1371 if style != PrintFmt::Full {
1372 if let Ok(cwd) = &cwd {
1373 if let Ok(suffix) = path.strip_prefix(cwd) {
1374 return fmt::Display::fmt(&suffix.display(), fmt);
1375 }
1376 }
1377 }
1378 fmt::Display::fmt(&path.display(), fmt)
1379 };
1380
1381 let mut f = BacktraceFmt::new(fmt, style, &mut print_path);
1382 f.add_context()?;
1383 let mut strip = true;
1384 for frame in self.0.frames() {
1385 if let [symbol, ..] = frame.symbols()
1386 && let Some(name) = symbol.name()
1387 && let Some(filename) = name.as_str()
1388 {
1389 match filename {
1390 "test::run_test_in_process"
1391 | "scheduler::executor::spawn_local_with_source_location::impl$1::poll<core::pin::Pin<alloc::boxed::Box<dyn$<core::future::future::Future<assoc$<Output,enum2$<core::result::Result<workspace::OpenResult,anyhow::Error> > > > >,alloc::alloc::Global> > >" => {
1392 strip = true
1393 }
1394 "gpui::app::entity_map::LeakDetector::handle_created" => {
1395 strip = false;
1396 continue;
1397 }
1398 "zed::main" => {
1399 strip = true;
1400 f.frame().backtrace_frame(frame)?;
1401 }
1402 _ => {}
1403 }
1404 }
1405 if strip {
1406 continue;
1407 }
1408 f.frame().backtrace_frame(frame)?;
1409 }
1410 f.finish()?;
1411 Ok(())
1412 }
1413}
1414
1415#[cfg(test)]
1416mod test {
1417 use crate::EntityMap;
1418
1419 struct TestEntity {
1420 pub i: i32,
1421 }
1422
1423 #[test]
1424 fn test_entity_map_slot_assignment_before_cleanup() {
1425 let mut entity_map = EntityMap::new();
1427
1428 let slot = entity_map.reserve::<TestEntity>();
1429 entity_map.insert(slot, TestEntity { i: 1 });
1430
1431 let slot = entity_map.reserve::<TestEntity>();
1432 entity_map.insert(slot, TestEntity { i: 2 });
1433
1434 let dropped = entity_map.take_dropped();
1435 assert_eq!(dropped.len(), 2);
1436
1437 assert_eq!(
1438 dropped
1439 .into_iter()
1440 .map(|(_, entity)| entity.downcast::<TestEntity>().unwrap().i)
1441 .collect::<Vec<i32>>(),
1442 vec![1, 2],
1443 );
1444 }
1445
1446 #[test]
1447 fn test_entity_map_weak_upgrade_before_cleanup() {
1448 let mut entity_map = EntityMap::new();
1450
1451 let slot = entity_map.reserve::<TestEntity>();
1452 let handle = entity_map.insert(slot, TestEntity { i: 1 });
1453 let weak = handle.downgrade();
1454 drop(handle);
1455
1456 let strong = weak.upgrade();
1457 assert_eq!(strong, None);
1458
1459 let dropped = entity_map.take_dropped();
1460 assert_eq!(dropped.len(), 1);
1461
1462 assert_eq!(
1463 dropped
1464 .into_iter()
1465 .map(|(_, entity)| entity.downcast::<TestEntity>().unwrap().i)
1466 .collect::<Vec<i32>>(),
1467 vec![1],
1468 );
1469 }
1470
1471 #[test]
1472 fn test_leak_detector_snapshot_no_leaks() {
1473 let mut entity_map = EntityMap::new();
1474
1475 let slot = entity_map.reserve::<TestEntity>();
1476 let pre_existing = entity_map.insert(slot, TestEntity { i: 1 });
1477
1478 let snapshot = entity_map.leak_detector_snapshot();
1479
1480 let slot = entity_map.reserve::<TestEntity>();
1481 let temporary = entity_map.insert(slot, TestEntity { i: 2 });
1482 drop(temporary);
1483
1484 entity_map.assert_no_new_leaks(&snapshot);
1485
1486 drop(pre_existing);
1487 }
1488
1489 #[test]
1490 #[should_panic(expected = "New entity leaks detected since snapshot")]
1491 fn test_leak_detector_snapshot_detects_new_leak() {
1492 let mut entity_map = EntityMap::new();
1493
1494 let slot = entity_map.reserve::<TestEntity>();
1495 let pre_existing = entity_map.insert(slot, TestEntity { i: 1 });
1496
1497 let snapshot = entity_map.leak_detector_snapshot();
1498
1499 let slot = entity_map.reserve::<TestEntity>();
1500 let leaked = entity_map.insert(slot, TestEntity { i: 2 });
1501
1502 entity_map.assert_no_new_leaks(&snapshot);
1504
1505 drop(pre_existing);
1506 drop(leaked);
1507 }
1508}