1use std::collections::{BTreeMap, BTreeSet};
2
3#[cfg(feature = "records")]
4use std::sync::{Arc, Mutex};
5
6#[cfg(feature = "records")]
7use spatialrust_records::{CancellationToken, MemoryBudget, MemoryReservation, MemoryTracker};
8
9use crate::{LodError, LodResult, NodeId};
10
11#[derive(Clone, Copy, Debug, PartialEq, Eq)]
13pub struct LodBudgets {
14 pub max_points: u64,
16 pub max_host_bytes: u64,
18 pub max_gpu_bytes: u64,
20 pub max_upload_bytes_per_frame: u64,
22 pub max_in_flight: usize,
24}
25
26impl LodBudgets {
27 pub fn validate(self) -> LodResult<()> {
29 if self.max_points == 0
30 || self.max_host_bytes == 0
31 || self.max_gpu_bytes == 0
32 || self.max_upload_bytes_per_frame == 0
33 || self.max_in_flight == 0
34 {
35 return Err(LodError::BudgetExceeded(
36 "all LOD resource limits must be positive".into(),
37 ));
38 }
39 Ok(())
40 }
41}
42
43#[derive(Clone, Copy, Debug, PartialEq, Eq)]
45pub struct ResidentNode {
46 pub id: NodeId,
48 pub point_count: u64,
50 pub gpu_bytes: u64,
52 pub last_used_generation: u64,
54}
55
56#[derive(Clone, Debug, PartialEq, Eq)]
58pub struct GpuCacheReceipt {
59 pub admitted: NodeId,
61 pub evicted: Vec<NodeId>,
63 pub resident_points: u64,
65 pub resident_gpu_bytes: u64,
67}
68
69#[derive(Clone, Debug)]
71pub struct LodGpuCache {
72 budgets: LodBudgets,
73 residents: BTreeMap<NodeId, ResidentNode>,
74 resident_points: u64,
75 resident_gpu_bytes: u64,
76}
77
78impl LodGpuCache {
79 pub fn try_new(budgets: LodBudgets) -> LodResult<Self> {
81 budgets.validate()?;
82 Ok(Self { budgets, residents: BTreeMap::new(), resident_points: 0, resident_gpu_bytes: 0 })
83 }
84
85 #[must_use]
87 pub fn resident_ids(&self) -> BTreeSet<NodeId> {
88 self.residents.keys().copied().collect()
89 }
90
91 #[must_use]
93 pub fn resident(&self, id: NodeId) -> Option<&ResidentNode> {
94 self.residents.get(&id)
95 }
96
97 pub fn touch(&mut self, id: NodeId, generation: u64) -> LodResult<()> {
99 let node = self.residents.get_mut(&id).ok_or(LodError::UnknownNode(id.0))?;
100 node.last_used_generation = generation;
101 Ok(())
102 }
103
104 pub fn admit(
108 &mut self,
109 id: NodeId,
110 point_count: u64,
111 gpu_bytes: u64,
112 generation: u64,
113 protected: &BTreeSet<NodeId>,
114 ) -> LodResult<GpuCacheReceipt> {
115 if point_count == 0
116 || gpu_bytes == 0
117 || point_count > self.budgets.max_points
118 || gpu_bytes > self.budgets.max_gpu_bytes
119 {
120 return Err(LodError::BudgetExceeded(format!(
121 "node {} cannot fit GPU point/byte budget",
122 id.0
123 )));
124 }
125 if self.residents.contains_key(&id) {
126 self.touch(id, generation)?;
127 return Ok(GpuCacheReceipt {
128 admitted: id,
129 evicted: Vec::new(),
130 resident_points: self.resident_points,
131 resident_gpu_bytes: self.resident_gpu_bytes,
132 });
133 }
134
135 let mut next_points = self
136 .resident_points
137 .checked_add(point_count)
138 .ok_or_else(|| LodError::BudgetExceeded("resident point count overflow".into()))?;
139 let mut next_bytes = self
140 .resident_gpu_bytes
141 .checked_add(gpu_bytes)
142 .ok_or_else(|| LodError::BudgetExceeded("resident GPU byte count overflow".into()))?;
143 let mut candidates: Vec<_> = self
144 .residents
145 .values()
146 .filter(|resident| !protected.contains(&resident.id))
147 .copied()
148 .collect();
149 candidates.sort_by_key(|resident| (resident.last_used_generation, resident.id));
150 let mut evicted = Vec::new();
151 for candidate in candidates {
152 if next_points <= self.budgets.max_points && next_bytes <= self.budgets.max_gpu_bytes {
153 break;
154 }
155 next_points -= candidate.point_count;
156 next_bytes -= candidate.gpu_bytes;
157 evicted.push(candidate.id);
158 }
159 if next_points > self.budgets.max_points || next_bytes > self.budgets.max_gpu_bytes {
160 return Err(LodError::BudgetExceeded(
161 "protected GPU residents leave insufficient capacity".into(),
162 ));
163 }
164 for evicted_id in &evicted {
165 self.residents.remove(evicted_id);
166 }
167 self.resident_points = next_points;
168 self.resident_gpu_bytes = next_bytes;
169 self.residents.insert(
170 id,
171 ResidentNode { id, point_count, gpu_bytes, last_used_generation: generation },
172 );
173 Ok(GpuCacheReceipt {
174 admitted: id,
175 evicted,
176 resident_points: next_points,
177 resident_gpu_bytes: next_bytes,
178 })
179 }
180
181 pub fn remove(&mut self, id: NodeId) -> Option<ResidentNode> {
183 let resident = self.residents.remove(&id)?;
184 self.resident_points -= resident.point_count;
185 self.resident_gpu_bytes -= resident.gpu_bytes;
186 Some(resident)
187 }
188}
189
190#[cfg(feature = "records")]
191#[derive(Debug, Default)]
192struct UploadState {
193 generation: u64,
194 admitted_upload_bytes: u64,
195 in_flight: BTreeSet<NodeId>,
196}
197
198#[cfg(feature = "records")]
200#[derive(Debug)]
201pub struct HostChunkLease {
202 node: NodeId,
203 point_count: u64,
204 upload_bytes: u64,
205 generation: u64,
206 cancellation: CancellationToken,
207 reservation: MemoryReservation,
208 state: Arc<Mutex<UploadState>>,
209 released: bool,
210}
211
212#[cfg(feature = "records")]
213impl HostChunkLease {
214 #[must_use]
216 pub const fn node(&self) -> NodeId {
217 self.node
218 }
219
220 #[must_use]
222 pub const fn host_bytes(&self) -> u64 {
223 self.reservation.bytes()
224 }
225
226 #[must_use]
228 pub fn cancellation_token(&self) -> CancellationToken {
229 self.cancellation.clone()
230 }
231
232 fn release(&mut self) {
233 if !self.released {
234 self.state.lock().expect("LOD upload state poisoned").in_flight.remove(&self.node);
235 self.released = true;
236 }
237 }
238}
239
240#[cfg(feature = "records")]
241impl Drop for HostChunkLease {
242 fn drop(&mut self) {
243 self.release();
244 }
245}
246
247#[cfg(feature = "records")]
249#[derive(Clone, Debug, PartialEq, Eq)]
250pub struct LodUploadReceipt {
251 pub node: NodeId,
253 pub generation: u64,
255 pub host_bytes: u64,
257 pub upload_bytes: u64,
259 pub gpu_bytes: u64,
261 pub evicted: Vec<NodeId>,
263 pub peak_host_bytes: u64,
265}
266
267#[cfg(feature = "records")]
269#[derive(Clone, Debug)]
270pub struct LodUploadSession {
271 budgets: LodBudgets,
272 memory: MemoryTracker,
273 state: Arc<Mutex<UploadState>>,
274}
275
276#[cfg(feature = "records")]
277impl LodUploadSession {
278 pub fn try_new(budgets: LodBudgets) -> LodResult<Self> {
280 budgets.validate()?;
281 let memory_budget = MemoryBudget::new(budgets.max_host_bytes)
282 .map_err(|error| LodError::Records(error.to_string()))?;
283 Ok(Self {
284 budgets,
285 memory: MemoryTracker::new(memory_budget),
286 state: Arc::new(Mutex::new(UploadState::default())),
287 })
288 }
289
290 pub fn begin_frame(&self, generation: u64) -> LodResult<()> {
292 let mut state = self.state.lock().expect("LOD upload state poisoned");
293 if generation <= state.generation {
294 return Err(LodError::InvalidPlanner(
295 "upload generations must increase monotonically".into(),
296 ));
297 }
298 state.generation = generation;
299 state.admitted_upload_bytes = 0;
300 Ok(())
301 }
302
303 pub fn try_lease(
305 &self,
306 node: NodeId,
307 point_count: u64,
308 host_bytes: u64,
309 upload_bytes: u64,
310 cancellation: CancellationToken,
311 ) -> LodResult<HostChunkLease> {
312 if point_count == 0 || host_bytes == 0 || upload_bytes == 0 {
313 return Err(LodError::BudgetExceeded(
314 "LOD lease counts and bytes must be positive".into(),
315 ));
316 }
317 let mut state = self.state.lock().expect("LOD upload state poisoned");
318 if state.in_flight.contains(&node) {
319 return Err(LodError::BudgetExceeded(format!("node {} is already in flight", node.0)));
320 }
321 if state.in_flight.len() >= self.budgets.max_in_flight {
322 return Err(LodError::BudgetExceeded("in-flight chunk budget exhausted".into()));
323 }
324 let next_upload = state
325 .admitted_upload_bytes
326 .checked_add(upload_bytes)
327 .ok_or_else(|| LodError::BudgetExceeded("frame upload byte overflow".into()))?;
328 if next_upload > self.budgets.max_upload_bytes_per_frame {
329 return Err(LodError::BudgetExceeded("per-frame upload byte budget exhausted".into()));
330 }
331 let reservation = self
332 .memory
333 .try_reserve(host_bytes)
334 .map_err(|error| LodError::Records(error.to_string()))?;
335 state.in_flight.insert(node);
336 state.admitted_upload_bytes = next_upload;
337 Ok(HostChunkLease {
338 node,
339 point_count,
340 upload_bytes,
341 generation: state.generation,
342 cancellation,
343 reservation,
344 state: Arc::clone(&self.state),
345 released: false,
346 })
347 }
348
349 pub fn complete_upload(
351 &self,
352 mut lease: HostChunkLease,
353 gpu_bytes: u64,
354 cache: &mut LodGpuCache,
355 protected: &BTreeSet<NodeId>,
356 ) -> LodResult<LodUploadReceipt> {
357 lease.cancellation.check().map_err(|error| LodError::Records(error.to_string()))?;
358 if gpu_bytes == 0 {
359 return Err(LodError::BudgetExceeded("resident GPU bytes must be positive".into()));
360 }
361 let cache_receipt =
362 cache.admit(lease.node, lease.point_count, gpu_bytes, lease.generation, protected)?;
363 let receipt = LodUploadReceipt {
364 node: lease.node,
365 generation: lease.generation,
366 host_bytes: lease.reservation.bytes(),
367 upload_bytes: lease.upload_bytes,
368 gpu_bytes,
369 evicted: cache_receipt.evicted,
370 peak_host_bytes: self.memory.snapshot().peak_bytes,
371 };
372 lease.release();
373 Ok(receipt)
374 }
375
376 #[must_use]
378 pub fn in_flight(&self) -> BTreeSet<NodeId> {
379 self.state.lock().expect("LOD upload state poisoned").in_flight.clone()
380 }
381
382 #[must_use]
384 pub fn host_memory(&self) -> spatialrust_records::MemorySnapshot {
385 self.memory.snapshot()
386 }
387}
388
389#[cfg(test)]
390mod tests {
391 use std::collections::BTreeSet;
392
393 use crate::{LodBudgets, LodGpuCache, NodeId};
394
395 fn budgets() -> LodBudgets {
396 LodBudgets {
397 max_points: 100,
398 max_host_bytes: 100,
399 max_gpu_bytes: 100,
400 max_upload_bytes_per_frame: 100,
401 max_in_flight: 2,
402 }
403 }
404
405 #[test]
406 fn gpu_cache_evicts_deterministic_lru_and_preserves_protected_nodes() {
407 let mut cache = LodGpuCache::try_new(budgets()).unwrap();
408 cache.admit(NodeId(2), 40, 40, 1, &BTreeSet::new()).unwrap();
409 cache.admit(NodeId(1), 40, 40, 1, &BTreeSet::new()).unwrap();
410 cache.touch(NodeId(2), 2).unwrap();
411 let receipt = cache.admit(NodeId(3), 50, 50, 3, &BTreeSet::new()).unwrap();
412 assert_eq!(receipt.evicted, vec![NodeId(1)]);
413 assert_eq!(cache.resident_ids(), BTreeSet::from([NodeId(2), NodeId(3)]));
414
415 let before = cache.resident_ids();
416 assert!(cache
417 .admit(NodeId(4), 90, 90, 4, &BTreeSet::from([NodeId(2), NodeId(3)]))
418 .is_err());
419 assert_eq!(cache.resident_ids(), before);
420 }
421
422 #[cfg(feature = "records")]
423 #[test]
424 fn leases_enforce_memory_upload_inflight_cancellation_and_cleanup() {
425 let session = super::LodUploadSession::try_new(budgets()).unwrap();
426 session.begin_frame(1).unwrap();
427 let first = session
428 .try_lease(NodeId(1), 20, 60, 60, spatialrust_records::CancellationToken::default())
429 .unwrap();
430 assert!(session
431 .try_lease(NodeId(2), 20, 50, 20, spatialrust_records::CancellationToken::default())
432 .is_err());
433 assert_eq!(session.host_memory().current_bytes, 60);
434 drop(first);
435 assert_eq!(session.host_memory().current_bytes, 0);
436 assert!(session.in_flight().is_empty());
437
438 session.begin_frame(2).unwrap();
439 let cancellation = spatialrust_records::CancellationToken::default();
440 let lease = session.try_lease(NodeId(3), 20, 40, 40, cancellation.clone()).unwrap();
441 cancellation.cancel();
442 let mut cache = LodGpuCache::try_new(budgets()).unwrap();
443 assert!(session.complete_upload(lease, 40, &mut cache, &BTreeSet::new()).is_err());
444 assert!(cache.resident_ids().is_empty());
445 assert!(session.in_flight().is_empty());
446
447 session.begin_frame(3).unwrap();
448 let lease = session
449 .try_lease(NodeId(4), 20, 40, 40, spatialrust_records::CancellationToken::default())
450 .unwrap();
451 let receipt = session.complete_upload(lease, 40, &mut cache, &BTreeSet::new()).unwrap();
452 assert_eq!(receipt.host_bytes, 40);
453 assert_eq!(receipt.upload_bytes, 40);
454 assert_eq!(session.host_memory().current_bytes, 0);
455 assert_eq!(cache.resident_ids(), BTreeSet::from([NodeId(4)]));
456 }
457}