Improved readback data handling (#1734)

* user or re_renderer is no longer forced to consume all data in one go, instead we'll search if the data you request is available and return it then
    * this leaves the possibility for stale data, which might cause heavy memory leaks, so we clean stale stuff up automatically! (this can happen easily in a healthy application! Think closing a view for which you request picking data)
* identifiers are user-chosen and not required to be unique (space view id is a good & valid identifier!)
* readback data can cary arbitrary user data
* the readback belt itself is a strictly internal datastructure now. Higher level systems like `ScreenshotProcessor` (new! shifting further towards a class-per-draw-pass; this is still an ongoing evolution!) or `PickingLayerProcessor` wrap your userdata and provide everything you need to know about for their readback data
   * this keeps the readback belt agnostic while presenting high level constructs where needed! 
* Do smaller readback chunks

crates/re_renderer/examples/multiview.rs

@@ -12,10 +12,9 @@ use re_renderer::{
         GenericSkyboxDrawData, LineDrawData, LineStripFlags, MeshDrawData, MeshInstance,
         TestTriangleDrawData,
     },
-    view_builder::{
-        OrthographicCameraMode, Projection, ScheduledScreenshot, TargetConfiguration, ViewBuilder,
-    },
-    Color32, LineStripSeriesBuilder, PointCloudBuilder, RenderContext, Rgba, Size,
+    view_builder::{OrthographicCameraMode, Projection, TargetConfiguration, ViewBuilder},
+    Color32, GpuReadbackIdentifier, LineStripSeriesBuilder, PointCloudBuilder, RenderContext, Rgba,
+    ScreenshotProcessor, Size,
 };
 use winit::event::{ElementState, VirtualKeyCode};
 
