hrx/
buffer_pool.rs

1//! Bounded automatic reuse of temporary device buffers on one ordered stream.
2
3use crate::{Buffer, Error, Result, Stream, View};
4use std::{
5    collections::BTreeMap,
6    sync::{Arc, Mutex},
7};
8
9const DEFAULT_LIMIT: usize = 512 << 20;
10
11#[derive(Default)]
12struct FreeList {
13    cached: usize,
14    blocks: BTreeMap<usize, Vec<Buffer>>,
15    stream: Option<usize>,
16}
17
18/// A bounded free list whose allocations are safe to reuse on one ordered stream.
19pub struct BufferPool {
20    limit: usize,
21    free: Mutex<FreeList>,
22}
23
24impl Default for BufferPool {
25    fn default() -> Self {
26        Self::with_limit(DEFAULT_LIMIT)
27    }
28}
29
30impl BufferPool {
31    /// Create a pool with the default 512 MiB retained-byte limit.
32    #[must_use]
33    pub fn new() -> Arc<Self> {
34        Arc::new(Self::default())
35    }
36
37    /// Create a pool with an explicit retained-byte limit.
38    #[must_use]
39    pub fn with_limit(limit: usize) -> Self {
40        Self {
41            limit,
42            free: Mutex::new(FreeList::default()),
43        }
44    }
45
46    /// Acquire at least `bytes`, reusing a block no larger than twice the request.
47    pub fn acquire(self: &Arc<Self>, stream: &Stream, bytes: usize) -> Result<PooledBuffer> {
48        if bytes == 0 {
49            return Err(Error::Message("empty pooled allocation".into()));
50        }
51        let mut free = self.free.lock().unwrap_or_else(|error| error.into_inner());
52        match free.stream {
53            Some(owner) if owner != stream.id() => {
54                return Err(Error::Message(
55                    "a buffer pool serves one stream because reuse follows that queue's order"
56                        .into(),
57                ));
58            }
59            Some(_) => {}
60            None => free.stream = Some(stream.id()),
61        }
62        let reusable = free
63            .blocks
64            .range_mut(bytes..=bytes.saturating_mul(2))
65            .next()
66            .map(|(&size, blocks)| (size, blocks.pop()));
67        if let Some((size, Some(buffer))) = reusable {
68            free.cached -= size;
69            if free.blocks[&size].is_empty() {
70                free.blocks.remove(&size);
71            }
72            return Ok(PooledBuffer {
73                buffer: Some(buffer),
74                pool: self.clone(),
75            });
76        }
77        drop(free);
78        Ok(PooledBuffer {
79            buffer: Some(stream.allocate(bytes)?),
80            pool: self.clone(),
81        })
82    }
83
84    /// Bytes currently retained for reuse.
85    #[must_use]
86    pub fn cached_bytes(&self) -> usize {
87        self.free
88            .lock()
89            .unwrap_or_else(|error| error.into_inner())
90            .cached
91    }
92
93    fn release(&self, buffer: Buffer) {
94        let mut free = self.free.lock().unwrap_or_else(|error| error.into_inner());
95        if buffer.bytes() > self.limit.saturating_sub(free.cached) {
96            return;
97        }
98        free.cached += buffer.bytes();
99        free.blocks.entry(buffer.bytes()).or_default().push(buffer);
100    }
101}
102
103/// Device storage returned to its [`BufferPool`] when the lease drops.
104pub struct PooledBuffer {
105    buffer: Option<Buffer>,
106    pool: Arc<BufferPool>,
107}
108
109impl PooledBuffer {
110    /// Borrow the allocation.
111    #[must_use]
112    pub fn buffer(&self) -> &Buffer {
113        self.buffer.as_ref().expect("live pooled buffer")
114    }
115
116    /// Borrow the whole allocation as a binding.
117    #[must_use]
118    pub fn binding(&self) -> View<'_> {
119        self.buffer().binding()
120    }
121}
122
123impl Drop for PooledBuffer {
124    fn drop(&mut self) {
125        if let Some(buffer) = self.buffer.take() {
126            self.pool.release(buffer);
127        }
128    }
129}
hrx/buffer_pool.rs

hrx/
buffer_pool.rs
