bevy_shader_plugin/
readback.rs

1use std::{
2    any::type_name,
3    marker::PhantomData,
4    sync::{
5        Arc, Mutex,
6        atomic::{AtomicBool, Ordering},
7    },
8};
9
10use bevy::{
11    core_pipeline::schedule::camera_driver,
12    prelude::*,
13    render::{
14        Render, RenderApp, RenderSystems,
15        render_asset::RenderAssets,
16        render_resource::{
17            Buffer, BufferDescriptor, BufferUsages, MapMode,
18            encase::{
19                self, ShaderSize, ShaderType,
20                internal::{CreateFrom, WriteInto},
21            },
22        },
23        renderer::{RenderContext, RenderDevice, RenderGraph, RenderGraphSystems},
24        storage::GpuShaderBuffer,
25    },
26};
27
28use crate::{
29    ArrayBufferChanges, ShaderPassSystems, ShaderSystems,
30    auto_array::{ArrayBufferState, array_stride},
31    bindings::AutoBufferKind,
32    cpu_buffer::CpuBuffer,
33    gpu_buffer::GpuBuffer,
34};
35
36/// When a read-back buffer is copied to the CPU.
37#[derive(Clone, Copy, Debug, PartialEq, Eq)]
38pub enum ReadbackMode {
39    /// Copy the buffer every frame (skipping frames while a previous copy is still
40    /// being mapped), so `latest()` keeps updating. Costs one buffer copy per frame.
41    EveryFrame,
42    /// Copy the buffer once per `request()` call.
43    OnRequest,
44}
45
46/// State shared between a main-world readback resource and the render world. Held in an
47/// `Arc` by both, so no extraction is needed and it works with pipelined rendering.
48#[derive(Default)]
49struct ReadbackShared {
50    /// Set by `request()`, consumed when a copy is queued.
51    requested: AtomicBool,
52    /// A copy has been queued and its bytes haven't arrived yet.
53    in_flight: AtomicBool,
54    /// Bytes delivered by the map callback, taken by the main world.
55    result: Mutex<Option<Vec<u8>>>,
56}
57
58impl ReadbackShared {
59    fn take_result(&self) -> Option<Vec<u8>> {
60        self.result.lock().ok()?.take()
61    }
62}
63
64/// The most recent CPU copy of an array buffer or GPU buffer, decoded into `T` values.
65///
66/// This resource only changes when new data arrives, so check `is_changed()` to react to
67/// fresh data. [`request`](Self::request) takes `&self`, so requesting through
68/// `Res<GpuReadback<..>>` doesn't count as a change:
69///
70/// ```rust,ignore
71/// fn request(readback: Res<GpuReadback<Particles, Particle>>) {
72///     readback.request();
73/// }
74///
75/// fn print(readback: Res<GpuReadback<Particles, Particle>>) {
76///     if !readback.is_changed() { return; }
77///     let Some(particles) = readback.latest() else { return };
78///     info!("first particle at {}", particles[0].pos);
79/// }
80/// ```
81///
82/// The data shows the buffer at the end of the frame it was copied in, and arrives
83/// 1–3 frames later. Use it for UI, stats, or logging — not for logic that needs the
84/// exact current GPU state.
85#[derive(Resource)]
86pub struct GpuReadback<Tag, T> {
87    shared: Arc<ReadbackShared>,
88    latest: Option<Vec<T>>,
89    _marker: PhantomData<Tag>,
90}
91
92impl<Tag, T> GpuReadback<Tag, T> {
93    /// The most recently read-back contents of the buffer, or `None` until the first
94    /// readback arrives.
95    pub fn latest(&self) -> Option<&[T]> {
96        self.latest.as_deref()
97    }
98
99    /// Ask for the buffer to be copied back. Only needed with
100    /// [`ReadbackMode::OnRequest`]. Requests made while a copy is in flight are merged
101    /// into one follow-up copy.
102    pub fn request(&self) {
103        self.shared.requested.store(true, Ordering::Release);
104    }
105}
106
107/// The most recent CPU copy of a storage buffer, decoded as the registered resource type
108/// `S` (which may itself contain arrays).
109///
110/// Works like [`GpuReadback`]: check `is_changed()` for fresh data, and call
111/// [`request`](Self::request) in [`ReadbackMode::OnRequest`] mode. The data arrives
112/// 1–3 frames after the copy.
113#[derive(Resource)]
114pub struct StorageReadback<S> {
115    shared: Arc<ReadbackShared>,
116    latest: Option<S>,
117}
118
119impl<S> StorageReadback<S> {
120    /// The most recently read-back contents of the buffer, or `None` until the first
121    /// readback arrives.
122    pub fn latest(&self) -> Option<&S> {
123        self.latest.as_ref()
124    }
125
126    /// Ask for the buffer to be copied back. Only needed with
127    /// [`ReadbackMode::OnRequest`]. Requests made while a copy is in flight are merged
128    /// into one follow-up copy.
129    pub fn request(&self) {
130        self.shared.requested.store(true, Ordering::Release);
131    }
132}
133
134// ---------------------------------------------------------------------------
135// Render world
136// ---------------------------------------------------------------------------
137
138/// One buffer copy recorded by [`readback_copy_pass`] this frame.
139struct QueuedCopy {
140    source: Buffer,
141    staging: Buffer,
142    size: u64,
143    shared: Arc<ReadbackShared>,
144}
145
146/// Copies queued this frame. Filled during `PrepareBindGroups`, recorded by
147/// [`readback_copy_pass`], then mapped and cleared by [`map_readbacks`].
148#[derive(Resource, Default)]
149struct ReadbackCopies(Vec<QueuedCopy>);
150
151/// Render-world state for one read-back registration. `Key` distinguishes registrations.
152#[derive(Resource)]
153struct ReadbackJob<Key> {
154    shared: Arc<ReadbackShared>,
155    mode: ReadbackMode,
156    staging: Option<Buffer>,
157    _marker: PhantomData<Key>,
158}
159
160// Keys keep registrations of different buffer kinds apart.
161struct GpuKey<Tag>(PhantomData<Tag>);
162struct ArrayKey<Tag>(PhantomData<Tag>);
163struct StorageKey<S>(PhantomData<S>);
164
165/// Whether to queue a copy this frame. Never copies while a previous copy is in flight;
166/// an on-request copy consumes the request, which otherwise stays pending.
167fn should_copy(mode: ReadbackMode, in_flight: bool, requested: &AtomicBool) -> bool {
168    if in_flight {
169        return false;
170    }
171    match mode {
172        ReadbackMode::EveryFrame => true,
173        ReadbackMode::OnRequest => requested.swap(false, Ordering::AcqRel),
174    }
175}
176
177impl<Key> ReadbackJob<Key> {
178    fn new(shared: Arc<ReadbackShared>, mode: ReadbackMode) -> Self {
179        Self {
180            shared,
181            mode,
182            staging: None,
183            _marker: PhantomData,
184        }
185    }
186
187    /// Queue a copy of the first `size` bytes of `source`, if this job wants one now.
188    fn queue_copy(
189        &mut self,
190        source: &Buffer,
191        size: u64,
192        render_device: &RenderDevice,
193        copies: &mut ReadbackCopies,
194    ) {
195        if size == 0
196            || !should_copy(
197                self.mode,
198                self.shared.in_flight.load(Ordering::Acquire),
199                &self.shared.requested,
200            )
201        {
202            return;
203        }
204
205        // Reuse the staging buffer unless the size changed. It is never still mapped
206        // here: `in_flight` is only cleared after unmapping.
207        let staging = match &self.staging {
208            Some(staging) if staging.size() == size => staging.clone(),
209            _ => {
210                let staging = render_device.create_buffer(&BufferDescriptor {
211                    label: Some("readback_staging_buffer"),
212                    size,
213                    usage: BufferUsages::MAP_READ | BufferUsages::COPY_DST,
214                    mapped_at_creation: false,
215                });
216                self.staging = Some(staging.clone());
217                staging
218            }
219        };
220
221        self.shared.in_flight.store(true, Ordering::Release);
222        copies.0.push(QueuedCopy {
223            source: source.clone(),
224            staging,
225            size,
226            shared: self.shared.clone(),
227        });
228    }
229}
230
231/// Records every queued readback copy. Runs in the root
232/// [`RenderGraph`] schedule after
233/// [`camera_driver`](bevy::core_pipeline::schedule::camera_driver), so copies see
234/// everything the frame's compute passes and cameras wrote.
235fn readback_copy_pass(copies: Res<ReadbackCopies>, mut ctx: RenderContext) {
236    if copies.0.is_empty() {
237        return;
238    }
239    let encoder = ctx.command_encoder();
240    for copy in &copies.0 {
241        encoder.copy_buffer_to_buffer(&copy.source, 0, &copy.staging, 0, copy.size);
242    }
243}
244
245/// After the frame is submitted, asks wgpu to map every staging buffer copied into this
246/// frame. The callback hands the bytes to the main world.
247fn map_readbacks(mut copies: ResMut<ReadbackCopies>) {
248    for copy in copies.0.drain(..) {
249        let QueuedCopy {
250            staging,
251            size,
252            shared,
253            ..
254        } = copy;
255        let mapped = staging.clone();
256        staging
257            .slice(..size)
258            .map_async(MapMode::Read, move |result| {
259                match result {
260                    Ok(()) => {
261                        let bytes = mapped.slice(..size).get_mapped_range().to_vec();
262                        mapped.unmap();
263                        if let Ok(mut slot) = shared.result.lock() {
264                            *slot = Some(bytes);
265                        }
266                    }
267                    Err(error) => warn!("buffer readback failed: {error}"),
268                }
269                shared.in_flight.store(false, Ordering::Release);
270            });
271    }
272}
273
274/// Copy-and-map machinery shared by every readback registration.
275struct ReadbackPlugin;
276
277impl Plugin for ReadbackPlugin {
278    fn build(&self, app: &mut App) {
279        let Some(render_app) = app.get_sub_app_mut(RenderApp) else {
280            return;
281        };
282        render_app
283            .init_resource::<ReadbackCopies>()
284            .add_systems(
285                RenderGraph,
286                readback_copy_pass
287                    .in_set(ShaderPassSystems::ReadbackCopy)
288                    .in_set(RenderGraphSystems::Render)
289                    .after(camera_driver),
290            )
291            // Cleanup runs after the frame's commands are submitted, which is also where
292            // Bevy maps its own readbacks.
293            .add_systems(Render, map_readbacks.in_set(RenderSystems::Cleanup));
294    }
295}
296
297fn add_readback_plugin(app: &mut App) {
298    if !app.is_plugin_added::<ReadbackPlugin>() {
299        app.add_plugins(ReadbackPlugin);
300    }
301}
302
303/// Adds the render-world job for `Key` and a system that resolves its source buffer each
304/// frame and queues copies.
305fn add_render_job<Key, M>(
306    app: &mut App,
307    shared: Arc<ReadbackShared>,
308    mode: ReadbackMode,
309    queue_system: impl IntoScheduleConfigs<bevy::ecs::system::ScheduleSystem, M>,
310) where
311    Key: Send + Sync + 'static,
312{
313    let Some(render_app) = app.get_sub_app_mut(RenderApp) else {
314        return;
315    };
316    render_app
317        .insert_resource(ReadbackJob::<Key>::new(shared, mode))
318        .add_systems(
319            Render,
320            queue_system
321                .in_set(RenderSystems::PrepareBindGroups)
322                .after(ShaderSystems::PrepareBindings),
323        );
324}
325
326fn decode<D: ShaderType + CreateFrom>(bytes: &[u8]) -> Option<D> {
327    match encase::StorageBuffer::new(bytes).create() {
328        Ok(value) => Some(value),
329        Err(error) => {
330            error!("could not decode read-back {}: {error}", type_name::<D>());
331            None
332        }
333    }
334}
335
336/// Main world: moves newly arrived bytes into the [`GpuReadback`] resource.
337fn deliver_slice<Tag, T>(mut readback: ResMut<GpuReadback<Tag, T>>)
338where
339    Tag: Send + Sync + 'static,
340    T: ShaderSize + CreateFrom + Send + Sync + 'static,
341{
342    if let Some(bytes) = readback.shared.take_result()
343        && let Some(values) = decode::<Vec<T>>(&bytes)
344    {
345        readback.latest = Some(values);
346    }
347}
348
349/// Main world: moves newly arrived bytes into the [`StorageReadback`] resource.
350fn deliver_storage<S>(mut readback: ResMut<StorageReadback<S>>)
351where
352    S: ShaderType + CreateFrom + Send + Sync + 'static,
353{
354    if let Some(bytes) = readback.shared.take_result()
355        && let Some(value) = decode::<S>(&bytes)
356    {
357        readback.latest = Some(value);
358    }
359}
360
361/// Inserts `GpuReadback<Tag, T>` and its delivery system, returning the shared state.
362fn insert_slice_readback<Tag, T>(app: &mut App) -> Arc<ReadbackShared>
363where
364    Tag: Send + Sync + 'static,
365    T: ShaderSize + CreateFrom + Send + Sync + 'static,
366{
367    assert!(
368        !app.world().contains_resource::<GpuReadback<Tag, T>>(),
369        "readback is already enabled for tag `{}`",
370        type_name::<Tag>()
371    );
372    let shared = Arc::new(ReadbackShared::default());
373    app.insert_resource(GpuReadback::<Tag, T> {
374        shared: shared.clone(),
375        latest: None,
376        _marker: PhantomData,
377    })
378    .add_systems(PreUpdate, deliver_slice::<Tag, T>);
379    shared
380}
381
382// ---------------------------------------------------------------------------
383// Registration
384// ---------------------------------------------------------------------------
385
386/// See [`ShaderAppExt::read_back_gpu_buffer`](crate::ShaderAppExt::read_back_gpu_buffer).
387pub(crate) fn register_gpu<Tag, T>(app: &mut App, mode: ReadbackMode)
388where
389    Tag: Send + Sync + 'static,
390    T: ShaderSize + CreateFrom + Send + Sync + 'static,
391{
392    let Some(gpu_buffer) = app.world().get_resource::<GpuBuffer<Tag>>() else {
393        panic!(
394            "read_back_gpu_buffer::<{}, _>: call register_gpu_buffer for this tag first",
395            type_name::<Tag>()
396        );
397    };
398    let asset_id = gpu_buffer.handle().id();
399
400    add_readback_plugin(app);
401    let shared = insert_slice_readback::<Tag, T>(app);
402    add_render_job::<GpuKey<Tag>, _>(
403        app,
404        shared,
405        mode,
406        move |buffers: Res<RenderAssets<GpuShaderBuffer>>,
407              mut job: ResMut<ReadbackJob<GpuKey<Tag>>>,
408              render_device: Res<RenderDevice>,
409              mut copies: ResMut<ReadbackCopies>| {
410            if let Some(gpu_buffer) = buffers.get(asset_id) {
411                let size = gpu_buffer.buffer.size();
412                job.queue_copy(&gpu_buffer.buffer, size, &render_device, &mut copies);
413            }
414        },
415    );
416}
417
418/// See [`ShaderAppExt::read_back_array_buffer`](crate::ShaderAppExt::read_back_array_buffer).
419pub(crate) fn register_array<Tag, T>(app: &mut App, mode: ReadbackMode)
420where
421    Tag: Send + Sync + 'static,
422    T: ShaderSize + CreateFrom + Send + Sync + 'static,
423{
424    assert!(
425        app.world().contains_resource::<ArrayBufferChanges<Tag>>(),
426        "read_back_array_buffer::<{}, _>: call register_array_buffer for this tag first",
427        type_name::<Tag>()
428    );
429
430    add_readback_plugin(app);
431    let shared = insert_slice_readback::<Tag, T>(app);
432    add_render_job::<ArrayKey<Tag>, _>(
433        app,
434        shared,
435        mode,
436        |state: Option<Res<ArrayBufferState<Tag>>>,
437         mut job: ResMut<ReadbackJob<ArrayKey<Tag>>>,
438         render_device: Res<RenderDevice>,
439         mut copies: ResMut<ReadbackCopies>| {
440            let Some(state) = state else { return };
441            if state.stride != array_stride::<T>() {
442                error_once!(
443                    "read_back_array_buffer::<{}, {}>: element type doesn't match the one passed to register_array_buffer",
444                    type_name::<Tag>(),
445                    type_name::<T>()
446                );
447                return;
448            }
449            let size = state.buffer.size();
450            job.queue_copy(&state.buffer, size, &render_device, &mut copies);
451        },
452    );
453}
454
455/// See [`ShaderAppExt::read_back_storage_buffer`](crate::ShaderAppExt::read_back_storage_buffer).
456pub(crate) fn register_storage<S>(app: &mut App, mode: ReadbackMode)
457where
458    S: ShaderType + WriteInto + Default + CreateFrom + Resource,
459{
460    let has_storage_slot = app
461        .get_sub_app(RenderApp)
462        .and_then(|render_app| render_app.world().get_resource::<CpuBuffer<S>>())
463        .is_some_and(|state| {
464            state
465                .slot_kinds()
466                .any(|kind| matches!(kind, AutoBufferKind::Storage { .. }))
467        });
468    assert!(
469        has_storage_slot,
470        "read_back_storage_buffer::<{}>: call register_storage_buffer for this type first",
471        type_name::<S>()
472    );
473    assert!(
474        !app.world().contains_resource::<StorageReadback<S>>(),
475        "read_back_storage_buffer::<{}>: readback is already enabled for this type",
476        type_name::<S>()
477    );
478
479    add_readback_plugin(app);
480    let shared = Arc::new(ReadbackShared::default());
481    app.insert_resource(StorageReadback::<S> {
482        shared: shared.clone(),
483        latest: None,
484    })
485    .add_systems(PreUpdate, deliver_storage::<S>);
486
487    add_render_job::<StorageKey<S>, _>(
488        app,
489        shared,
490        mode,
491        |state: Res<CpuBuffer<S>>,
492         mut job: ResMut<ReadbackJob<StorageKey<S>>>,
493         render_device: Res<RenderDevice>,
494         mut copies: ResMut<ReadbackCopies>| {
495            if let Some((buffer, size)) = state.readback_source() {
496                job.queue_copy(buffer, size, &render_device, &mut copies);
497            }
498        },
499    );
500}
501
502#[cfg(test)]
503mod tests {
504    use super::*;
505
506    #[test]
507    fn every_frame_copies_unless_in_flight() {
508        let requested = AtomicBool::new(false);
509        assert!(should_copy(ReadbackMode::EveryFrame, false, &requested));
510        assert!(!should_copy(ReadbackMode::EveryFrame, true, &requested));
511    }
512
513    #[test]
514    fn on_request_consumes_the_request() {
515        let requested = AtomicBool::new(true);
516        assert!(should_copy(ReadbackMode::OnRequest, false, &requested));
517        assert!(!requested.load(Ordering::Acquire));
518        assert!(!should_copy(ReadbackMode::OnRequest, false, &requested));
519    }
520
521    #[test]
522    fn on_request_stays_pending_while_in_flight() {
523        let requested = AtomicBool::new(true);
524        assert!(!should_copy(ReadbackMode::OnRequest, true, &requested));
525        assert!(requested.load(Ordering::Acquire));
526    }
527
528    #[test]
529    fn decodes_what_encase_wrote() {
530        #[derive(ShaderType, Debug, PartialEq)]
531        struct Sample {
532            a: Vec2,
533            b: f32,
534        }
535        let values = vec![
536            Sample {
537                a: Vec2::new(1.0, 2.0),
538                b: 3.0,
539            },
540            Sample {
541                a: Vec2::ZERO,
542                b: -1.0,
543            },
544        ];
545        let mut bytes = encase::StorageBuffer::new(Vec::new());
546        bytes.write(&values).unwrap();
547        assert_eq!(decode::<Vec<Sample>>(&bytes.into_inner()), Some(values));
548    }
549}
bevy_shader_plugin/readback.rs

bevy_shader_plugin/
readback.rs
