1use std::{
2 collections::HashMap,
3 fmt::Debug,
4 ops::Deref,
5 sync::{atomic::AtomicU64, Arc, RwLock, Weak},
6};
7
8use serde::{Deserialize, Serialize};
9
10pub enum ResourceUpdateAction {
11 Insert(MistyResourceId, Vec<u8>),
12 Remove(MistyResourceId),
13}
14
15#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
16pub struct MistyResourceId(u64);
17
18impl MistyResourceId {
19 pub fn wrap(id: u64) -> Self {
20 Self(id)
21 }
22 pub fn invalid() -> Self {
23 Self(0)
24 }
25}
26
27impl Default for MistyResourceId {
28 fn default() -> Self {
29 Self::invalid()
30 }
31}
32
33const _: () = {
34 static ALLOCATED: AtomicU64 = AtomicU64::new(1);
35 impl MistyResourceId {
36 pub fn alloc() -> Self {
37 let id = ALLOCATED.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
38 return Self(id);
39 }
40 }
41 impl Deref for MistyResourceId {
42 type Target = u64;
43
44 fn deref(&self) -> &Self::Target {
45 &self.0
46 }
47 }
48};
49
50#[derive(Debug)]
51enum ToFlushResourceAction {
52 Insert(MistyResourceHandle),
53 Remove,
54}
55
56impl PartialEq for ToFlushResourceAction {
57 fn eq(&self, other: &Self) -> bool {
58 match (self, other) {
59 (Self::Remove, Self::Remove) => true,
60 (Self::Insert(l0), Self::Insert(r0)) => l0.id() == r0.id(),
61 _ => false,
62 }
63 }
64}
65
66struct MistyResourceManagerStore {
67 pending_actions: Arc<RwLock<Option<HashMap<MistyResourceId, ToFlushResourceAction>>>>,
68 weak_map: RwLock<HashMap<MistyResourceId, Weak<MistyResourceHandleInner>>>,
69}
70
71pub struct MistyResourceManager {
72 store: Arc<MistyResourceManagerStore>,
73}
74
75impl Debug for MistyResourceManager {
76 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
77 f.debug_struct("ResourceManager").finish()
78 }
79}
80
81#[derive(Debug, Clone)]
82struct MistyResourceHandleInner {
83 id: MistyResourceId,
84 store_ref: Weak<MistyResourceManagerStore>,
85 buf: Vec<u8>,
86}
87
88#[derive(Debug, Clone)]
89pub struct MistyResourceHandle {
90 ptr: Arc<MistyResourceHandleInner>,
91}
92
93impl MistyResourceHandle {
94 pub fn id(&self) -> MistyResourceId {
95 self.ptr.id
96 }
97
98 pub fn load(&self) -> &Vec<u8> {
99 &self.ptr.buf
100 }
101}
102
103impl Drop for MistyResourceHandleInner {
104 fn drop(&mut self) {
105 let store_ref = self.store_ref.upgrade();
106 if store_ref.is_none() {
107 return;
108 }
109 let store_ref = store_ref.unwrap();
110
111 {
112 let mut writter = store_ref.pending_actions.write().unwrap();
113 let writer = writter.as_mut().unwrap();
114 if writer.contains_key(&self.id) {
115 debug_assert!(writer.get(&self.id).unwrap() != &ToFlushResourceAction::Remove);
116 writer.remove(&self.id);
117 } else {
118 writer.insert(self.id, ToFlushResourceAction::Remove);
119 }
120 }
121 {
122 let mut writter = store_ref.weak_map.write().unwrap();
123 writter.remove(&self.id);
124 }
125 }
126}
127
128impl MistyResourceManager {
129 pub fn new() -> Self {
130 Self {
131 store: Arc::new(MistyResourceManagerStore {
132 pending_actions: Arc::new(RwLock::new(Some(Default::default()))),
133 weak_map: Default::default(),
134 }),
135 }
136 }
137
138 pub fn get_handle(&self, id: MistyResourceId) -> Option<MistyResourceHandle> {
139 let reader = self.store.weak_map.read().unwrap();
140 reader
141 .get(&id)
142 .map(|ptr| ptr.upgrade())
143 .unwrap_or_default()
144 .map(|ptr| MistyResourceHandle { ptr })
145 }
146
147 pub fn insert(&self, buf: impl Into<Vec<u8>> + 'static) -> MistyResourceHandle {
148 let id = MistyResourceId::alloc();
149 let buf: Vec<u8> = buf.into();
150
151 let ptr = Arc::new(MistyResourceHandleInner {
152 id,
153 store_ref: Arc::downgrade(&self.store),
154 buf,
155 });
156 let handle = MistyResourceHandle { ptr: ptr.clone() };
157
158 {
159 let mut writter = self.store.pending_actions.write().unwrap();
160 let writer = writter.as_mut().unwrap();
161 writer.insert(id, ToFlushResourceAction::Insert(handle.clone()));
162 }
163 {
164 let mut writter = self.store.weak_map.write().unwrap();
165 writter.insert(id, Arc::downgrade(&ptr));
166 }
167
168 handle
169 }
170
171 pub(crate) fn take_all_actions(&self) -> Vec<ResourceUpdateAction> {
172 let pending_ids = {
173 let mut writer = self.store.pending_actions.write().unwrap();
174 let cloned = writer.take();
175 *writer = Some(Default::default());
176 cloned.unwrap()
177 };
178
179 let mut ret: Vec<ResourceUpdateAction> = Default::default();
180 for (id, action) in pending_ids {
181 match action {
182 ToFlushResourceAction::Insert(handle) => {
183 ret.push(ResourceUpdateAction::Insert(id, handle.load().clone()));
184 }
185 ToFlushResourceAction::Remove => {
186 ret.push(ResourceUpdateAction::Remove(id));
187 }
188 }
189 }
190 ret
191 }
192}