wgpu_core/
registry.rs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
use alloc::sync::Arc;

use crate::{
    id::Id,
    identity::IdentityManager,
    lock::{rank, RwLock, RwLockReadGuard, RwLockWriteGuard},
    storage::{Element, Storage, StorageItem},
};

#[derive(Copy, Clone, Debug, Default, PartialEq, Eq)]
pub struct RegistryReport {
    pub num_allocated: usize,
    pub num_kept_from_user: usize,
    pub num_released_from_user: usize,
    pub element_size: usize,
}

impl RegistryReport {
    pub fn is_empty(&self) -> bool {
        self.num_allocated + self.num_kept_from_user == 0
    }
}

/// Registry is the primary holder of each resource type
/// Every resource is now arcanized so the last arc released
/// will in the end free the memory and release the inner raw resource
///
/// Registry act as the main entry point to keep resource alive
/// when created and released from user land code
///
/// A resource may still be alive when released from user land code
/// if it's used in active submission or anyway kept alive from
/// any other dependent resource
///
#[derive(Debug)]
pub(crate) struct Registry<T: StorageItem> {
    // Must only contain an id which has either never been used or has been released from `storage`
    identity: Arc<IdentityManager<T::Marker>>,
    storage: RwLock<Storage<T>>,
}

impl<T: StorageItem> Registry<T> {
    pub(crate) fn new() -> Self {
        Self {
            identity: Arc::new(IdentityManager::new()),
            storage: RwLock::new(rank::REGISTRY_STORAGE, Storage::new()),
        }
    }
}

#[must_use]
pub(crate) struct FutureId<'a, T: StorageItem> {
    id: Id<T::Marker>,
    data: &'a RwLock<Storage<T>>,
}

impl<T: StorageItem> FutureId<'_, T> {
    pub fn id(&self) -> Id<T::Marker> {
        self.id
    }

    /// Assign a new resource to this ID.
    ///
    /// Registers it with the registry.
    pub fn assign(self, value: T) -> Id<T::Marker> {
        let mut data = self.data.write();
        data.insert(self.id, value);
        self.id
    }
}

impl<T: StorageItem> Registry<T> {
    pub(crate) fn prepare(&self, id_in: Option<Id<T::Marker>>) -> FutureId<T> {
        FutureId {
            id: match id_in {
                Some(id_in) => {
                    self.identity.mark_as_used(id_in);
                    id_in
                }
                None => self.identity.process(),
            },
            data: &self.storage,
        }
    }

    #[track_caller]
    pub(crate) fn read<'a>(&'a self) -> RwLockReadGuard<'a, Storage<T>> {
        self.storage.read()
    }
    #[track_caller]
    pub(crate) fn write<'a>(&'a self) -> RwLockWriteGuard<'a, Storage<T>> {
        self.storage.write()
    }
    pub(crate) fn remove(&self, id: Id<T::Marker>) -> T {
        let value = self.storage.write().remove(id);
        // This needs to happen *after* removing it from the storage, to maintain the
        // invariant that `self.identity` only contains ids which are actually available
        // See https://github.com/gfx-rs/wgpu/issues/5372
        self.identity.free(id);
        //Returning None is legal if it's an error ID
        value
    }

    pub(crate) fn generate_report(&self) -> RegistryReport {
        let storage = self.storage.read();
        let mut report = RegistryReport {
            element_size: size_of::<T>(),
            ..Default::default()
        };
        report.num_allocated = self.identity.values.lock().count();
        for element in storage.map.iter() {
            match *element {
                Element::Occupied(..) => report.num_kept_from_user += 1,
                Element::Vacant => report.num_released_from_user += 1,
            }
        }
        report
    }
}

impl<T: StorageItem + Clone> Registry<T> {
    pub(crate) fn get(&self, id: Id<T::Marker>) -> T {
        self.read().get(id)
    }
}

#[cfg(test)]
mod tests {
    use super::Registry;
    use crate::{id::Marker, resource::ResourceType, storage::StorageItem};
    use alloc::sync::Arc;

    struct TestData;
    struct TestDataId;
    impl Marker for TestDataId {}

    impl ResourceType for TestData {
        const TYPE: &'static str = "TestData";
    }
    impl StorageItem for TestData {
        type Marker = TestDataId;
    }

    #[test]
    fn simultaneous_registration() {
        let registry = Registry::new();
        std::thread::scope(|s| {
            for _ in 0..5 {
                s.spawn(|| {
                    for _ in 0..1000 {
                        let value = Arc::new(TestData);
                        let new_id = registry.prepare(None);
                        let id = new_id.assign(value);
                        registry.remove(id);
                    }
                });
            }
        })
    }
}