rafx_framework/resources/
pool.rs

1use rafx_api::{RafxDescriptorSetArray, RafxDeviceContext, RafxResult};
2use std::collections::VecDeque;
3use std::num::Wrapping;
4
5// The u32 is the index of the pool being created (i.e. increases by 1 every time callback is invoked)
6pub type PoolResourceAllocatorAllocFn<T> =
7    dyn Fn(&RafxDeviceContext, u32) -> RafxResult<T> + Send + Sync;
8
9/// Implement to customize how PoolAllocator resets and destroys pools
10pub trait PooledResourceImpl {
11    fn reset(&mut self) -> RafxResult<()>;
12}
13
14struct PoolResourceInFlight<T: PooledResourceImpl> {
15    pool: T,
16    live_until_frame: Wrapping<u32>,
17}
18
19/// This handles waiting for N frames to pass before resetting the pool. "Restting" could mean
20/// different things depending on the resource. This allocator also has a callback for allocating
21/// new pools for use. A maximum pool count should be provided so that an unbounded leak of pools
22/// can be detected.
23pub struct PooledResourceAllocator<T: PooledResourceImpl> {
24    device_context: RafxDeviceContext,
25
26    // Allocates a new pool
27    allocate_fn: Box<PoolResourceAllocatorAllocFn<T>>,
28
29    // We are assuming that all pools can survive for the same amount of time so the data in
30    // this VecDeque will naturally be orderered such that things that need to be reset sooner
31    // are at the front
32    in_flight_pools: VecDeque<PoolResourceInFlight<T>>,
33
34    // Pools that have been reset and are ready for allocation
35    reset_pools: Vec<T>,
36
37    // All pools that are retired will be reset after N frames
38    max_in_flight_frames: Wrapping<u32>,
39
40    // Incremented when on_frame_complete is called
41    frame_index: Wrapping<u32>,
42
43    // Number of pools we have created in total
44    created_pool_count: u32,
45
46    // Max number of pools to create (sum includes allocated pools, pools in flight, and reset pools
47    max_pool_count: u32,
48}
49
50impl<T: PooledResourceImpl> PooledResourceAllocator<T> {
51    /// Create a pool allocator that will reset resources after N frames. Keep in mind that if for
52    /// example you want to push a single resource per frame, up to N+1 resources will exist
53    /// in the sink. If max_in_flight_frames is 2, then you would have a resource that has
54    /// likely not been submitted to the GPU yet, plus a resource per the N frames that have
55    /// been submitted
56    pub fn new<F: Fn(&RafxDeviceContext, u32) -> RafxResult<T> + Send + Sync + 'static>(
57        device_context: &RafxDeviceContext,
58        max_in_flight_frames: u32,
59        max_pool_count: u32,
60        allocate_fn: F,
61    ) -> Self {
62        PooledResourceAllocator {
63            device_context: device_context.clone(),
64            allocate_fn: Box::new(allocate_fn),
65            in_flight_pools: Default::default(),
66            reset_pools: Default::default(),
67            max_in_flight_frames: Wrapping(max_in_flight_frames),
68            frame_index: Wrapping(0),
69            created_pool_count: 0,
70            max_pool_count,
71        }
72    }
73
74    /// Allocate a pool - either reusing an old one that has been reset or creating a new one. Will
75    /// assert that we do not exceed max_pool_count. The pool is allowed to exist until retire_pool
76    /// is called. After this point, we will wait for N frames before restting it.
77    pub fn allocate_pool(&mut self) -> RafxResult<T> {
78        self.reset_pools.pop().map(Ok).unwrap_or_else(|| {
79            let pool_index = self.created_pool_count;
80            self.created_pool_count += 1;
81            assert!(self.created_pool_count <= self.max_pool_count);
82            (self.allocate_fn)(&self.device_context, pool_index)
83        })
84    }
85
86    /// Schedule the pool to reset after we complete N frames
87    pub fn retire_pool(
88        &mut self,
89        pool: T,
90    ) {
91        self.in_flight_pools.push_back(PoolResourceInFlight {
92            pool,
93            live_until_frame: self.frame_index + self.max_in_flight_frames,
94        });
95    }
96
97    /// Call when we are ready to reset another set of resources, most likely when a frame is
98    /// presented or a new frame begins
99    pub fn update(&mut self) -> RafxResult<()> {
100        self.frame_index += Wrapping(1);
101
102        // Determine how many pools we can drain
103        let mut pools_to_drain = 0;
104        for in_flight_pool in &self.in_flight_pools {
105            // If frame_index matches or exceeds live_until_frame, then the result will be a very
106            // high value due to wrapping a negative value to u32::MAX
107            if in_flight_pool.live_until_frame - self.frame_index > Wrapping(std::u32::MAX / 2) {
108                pools_to_drain += 1;
109            } else {
110                break;
111            }
112        }
113
114        // Reset them and add them to the list of pools ready to be allocated
115        let pools_to_reset: Vec<_> = self.in_flight_pools.drain(0..pools_to_drain).collect();
116        for mut pool_to_reset in pools_to_reset {
117            T::reset(&mut pool_to_reset.pool)?;
118            self.reset_pools.push(pool_to_reset.pool);
119        }
120
121        Ok(())
122    }
123
124    /// Immediately destroy everything. We assume the device is idle and nothing is in flight.
125    /// Calling this function when the device is not idle could result in a deadlock
126    pub fn destroy(&mut self) -> RafxResult<()> {
127        for pool in self.in_flight_pools.drain(..) {
128            std::mem::drop(pool.pool);
129        }
130
131        for pool in self.reset_pools.drain(..) {
132            std::mem::drop(pool);
133        }
134
135        Ok(())
136    }
137}
138
139// We assume destroy was called
140impl<T: PooledResourceImpl> Drop for PooledResourceAllocator<T> {
141    fn drop(&mut self) {
142        assert!(self.in_flight_pools.is_empty());
143        assert!(self.reset_pools.is_empty())
144    }
145}
146
147//
148// Implementation for descriptor pools
149//
150impl PooledResourceImpl for RafxDescriptorSetArray {
151    fn reset(&mut self) -> RafxResult<()> {
152        Ok(())
153    }
154}
155
156pub type DescriptorSetArrayPoolAllocator = PooledResourceAllocator<RafxDescriptorSetArray>;
rafx_framework/resources/pool.rs

rafx_framework/resources/
pool.rs
