wgpu_core/
registry.rs

1use alloc::sync::Arc;
2
3use crate::{
4    id::Id,
5    identity::IdentityManager,
6    lock::{
7        rank::{self, LockRank},
8        RwLock, RwLockReadGuard,
9    },
10    storage::{Element, Storage, StorageItem},
11};
12
13#[derive(Copy, Clone, Debug, Default, PartialEq, Eq)]
14pub struct RegistryReport {
15    /// Count of IDs allocated by [`IdentityManager`]
16    ///
17    /// This may be inconsistent with other fields of the report if
18    /// IDs are allocated or released concurrently with report
19    /// generation.
20    pub num_allocated: usize,
21    pub num_kept_from_user: usize,
22    pub num_released_from_user: usize,
23    pub element_size: usize,
24}
25
26impl RegistryReport {
27    pub fn is_empty(&self) -> bool {
28        self.num_allocated + self.num_kept_from_user == 0
29    }
30}
31
32/// Registry is the primary holder of each resource type
33/// Every resource is now arcanized so the last arc released
34/// will in the end free the memory and release the inner raw resource
35///
36/// Registry act as the main entry point to keep resource alive
37/// when created and released from user land code
38///
39/// A resource may still be alive when released from user land code
40/// if it's used in active submission or anyway kept alive from
41/// any other dependent resource
42///
43#[derive(Debug)]
44pub(crate) struct Registry<T: StorageItem> {
45    // Must only contain an id which has either never been used or has been released from `storage`
46    identity: Arc<IdentityManager<T::Marker>>,
47    storage: RwLock<Storage<T>>,
48}
49
50impl<T: StorageItem> Registry<T> {
51    pub(crate) fn new() -> Self {
52        Self::with_rank(rank::HUB_OTHER)
53    }
54
55    pub(crate) fn with_rank(rank: LockRank) -> Self {
56        Self {
57            identity: Arc::new(IdentityManager::new()),
58            storage: RwLock::new(rank, Storage::new()),
59        }
60    }
61}
62
63#[must_use]
64pub(crate) struct FutureId<'a, T: StorageItem> {
65    id: Id<T::Marker>,
66    data: &'a RwLock<Storage<T>>,
67}
68
69impl<T: StorageItem> FutureId<'_, T> {
70    /// Assign a new resource to this ID.
71    ///
72    /// Registers it with the registry.
73    pub fn assign(self, value: T) -> Id<T::Marker> {
74        let mut data = self.data.write();
75        data.insert(self.id, value);
76        self.id
77    }
78}
79
80impl<T: StorageItem> Registry<T> {
81    pub(crate) fn prepare(&self, id_in: Option<Id<T::Marker>>) -> FutureId<'_, T> {
82        FutureId {
83            id: match id_in {
84                Some(id_in) => {
85                    self.identity.mark_as_used(id_in);
86                    id_in
87                }
88                None => self.identity.process(),
89            },
90            data: &self.storage,
91        }
92    }
93
94    #[track_caller]
95    pub(crate) fn read<'a>(&'a self) -> RwLockReadGuard<'a, Storage<T>> {
96        self.storage.read()
97    }
98    pub(crate) fn remove(&self, id: Id<T::Marker>) -> T {
99        let value = self.storage.write().remove(id);
100        // This needs to happen *after* removing it from the storage, to maintain the
101        // invariant that `self.identity` only contains ids which are actually available
102        // See https://github.com/gfx-rs/wgpu/issues/5372
103        self.identity.free(id);
104        //Returning None is legal if it's an error ID
105        value
106    }
107
108    pub(crate) fn generate_report(&self) -> RegistryReport {
109        let mut report = RegistryReport {
110            element_size: size_of::<T>(),
111            ..Default::default()
112        };
113
114        // Acquiring the identity values lock while holding the storage lock
115        // would require adding an edge in the lock graph, and would still not
116        // ensure a consistent report, because the storage lock is not otherwise
117        // acquired when managing IDs.
118        report.num_allocated = self.identity.values.lock().count();
119
120        let storage = self.storage.read();
121        for element in storage.map.iter() {
122            match *element {
123                Element::Occupied(..) => report.num_kept_from_user += 1,
124                Element::Vacant => report.num_released_from_user += 1,
125            }
126        }
127        report
128    }
129}
130
131impl<T: StorageItem + Clone> Registry<T> {
132    pub(crate) fn get(&self, id: Id<T::Marker>) -> T {
133        self.read().get(id)
134    }
135}
136
137#[cfg(test)]
138mod tests {
139    use super::Registry;
140    use crate::{id::Marker, resource::ResourceType, storage::StorageItem};
141    use alloc::sync::Arc;
142
143    struct TestData;
144    struct TestDataId;
145    impl Marker for TestDataId {
146        const TYPE: &'static str = "TestData";
147    }
148
149    impl ResourceType for TestData {
150        const TYPE: &'static str = "TestData";
151    }
152    impl StorageItem for TestData {
153        type Marker = TestDataId;
154    }
155
156    #[test]
157    fn simultaneous_registration() {
158        let registry = Registry::new();
159        std::thread::scope(|s| {
160            for _ in 0..5 {
161                s.spawn(|| {
162                    for _ in 0..1000 {
163                        let value = Arc::new(TestData);
164                        let new_id = registry.prepare(None);
165                        let id = new_id.assign(value);
166                        registry.remove(id);
167                    }
168                });
169            }
170        })
171    }
172}