@@ -151,7 +150,6 @@ struct Multiview {
     random_points_colors: Vec<Color32>,
 
     take_screenshot_next_frame_for_view: Option<u32>,
-    scheduled_screenshots: Vec<ScheduledScreenshot>,
 }
 
 fn random_color(rnd: &mut impl rand::Rng) -> Color32 {
@@ -164,6 +162,45 @@ fn random_color(rnd: &mut impl rand::Rng) -> Color32 {
     .into()
 }
 
+/// Readback identifier for screenshots.
+/// Identifiers don't need to be unique and we don't have anything interesting to distinguish here!
+const READBACK_IDENTIFIER: GpuReadbackIdentifier = 0;
+
+fn handle_incoming_screenshots(re_ctx: &RenderContext) {
+    ScreenshotProcessor::next_readback_result(
+        re_ctx,
+        READBACK_IDENTIFIER,
+        |data, _extent, view_idx: u32| {
+            re_log::info!(
+                "Received screenshot for view {view_idx}. Total bytes {:?}",
+                data.len()
+            );
+
+            #[cfg(not(target_arch = "wasm32"))]
+            {
+                // Get next available file name.
+                let mut i = 1;
+                let filename = loop {
+                    let filename = format!("screenshot_{i}.png");
+                    if !std::path::Path::new(&filename).exists() {
+                        break filename;
+                    }
+                    i += 1;
+                };
+
+                image::save_buffer(
+                    filename,
+                    data,
+                    _extent.x,
+                    _extent.y,
+                    image::ColorType::Rgba8,
+                )
+                .expect("Failed to save screenshot");
+            }
+        },
+    );
+}
+
 impl Multiview {
     fn draw_view<D: 'static + re_renderer::renderer::DrawData + Sync + Send + Clone>(
         &mut self,
@@ -180,8 +217,9 @@ impl Multiview {
             .take_screenshot_next_frame_for_view
             .map_or(false, |i| i == index)
         {
-            self.scheduled_screenshots
-                .push(view_builder.schedule_screenshot(re_ctx).unwrap());
+            view_builder
+                .schedule_screenshot(re_ctx, READBACK_IDENTIFIER, index)
+                .unwrap();
             re_log::info!("Scheduled screenshot for view {}", index);
         }
 
@@ -193,54 +231,6 @@ impl Multiview {
 
         (view_builder, command_buffer)
     }
-
-    fn handle_incoming_screenshots(&mut self, re_ctx: &mut RenderContext) {
-        re_ctx
-            .gpu_readback_belt
-            .lock()
-            .receive_data(|data, identifier| {
-                if let Some(index) = self
-                    .scheduled_screenshots
-                    .iter()
-                    .position(|s| s.identifier == identifier)
-                {
-                    re_log::info!(
-                        "Received screenshot. Total bytes {:?}. Identifier {identifier}",
-                        data.len()
-                    );
-
-                    #[cfg(target_arch = "wasm32")]
-                    self.scheduled_screenshots.remove(index);
-
-                    #[cfg(not(target_arch = "wasm32"))]
-                    {
-                        let screenshot = self.scheduled_screenshots.swap_remove(index);
-
-                        re_log::info!("Received screenshot. Total bytes {}", data.len());
-
-                        // Get next available file name.
-                        let mut i = 1;
-                        let filename = loop {
-                            let filename = format!("screenshot_{i}.png");
-                            if !std::path::Path::new(&filename).exists() {
-                                break filename;
-                            }
-                            i += 1;
-                        };
-
-                        #[cfg(not(target_arch = "wasm32"))]
-                        image::save_buffer(
-                            filename,
-                            &screenshot.row_info.remove_padding(data),
-                            screenshot.extent.x,
-                            screenshot.extent.y,
-                            image::ColorType::Rgba8,
-                        )
-                        .expect("Failed to save screenshot");
-                    }
-                }
-            });
-    }
 }
 
 impl Example for Multiview {
@@ -297,7 +287,6 @@ impl Example for Multiview {
             random_points_colors,
 
             take_screenshot_next_frame_for_view: None,
-            scheduled_screenshots: Vec::new(),
         }
     }
 
@@ -317,7 +306,7 @@ impl Example for Multiview {
             ) * 10.0;
         }
 
-        self.handle_incoming_screenshots(re_ctx);
+        handle_incoming_screenshots(re_ctx);
 
         let seconds_since_startup = time.seconds_since_startup();
         let view_from_world =

crates/re_renderer/examples/picking.rs

@@ -1,10 +1,9 @@
-use ahash::HashMap;
 use itertools::Itertools as _;
 use rand::Rng;
 use re_renderer::{
     view_builder::{Projection, TargetConfiguration, ViewBuilder},
-    Color32, GpuReadbackBufferIdentifier, IntRect, PickingLayerId, PickingLayerInstanceId,
-    PointCloudBuilder, RenderContext, ScheduledPickingRect, Size,
+    Color32, GpuReadbackIdentifier, IntRect, PickingLayerInstanceId, PickingLayerProcessor,
+    PointCloudBuilder, Size,
 };
 
 mod framework;
@@ -18,7 +17,6 @@ struct PointSet {
 
 struct Picking {
     point_sets: Vec<PointSet>,
-    scheduled_picking_rects: HashMap<GpuReadbackBufferIdentifier, ScheduledPickingRect>,
     picking_position: glam::UVec2,
 }
 
@@ -32,36 +30,9 @@ fn random_color(rnd: &mut impl rand::Rng) -> Color32 {
     .into()
 }
 
-impl Picking {
-    #[allow(clippy::unused_self)]
-    fn handle_incoming_picking_data(&mut self, re_ctx: &mut RenderContext, _time: f32) {
-        re_ctx
-            .gpu_readback_belt
-            .lock()
-            .receive_data(|data, identifier| {
-                if let Some(picking_rect_info) = self.scheduled_picking_rects.remove(&identifier) {
-                    // TODO(andreas): Move this into a utility function?
-                    let picking_data_without_padding =
-                        picking_rect_info.row_info.remove_padding(data);
-                    let picking_data: &[PickingLayerId] =
-                        bytemuck::cast_slice(&picking_data_without_padding);
-
-                    // Grab the middle pixel. usually we'd want to do something clever that snaps the the closest object of interest.
-                    let picked_pixel = picking_data[(picking_rect_info.rect.extent.x / 2
-                        + (picking_rect_info.rect.extent.y / 2) * picking_rect_info.rect.extent.x)
-                        as usize];
-
-                    if picked_pixel.object.0 != 0 {
-                        let point_set = &mut self.point_sets[picked_pixel.object.0 as usize - 1];
-                        point_set.radii[picked_pixel.instance.0 as usize] = Size::new_scene(0.1);
-                        point_set.colors[picked_pixel.instance.0 as usize] = Color32::DEBUG_COLOR;
-                    }
-                } else {
-                    re_log::error!("Received picking data for unknown identifier");
-                }
-            });
-    }
-}
+/// Readback identifier for picking rects.
+/// Identifiers don't need to be unique and we don't have anything interesting to distinguish here!
+const READBACK_IDENTIFIER: GpuReadbackIdentifier = 0;
 
 impl framework::Example for Picking {
     fn title() -> &'static str {
@@ -103,7 +74,6 @@ impl framework::Example for Picking {
 
         Picking {
             point_sets,
-            scheduled_picking_rects: HashMap::default(),
             picking_position: glam::UVec2::ZERO,
         }
     }
@@ -112,10 +82,23 @@ impl framework::Example for Picking {
         &mut self,
         re_ctx: &mut re_renderer::RenderContext,
         resolution: [u32; 2],
-        time: &framework::Time,
+        _time: &framework::Time,
         pixels_from_point: f32,
     ) -> Vec<framework::ViewDrawResult> {
-        self.handle_incoming_picking_data(re_ctx, time.seconds_since_startup());
+        while let Some(picking_result) =
+            PickingLayerProcessor::next_readback_result::<()>(re_ctx, READBACK_IDENTIFIER)
+        {
+            // Grab the middle pixel. usually we'd want to do something clever that snaps the the closest object of interest.
+            let picked_pixel = picking_result.picking_data[(picking_result.rect.extent.x / 2
+                + (picking_result.rect.extent.y / 2) * picking_result.rect.extent.x)
+                as usize];
+
+            if picked_pixel.object.0 != 0 {
+                let point_set = &mut self.point_sets[picked_pixel.object.0 as usize - 1];
+                point_set.radii[picked_pixel.instance.0 as usize] = Size::new_scene(0.1);
+                point_set.colors[picked_pixel.instance.0 as usize] = Color32::DEBUG_COLOR;
+            }
+        }
 
         let mut view_builder = ViewBuilder::default();
 
@@ -148,19 +131,13 @@ impl framework::Example for Picking {
         // Use an uneven number of pixels for the picking rect so that there is a clearly defined middle-pixel.
         // (for this sample a size of 1 would be sufficient, but for a real application you'd want to use a larger size to allow snapping)
         let picking_rect_size = 31;
-
-        let picking_rect = view_builder
-            .schedule_picking_readback(
-                re_ctx,
-                IntRect::from_middle_and_extent(
-                    self.picking_position.as_ivec2(),
-                    glam::uvec2(picking_rect_size, picking_rect_size),
-                ),
-                false,
-            )
+        let picking_rect = IntRect::from_middle_and_extent(
+            self.picking_position.as_ivec2(),
+            glam::uvec2(picking_rect_size, picking_rect_size),
+        );
+        view_builder
+            .schedule_picking_rect(re_ctx, picking_rect, READBACK_IDENTIFIER, (), false)
             .unwrap();
-        self.scheduled_picking_rects
-            .insert(picking_rect.identifier, picking_rect);
 
         let mut builder = PointCloudBuilder::<()>::new(re_ctx);