oximedia_gpu/
buffer_pool.rs

1//! Zero-copy buffer pool for GPU-style memory management.
2//!
3//! Provides a reuse-oriented pool of byte buffers, inspired by GPU memory
4//! management patterns.  Buffers are acquired by size and alignment,
5//! used by the caller, and released back to the pool rather than freed.
6//! Unused buffers older than 60 seconds are evicted by [`BufferPool::defragment`].
7
8#![allow(clippy::cast_precision_loss)]
9
10use std::time::Instant;
11
12// ---------------------------------------------------------------------------
13// GpuBuffer
14// ---------------------------------------------------------------------------
15
16/// A raw byte buffer managed by a [`BufferPool`].
17pub struct GpuBuffer {
18    /// Unique identifier assigned by the owning pool.
19    pub id: u64,
20    /// Allocated capacity in bytes.
21    pub size_bytes: usize,
22    /// Alignment guarantee (in bytes).
23    pub alignment: usize,
24    /// Backing storage.
25    data: Vec<u8>,
26    /// Whether this buffer is currently checked out by a caller.
27    pub(crate) in_use: bool,
28    /// Monotonic timestamp of the most recent acquisition or release.
29    pub(crate) created_at: Instant,
30    /// Monotonic timestamp of last release (used for eviction).
31    pub(crate) last_released_at: Option<Instant>,
32}
33
34impl GpuBuffer {
35    /// Allocate a new buffer with the given `size` and `alignment`.
36    ///
37    /// The alignment hint is recorded but the backing `Vec<u8>` uses the
38    /// default allocator.  For truly aligned allocations a custom allocator
39    /// would be required; the pool still respects the alignment in
40    /// compatibility checks.
41    #[must_use]
42    pub fn new(id: u64, size: usize, alignment: usize) -> Self {
43        let effective_alignment = alignment.max(1);
44        Self {
45            id,
46            size_bytes: size,
47            alignment: effective_alignment,
48            data: vec![0u8; size],
49            in_use: false,
50            created_at: Instant::now(),
51            last_released_at: None,
52        }
53    }
54
55    /// View the buffer contents as a byte slice.
56    #[must_use]
57    pub fn as_slice(&self) -> &[u8] {
58        &self.data
59    }
60
61    /// View the buffer contents as a mutable byte slice.
62    #[must_use]
63    pub fn as_mut_slice(&mut self) -> &mut [u8] {
64        &mut self.data
65    }
66
67    /// Fill the entire buffer with `value` (memset equivalent).
68    pub fn fill(&mut self, value: u8) {
69        self.data.fill(value);
70    }
71
72    /// Whether this buffer is currently checked out.
73    #[must_use]
74    pub fn is_in_use(&self) -> bool {
75        self.in_use
76    }
77}
78
79impl std::fmt::Debug for GpuBuffer {
80    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
81        f.debug_struct("GpuBuffer")
82            .field("id", &self.id)
83            .field("size_bytes", &self.size_bytes)
84            .field("alignment", &self.alignment)
85            .field("in_use", &self.in_use)
86            .finish()
87    }
88}
89
90// ---------------------------------------------------------------------------
91// Pool statistics
92// ---------------------------------------------------------------------------
93
94/// Snapshot of pool health metrics.
95#[derive(Debug, Clone)]
96pub struct PoolStats {
97    /// Total number of buffers held by the pool (in-use + available).
98    pub total_buffers: usize,
99    /// Buffers currently checked out by callers.
100    pub in_use_buffers: usize,
101    /// Buffers available for immediate reuse.
102    pub available_buffers: usize,
103    /// Sum of all allocated buffer capacities in bytes.
104    pub total_allocated_bytes: usize,
105    /// Fraction of acquisitions satisfied from the pool (0.0 – 1.0).
106    pub reuse_rate: f64,
107}
108
109// ---------------------------------------------------------------------------
110// BufferPool
111// ---------------------------------------------------------------------------
112
113/// A pool of reusable GPU-style byte buffers.
114///
115/// Callers acquire a buffer via [`acquire`][BufferPool::acquire] (receiving its
116/// ID), read/write via [`get_mut`][BufferPool::get_mut], then return it to the
117/// pool with [`release`][BufferPool::release].
118pub struct BufferPool {
119    buffers: Vec<GpuBuffer>,
120    next_id: u64,
121    total_allocated: usize,
122    max_pool_bytes: usize,
123    /// Number of acquisitions satisfied by reusing an existing buffer.
124    reuse_count: u64,
125    /// Total acquisitions ever made (reused + newly allocated).
126    alloc_count: u64,
127}
128
129impl BufferPool {
130    /// Create a new pool that will hold at most `max_pool_bytes` of backing
131    /// storage before refusing new allocations.
132    #[must_use]
133    pub fn new(max_pool_bytes: usize) -> Self {
134        Self {
135            buffers: Vec::new(),
136            next_id: 1,
137            total_allocated: 0,
138            max_pool_bytes,
139            reuse_count: 0,
140            alloc_count: 0,
141        }
142    }
143
144    // -----------------------------------------------------------------------
145    // Acquire
146    // -----------------------------------------------------------------------
147
148    /// Check out a buffer of at least `size_bytes` with at least `alignment`.
149    ///
150    /// Strategy: find the *smallest* existing compatible free buffer to
151    /// minimise fragmentation.  If none exists, allocate a new one (provided
152    /// the pool is below its byte budget).
153    ///
154    /// Returns the buffer `id` on success, or `None` if no buffer is available
155    /// and allocating a new one would exceed the pool's byte budget.
156    pub fn acquire(&mut self, size_bytes: usize, alignment: usize) -> Option<u64> {
157        self.alloc_count += 1;
158
159        // Find the best (smallest compatible) free buffer.
160        let best_idx = self
161            .buffers
162            .iter()
163            .enumerate()
164            .filter(|(_, b)| {
165                !b.in_use && b.size_bytes >= size_bytes && b.alignment >= alignment.max(1)
166            })
167            .min_by_key(|(_, b)| b.size_bytes)
168            .map(|(idx, _)| idx);
169
170        if let Some(idx) = best_idx {
171            self.buffers[idx].in_use = true;
172            self.buffers[idx].created_at = Instant::now();
173            self.reuse_count += 1;
174            return Some(self.buffers[idx].id);
175        }
176
177        // No compatible free buffer — try to allocate a new one.
178        let effective_alignment = alignment.max(1);
179        let new_size = self.total_allocated + size_bytes;
180        if new_size > self.max_pool_bytes {
181            return None; // over budget
182        }
183
184        let id = self.next_id;
185        self.next_id += 1;
186
187        let mut buf = GpuBuffer::new(id, size_bytes, effective_alignment);
188        buf.in_use = true;
189        self.total_allocated += size_bytes;
190        self.buffers.push(buf);
191
192        Some(id)
193    }
194
195    // -----------------------------------------------------------------------
196    // Release
197    // -----------------------------------------------------------------------
198
199    /// Return a buffer to the pool by `id`.
200    ///
201    /// The buffer is kept for future reuse but marked as available.
202    /// Returns `true` if the buffer was found and released, `false` otherwise.
203    pub fn release(&mut self, id: u64) -> bool {
204        if let Some(buf) = self.buffers.iter_mut().find(|b| b.id == id) {
205            buf.in_use = false;
206            buf.last_released_at = Some(Instant::now());
207            true
208        } else {
209            false
210        }
211    }
212
213    // -----------------------------------------------------------------------
214    // Accessors
215    // -----------------------------------------------------------------------
216
217    /// Borrow the buffer with the given `id`.
218    #[must_use]
219    pub fn get(&self, id: u64) -> Option<&GpuBuffer> {
220        self.buffers.iter().find(|b| b.id == id)
221    }
222
223    /// Mutably borrow the buffer with the given `id`.
224    #[must_use]
225    pub fn get_mut(&mut self, id: u64) -> Option<&mut GpuBuffer> {
226        self.buffers.iter_mut().find(|b| b.id == id)
227    }
228
229    // -----------------------------------------------------------------------
230    // Defragmentation
231    // -----------------------------------------------------------------------
232
233    /// Evict all free buffers that have not been used for more than 60 seconds.
234    ///
235    /// In-use buffers are never evicted.
236    pub fn defragment(&mut self) {
237        let now = Instant::now();
238        let eviction_threshold = std::time::Duration::from_secs(60);
239
240        let mut bytes_freed = 0usize;
241        self.buffers.retain(|buf| {
242            if buf.in_use {
243                return true; // never evict live buffers
244            }
245            let idle_since = buf.last_released_at.unwrap_or(buf.created_at);
246            if now.duration_since(idle_since) > eviction_threshold {
247                bytes_freed += buf.size_bytes;
248                false // evict
249            } else {
250                true
251            }
252        });
253        self.total_allocated = self.total_allocated.saturating_sub(bytes_freed);
254    }
255
256    // -----------------------------------------------------------------------
257    // Stats
258    // -----------------------------------------------------------------------
259
260    /// Snapshot of pool metrics.
261    #[must_use]
262    pub fn stats(&self) -> PoolStats {
263        let in_use = self.buffers.iter().filter(|b| b.in_use).count();
264        let available = self.buffers.len() - in_use;
265        let reuse_rate = if self.alloc_count == 0 {
266            0.0
267        } else {
268            self.reuse_count as f64 / self.alloc_count as f64
269        };
270        PoolStats {
271            total_buffers: self.buffers.len(),
272            in_use_buffers: in_use,
273            available_buffers: available,
274            total_allocated_bytes: self.total_allocated,
275            reuse_rate,
276        }
277    }
278
279    /// Total bytes currently under management.
280    #[must_use]
281    pub fn total_allocated_bytes(&self) -> usize {
282        self.total_allocated
283    }
284
285    /// Maximum pool capacity in bytes.
286    #[must_use]
287    pub fn max_pool_bytes(&self) -> usize {
288        self.max_pool_bytes
289    }
290}
291
292impl std::fmt::Debug for BufferPool {
293    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
294        f.debug_struct("BufferPool")
295            .field("buffers", &self.buffers.len())
296            .field("total_allocated", &self.total_allocated)
297            .field("max_pool_bytes", &self.max_pool_bytes)
298            .field("alloc_count", &self.alloc_count)
299            .field("reuse_count", &self.reuse_count)
300            .finish()
301    }
302}
303
304// ---------------------------------------------------------------------------
305// Tests
306// ---------------------------------------------------------------------------
307
308#[cfg(test)]
309mod tests {
310    use super::*;
311
312    // --- GpuBuffer ---
313
314    #[test]
315    fn test_gpu_buffer_new() {
316        let buf = GpuBuffer::new(1, 1024, 64);
317        assert_eq!(buf.id, 1);
318        assert_eq!(buf.size_bytes, 1024);
319        assert_eq!(buf.alignment, 64);
320        assert_eq!(buf.as_slice().len(), 1024);
321        assert!(!buf.is_in_use());
322    }
323
324    #[test]
325    fn test_gpu_buffer_fill() {
326        let mut buf = GpuBuffer::new(2, 16, 4);
327        buf.fill(0xAB);
328        assert!(buf.as_slice().iter().all(|&b| b == 0xAB));
329    }
330
331    #[test]
332    fn test_gpu_buffer_as_mut_slice() {
333        let mut buf = GpuBuffer::new(3, 8, 1);
334        buf.as_mut_slice()[0] = 42;
335        assert_eq!(buf.as_slice()[0], 42);
336    }
337
338    // --- BufferPool::new ---
339
340    #[test]
341    fn test_pool_new_empty() {
342        let pool = BufferPool::new(1024 * 1024);
343        let stats = pool.stats();
344        assert_eq!(stats.total_buffers, 0);
345        assert_eq!(stats.reuse_rate, 0.0);
346    }
347
348    // --- acquire / release ---
349
350    #[test]
351    fn test_pool_acquire_and_release() {
352        let mut pool = BufferPool::new(1024 * 1024);
353        let id = pool.acquire(256, 4).expect("acquire failed");
354        assert!(pool.get(id).expect("missing").is_in_use());
355
356        let released = pool.release(id);
357        assert!(released, "release should succeed");
358        assert!(!pool.get(id).expect("missing").is_in_use());
359    }
360
361    #[test]
362    fn test_pool_reuse() {
363        let mut pool = BufferPool::new(1024 * 1024);
364        let id1 = pool.acquire(512, 4).expect("first acquire");
365        pool.release(id1);
366        let id2 = pool.acquire(512, 4).expect("second acquire");
367        // The pool should have reused the same buffer.
368        assert_eq!(id1, id2, "expected buffer reuse");
369        let stats = pool.stats();
370        assert!(stats.reuse_rate > 0.0);
371    }
372
373    #[test]
374    fn test_pool_smallest_compatible_preferred() {
375        let mut pool = BufferPool::new(4 * 1024 * 1024);
376        // Allocate two free buffers of different sizes.
377        let big = pool.acquire(4096, 4).expect("big");
378        let small = pool.acquire(256, 4).expect("small");
379        pool.release(big);
380        pool.release(small);
381        // Requesting 128 bytes: should get the 256-byte buffer (smallest compat).
382        let id = pool.acquire(128, 4).expect("reacquire");
383        assert_eq!(id, small, "should prefer smaller buffer");
384    }
385
386    #[test]
387    fn test_pool_budget_exceeded() {
388        let mut pool = BufferPool::new(100);
389        // First acquisition should succeed.
390        let id = pool.acquire(80, 1).expect("first");
391        // Second would exceed budget while first is in use.
392        let result = pool.acquire(80, 1);
393        assert!(result.is_none(), "should fail over budget");
394        pool.release(id);
395    }
396
397    #[test]
398    fn test_pool_release_unknown_id() {
399        let mut pool = BufferPool::new(1024);
400        assert!(
401            !pool.release(9999),
402            "releasing unknown id should return false"
403        );
404    }
405
406    #[test]
407    fn test_pool_get_missing() {
408        let pool = BufferPool::new(1024);
409        assert!(pool.get(42).is_none());
410    }
411
412    // --- get_mut ---
413
414    #[test]
415    fn test_pool_get_mut_write() {
416        let mut pool = BufferPool::new(1024 * 1024);
417        let id = pool.acquire(64, 1).expect("acquire");
418        {
419            let buf = pool.get_mut(id).expect("get_mut");
420            buf.as_mut_slice()[0] = 0xFF;
421        }
422        assert_eq!(pool.get(id).expect("get").as_slice()[0], 0xFF);
423    }
424
425    // --- stats ---
426
427    #[test]
428    fn test_pool_stats_in_use_count() {
429        let mut pool = BufferPool::new(1024 * 1024);
430        let id1 = pool.acquire(128, 1).expect("a1");
431        let _id2 = pool.acquire(128, 1).expect("a2");
432        pool.release(id1);
433        let stats = pool.stats();
434        assert_eq!(stats.total_buffers, 2);
435        assert_eq!(stats.in_use_buffers, 1);
436        assert_eq!(stats.available_buffers, 1);
437    }
438
439    // --- defragment ---
440
441    #[test]
442    fn test_pool_defragment_keeps_in_use() {
443        let mut pool = BufferPool::new(1024 * 1024);
444        let id = pool.acquire(64, 1).expect("acquire");
445        // Run defragment while buffer is in use — it should survive.
446        pool.defragment();
447        assert!(
448            pool.get(id).is_some(),
449            "in-use buffer should not be evicted"
450        );
451    }
452
453    #[test]
454    fn test_pool_defragment_recently_released_kept() {
455        let mut pool = BufferPool::new(1024 * 1024);
456        let id = pool.acquire(64, 1).expect("acquire");
457        pool.release(id);
458        // Buffer was just released — defragment should keep it (not 60s old).
459        pool.defragment();
460        assert!(
461            pool.get(id).is_some(),
462            "recently released buffer should survive"
463        );
464    }
465}
oximedia_gpu/buffer_pool.rs

oximedia_gpu/
buffer_pool.rs
