1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362
use crate as wgpu;
use std::fmt;
use std::future::Future;
use std::sync::atomic::{self, AtomicU32};
use std::sync::{Arc, Mutex};
use std::time::Duration;
use image::{GenericImage, GenericImageView};
/// A type dedicated to capturing a texture as a non-linear sRGBA image that can be read on the
/// CPU.
///
/// Calling **capture** will return a **Snapshot** that may be read after the given command encoder
/// has been submitted. **Snapshot**s can be read on the current thread via **read** or on a thread
/// pool via **read_threaded**.
///
/// If the **Capturer** is dropped while threaded callbacks are still being processed, the drop
/// implementation will block the current thread.
#[derive(Debug, Default)]
pub struct Capturer {
converter_data_pair: Mutex<Option<ConverterDataPair>>,
thread_pool: Arc<Mutex<Option<Arc<ThreadPool>>>>,
workers: Option<u32>,
timeout: Option<Duration>,
}
/// A wrapper around the futures thread pool that counts active futures.
#[derive(Debug)]
struct ThreadPool {
active_futures: Arc<AtomicU32>,
workers: u32,
timeout: Option<Duration>,
}
/// A snapshot captured by a **Capturer**.
///
/// A snapshot is a thin wrapper around a **wgpu::BufferImage** that knows that the image format is
/// specifically non-linear sRGBA8.
pub struct Snapshot {
buffer: wgpu::RowPaddedBuffer,
thread_pool: Arc<Mutex<Option<Arc<ThreadPool>>>>,
workers: Option<u32>,
timeout: Option<Duration>,
}
/// An error indicating that the threadpool timed out while waiting for a worker to become
/// available.
pub struct AwaitWorkerTimeout<F>(pub F);
#[derive(Debug)]
struct ConverterDataPair {
src_descriptor: wgpu::TextureDescriptor<'static>,
reshaper: wgpu::TextureReshaper,
dst_texture: wgpu::Texture,
}
/// A wrapper around a slice of bytes representing a non-linear sRGBA image.
///
/// Can be read from a captured `Snapshot`.
pub struct Rgba8AsyncMappedImageBuffer<'buffer>(wgpu::ImageReadMapping<'buffer>);
impl ThreadPool {
/// Spawns the given future if a worker is available. Otherwise, blocks and waits for a worker
/// to become available before spawning the future.
fn spawn_when_worker_available<F>(&self, future: F) -> Result<(), AwaitWorkerTimeout<F>>
where
F: 'static + Future<Output = ()> + Send,
{
// Wait until the number of active futures is less than the number of threads.
// If we don't wait, the capture futures may quickly fall far behind the main
// swapchain thread resulting in an out of memory error.
let mut start = None;
let mut interval_us = 128;
while self.active_futures() >= self.workers() {
if let Some(timeout) = self.timeout {
let start = start.get_or_insert_with(instant::Instant::now);
if start.elapsed() > timeout {
return Err(AwaitWorkerTimeout(future));
}
}
let duration = Duration::from_micros(interval_us);
std::thread::sleep(duration);
interval_us *= 2;
}
// Wrap the future with the counter.
let active_futures = self.active_futures.clone();
let future = async move {
active_futures.fetch_add(1, atomic::Ordering::SeqCst);
future.await;
active_futures.fetch_sub(1, atomic::Ordering::SeqCst);
};
tokio::spawn(future);
Ok(())
}
fn active_futures(&self) -> u32 {
self.active_futures.load(atomic::Ordering::SeqCst)
}
fn workers(&self) -> u32 {
self.workers
}
/// Await for the completion of all active futures, polling the device as necessary until all
/// futures have completed.
fn await_active_futures(&self, device: &wgpu::Device) -> Result<(), AwaitWorkerTimeout<()>> {
let mut start = None;
let mut interval_us = 128;
while self.active_futures() > 0 {
if let Some(timeout) = self.timeout {
let start = start.get_or_insert_with(instant::Instant::now);
if start.elapsed() > timeout {
return Err(AwaitWorkerTimeout(()));
}
}
device.poll(wgpu::Maintain::Wait);
let duration = Duration::from_micros(interval_us);
std::thread::sleep(duration);
interval_us *= 2;
}
Ok(())
}
}
impl Capturer {
/// The format to which textures will be converted before being mapped back to the CPU.
pub const DST_FORMAT: wgpu::TextureFormat = wgpu::TextureFormat::Rgba8UnormSrgb;
/// Create a new **TextureCapturer**.
///
/// Note that a **TextureCapturer** must only be used with a single texture. If you require
/// capturing multiple textures, you may create multiple **TextureCapturers**.
///
/// `workers` refers to the number of worker threads used to await GPU buffers to be mapped for
/// reading and for running user callbacks. If `None` is specified, a threadpool will be
/// spawned with a number of threads equal to the number of CPUs available on the system.
///
/// `timeout` specifies how long to block and wait for an available worker in the case that all
/// workers are busy at the time a `Snapshot::read` occurs. If `None` is specified, calls to
/// `Snapshot::read` will never time out (the default behaviour).
///
/// Note that the specified parameters are only relevant to calls to `Snapshot::read`. In the
/// case that the user uses `Snapshot::read_async`, it is the responsibility of the user to
/// execute the future.
pub fn new(workers: Option<u32>, timeout: Option<Duration>) -> Self {
Capturer {
converter_data_pair: Default::default(),
thread_pool: Default::default(),
workers,
timeout,
}
}
/// The number of futures currently running on the inner `ThreadPool`.
///
/// Note that futures are only run on the threadpool when the `Snapshot::read` method is used.
/// In the case that `Snapshot::read_async` is used it is up to the user to track their
/// futures.
///
/// If the inner thread pool mutex has been poisoned, or if the thread pool has not been
/// created due to no calls to `read`, this will return `0`.
pub fn active_snapshots(&self) -> u32 {
if let Ok(guard) = self.thread_pool.lock() {
if let Some(tp) = guard.as_ref() {
return tp.active_futures.load(atomic::Ordering::SeqCst);
}
}
0
}
/// The number of worker threads used to await GPU buffers to be mapped for reading and for
/// running user callbacks.
pub fn workers(&self) -> u32 {
if let Ok(guard) = self.thread_pool.lock() {
if let Some(tp) = guard.as_ref() {
return tp.workers();
}
}
self.workers.unwrap_or(num_cpus::get() as u32)
}
/// Capture the given texture at the state of the given command encoder.
pub fn capture(
&self,
device: &wgpu::Device,
encoder: &mut wgpu::CommandEncoder,
src_texture: &wgpu::Texture,
) -> Snapshot {
let buffer = if src_texture.format() != Self::DST_FORMAT {
let mut converter_data_pair = self
.converter_data_pair
.lock()
.expect("failed to lock converter");
// Create converter and target texture if they don't exist.
let converter_data_pair = converter_data_pair
.get_or_insert_with(|| create_converter_data_pair(device, src_texture));
// If the texture has changed in some way, recreate the converter.
if !wgpu::texture_descriptor_eq(
src_texture.descriptor(),
&converter_data_pair.src_descriptor,
) {
*converter_data_pair = create_converter_data_pair(device, src_texture);
}
// Encode the texture format conversion.
let dst_view = converter_data_pair.dst_texture.view();
converter_data_pair
.reshaper
.encode_render_pass(&dst_view.build(), encoder);
converter_data_pair.dst_texture.to_buffer(device, encoder)
} else {
src_texture.to_buffer(device, encoder)
};
Snapshot {
buffer,
thread_pool: self.thread_pool.clone(),
workers: self.workers,
timeout: self.timeout,
}
}
/// Await for the completion of all `Snapshot::read` active futures, polling the device as
/// necessary until all futures have reached completion or until a timeout is reached.
pub fn await_active_snapshots(
&self,
device: &wgpu::Device,
) -> Result<(), AwaitWorkerTimeout<()>> {
if let Ok(guard) = self.thread_pool.lock() {
if let Some(tp) = guard.as_ref() {
return tp.await_active_futures(device);
}
}
Ok(())
}
}
impl Snapshot {
/// Reads the non-linear sRGBA image from mapped memory and convert it to an owned buffer.
pub async fn read_async<'buffer>(
&'buffer self,
) -> Result<Rgba8AsyncMappedImageBuffer<'buffer>, wgpu::BufferAsyncError> {
let mapping = self.buffer.read().await?;
Ok(Rgba8AsyncMappedImageBuffer(mapping))
}
/// The same as `read_async`, but runs the resulting future on an inner threadpool and calls
/// the given callback with the mapped image buffer once complete.
///
/// Note: The given callback will not be called until the memory is mapped and the device is
/// polled. You should not rely on the callback being called immediately.
///
/// Note: The given callback will be called on the inner thread pool and will not be called on
/// the current thread.
///
/// Note: **This method may block** if the associated `wgpu::TextureCapturer` has an
/// `active_futures` count that is greater than the number of worker threads with which it was
/// created. This is necessary in order to avoid "out of memory" errors resulting from an
/// accumulating queue of pending texture buffers waiting to be mapped. To avoid blocking, you
/// can try using a higher thread count, capturing a smaller texture, or using `read_async`
/// instead and running the resulting future on a custom runtime or threadpool.
pub fn read<F>(self, callback: F) -> Result<(), AwaitWorkerTimeout<impl Future<Output = ()>>>
where
F: 'static + Send + FnOnce(Result<Rgba8AsyncMappedImageBuffer, wgpu::BufferAsyncError>),
{
let thread_pool = self.thread_pool();
let read_future = async move {
let res = self.read_async().await;
callback(res);
};
thread_pool.spawn_when_worker_available(read_future)
}
fn thread_pool(&self) -> Arc<ThreadPool> {
let mut guard = self
.thread_pool
.lock()
.expect("failed to acquire thread handle");
let thread_pool = guard.get_or_insert_with(|| {
let workers = self.workers.unwrap_or(num_cpus::get() as u32);
let thread_pool = ThreadPool {
active_futures: Arc::new(AtomicU32::new(0)),
workers,
timeout: self.timeout,
};
Arc::new(thread_pool)
});
thread_pool.clone()
}
}
impl<'b> Rgba8AsyncMappedImageBuffer<'b> {
pub fn as_image(&self) -> image::SubImage<wgpu::ImageHolder<image::Rgba<u8>>> {
// safe: we know it's Rgba<u8>
unsafe { self.0.as_image::<image::Rgba<u8>>() }
}
/// Convert the mapped image buffer to an owned buffer.
pub fn to_owned(&self) -> image::ImageBuffer<image::Rgba<u8>, Vec<u8>> {
let view = self.as_image();
let mut result = image::ImageBuffer::new(view.width(), view.height());
result
.copy_from(&view, 0, 0)
.expect("nannou internal error: image copy failed");
result
}
}
impl<T> std::error::Error for AwaitWorkerTimeout<T> {}
impl<T> fmt::Debug for AwaitWorkerTimeout<T> {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
f.debug_struct("AwaitWorkerTimeout").finish()
}
}
impl<T> fmt::Display for AwaitWorkerTimeout<T> {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
f.debug_struct("AwaitWorkerTimeout").finish()
}
}
// Create the format converter and the target texture.
fn create_converter_data_pair(
device: &wgpu::Device,
src_texture: &wgpu::Texture,
) -> ConverterDataPair {
// Create the destination format texture.
let dst_texture = wgpu::TextureBuilder::from(src_texture.descriptor.clone())
.sample_count(1)
.format(Capturer::DST_FORMAT)
.usage(wgpu::TextureUsages::RENDER_ATTACHMENT | wgpu::TextureUsages::COPY_SRC)
.build(device);
// Create the converter.
let src_sample_count = src_texture.sample_count();
let src_sample_type = src_texture.sample_type();
let src_view = src_texture.create_view(&wgpu::TextureViewDescriptor::default());
let dst_sample_count = 1;
let dst_format = dst_texture.format();
let reshaper = wgpu::TextureReshaper::new(
device,
&src_view,
src_sample_count,
src_sample_type,
dst_sample_count,
dst_format,
);
// Keep track of the `src_descriptor` to check if we need to recreate the converter.
let src_descriptor = src_texture.descriptor.clone();
ConverterDataPair {
src_descriptor,
reshaper,
dst_texture,
}
}