Skip to main content

spatialrust_lod/
residency.rs

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/// Hard limits shared by selection, decode/upload admission, and GPU residency.
12#[derive(Clone, Copy, Debug, PartialEq, Eq)]
13pub struct LodBudgets {
14    /// Maximum selected/resident points.
15    pub max_points: u64,
16    /// Maximum simultaneously leased decoded host bytes.
17    pub max_host_bytes: u64,
18    /// Maximum resident GPU bytes.
19    pub max_gpu_bytes: u64,
20    /// Maximum upload bytes admitted in one frame generation.
21    pub max_upload_bytes_per_frame: u64,
22    /// Maximum concurrently leased/in-flight chunks.
23    pub max_in_flight: usize,
24}
25
26impl LodBudgets {
27    /// Requires every hard limit to be positive.
28    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/// One device-resident node tracked without owning backend-specific handles.
44#[derive(Clone, Copy, Debug, PartialEq, Eq)]
45pub struct ResidentNode {
46    /// Node identity.
47    pub id: NodeId,
48    /// Resident points.
49    pub point_count: u64,
50    /// Resident allocation bytes.
51    pub gpu_bytes: u64,
52    /// Last frame generation that displayed/touched this node.
53    pub last_used_generation: u64,
54}
55
56/// Exact result of one GPU cache admission.
57#[derive(Clone, Debug, PartialEq, Eq)]
58pub struct GpuCacheReceipt {
59    /// Newly admitted node.
60    pub admitted: NodeId,
61    /// Nodes explicitly evicted before admission.
62    pub evicted: Vec<NodeId>,
63    /// Resident points after admission.
64    pub resident_points: u64,
65    /// Resident GPU bytes after admission.
66    pub resident_gpu_bytes: u64,
67}
68
69/// Backend-neutral GPU residency ledger with deterministic LRU eviction.
70#[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    /// Creates an empty bounded cache.
80    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    /// Current resident IDs.
86    #[must_use]
87    pub fn resident_ids(&self) -> BTreeSet<NodeId> {
88        self.residents.keys().copied().collect()
89    }
90
91    /// Finds a resident.
92    #[must_use]
93    pub fn resident(&self, id: NodeId) -> Option<&ResidentNode> {
94        self.residents.get(&id)
95    }
96
97    /// Marks a resident as used by `generation`.
98    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    /// Admits a completed explicit upload, evicting only unprotected LRU nodes.
105    ///
106    /// Failure leaves the cache unchanged.
107    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    /// Explicitly removes a node.
182    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/// Explicit host-memory lease for one decoded LOD chunk.
199#[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    /// Node identity.
215    #[must_use]
216    pub const fn node(&self) -> NodeId {
217        self.node
218    }
219
220    /// Exact reserved host bytes.
221    #[must_use]
222    pub const fn host_bytes(&self) -> u64 {
223        self.reservation.bytes()
224    }
225
226    /// Cooperative cancellation token checked before upload completion.
227    #[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/// Exact receipt for a host lease followed by caller-performed GPU upload.
248#[cfg(feature = "records")]
249#[derive(Clone, Debug, PartialEq, Eq)]
250pub struct LodUploadReceipt {
251    /// Uploaded node.
252    pub node: NodeId,
253    /// Plan/frame generation.
254    pub generation: u64,
255    /// Decoded host bytes held until upload completion.
256    pub host_bytes: u64,
257    /// Caller-declared bytes transferred to GPU.
258    pub upload_bytes: u64,
259    /// Caller-declared resident GPU bytes.
260    pub gpu_bytes: u64,
261    /// Nodes evicted to fit.
262    pub evicted: Vec<NodeId>,
263    /// Peak host lease bytes observed by the shared tracker.
264    pub peak_host_bytes: u64,
265}
266
267/// Bounded lease/upload admission state.
268#[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    /// Creates an empty upload session using the exact records memory tracker.
279    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    /// Starts a new frame/generation and resets only the per-frame upload ledger.
291    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    /// Reserves host memory, in-flight capacity, and this frame's upload bytes.
304    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    /// Records a caller-completed upload and admits it to the GPU cache.
350    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    /// Current in-flight IDs.
377    #[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    /// Current and peak host lease bytes.
383    #[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}