From 9960eb9147305c977d9e63159cee17d99b5ba4f2 Mon Sep 17 00:00:00 2001
From: Andreas Reich <andreas@rerun.io>
Date: Mon, 20 Feb 2023 09:54:00 +0100
Subject: [PATCH] Texture & buffer call now `destroy` on removal from pool
 Fixes #592

---
 .../src/wgpu_resources/bind_group_pool.rs     |  2 +-
 .../src/wgpu_resources/buffer_pool.rs         |  3 +-
 .../wgpu_resources/dynamic_resource_pool.rs   | 32 +++++++++++++------
 .../src/wgpu_resources/texture_pool.rs        |  3 +-
 4 files changed, 28 insertions(+), 12 deletions(-)

diff --git a/crates/re_renderer/src/wgpu_resources/bind_group_pool.rs b/crates/re_renderer/src/wgpu_resources/bind_group_pool.rs
index 701c1ec31a8d..460810a9db53 100644
--- a/crates/re_renderer/src/wgpu_resources/bind_group_pool.rs
+++ b/crates/re_renderer/src/wgpu_resources/bind_group_pool.rs
@@ -185,7 +185,7 @@ impl GpuBindGroupPool {
         _buffers: &mut GpuBufferPool,
         _samplers: &mut GpuSamplerPool,
     ) {
-        self.pool.frame_maintenance(frame_index);
+        self.pool.frame_maintenance(frame_index, |_res| {});
         // TODO(andreas): Update usage counter on dependent resources.
     }
 
diff --git a/crates/re_renderer/src/wgpu_resources/buffer_pool.rs b/crates/re_renderer/src/wgpu_resources/buffer_pool.rs
index 497e3299c35a..d29d73e52755 100644
--- a/crates/re_renderer/src/wgpu_resources/buffer_pool.rs
+++ b/crates/re_renderer/src/wgpu_resources/buffer_pool.rs
@@ -56,7 +56,8 @@ impl GpuBufferPool {
 
     /// Called by `RenderContext` every frame. Updates statistics and may free unused buffers.
     pub fn frame_maintenance(&mut self, frame_index: u64) {
-        self.pool.frame_maintenance(frame_index);
+        self.pool
+            .frame_maintenance(frame_index, |res| res.destroy());
     }
 
     /// Takes strong buffer handle to ensure the user is still holding on to the buffer.
diff --git a/crates/re_renderer/src/wgpu_resources/dynamic_resource_pool.rs b/crates/re_renderer/src/wgpu_resources/dynamic_resource_pool.rs
index f82842fdb30b..46e96cd4f5bb 100644
--- a/crates/re_renderer/src/wgpu_resources/dynamic_resource_pool.rs
+++ b/crates/re_renderer/src/wgpu_resources/dynamic_resource_pool.rs
@@ -91,7 +91,11 @@ where
             })
     }
 
-    pub fn frame_maintenance(&mut self, frame_index: u64) {
+    pub fn frame_maintenance(
+        &mut self,
+        frame_index: u64,
+        mut on_destroy_resource: impl FnMut(&Res),
+    ) {
         self.current_frame_index = frame_index;
 
         // Throw out any resources that we haven't reclaimed last frame.
@@ -102,7 +106,9 @@ where
                 "Drained dangling resources from last frame",
             );
             for handle in handles {
-                self.resources.remove(*handle);
+                if let Some((_, res)) = self.resources.remove(*handle) {
+                    on_destroy_resource(&res);
+                }
                 self.total_resource_size_in_bytes -= desc.resource_size_in_bytes();
             }
         }
@@ -199,8 +205,10 @@ mod tests {
         // Still, no resources were dropped.
         {
             let drop_counter_before = drop_counter.load(Ordering::Acquire);
-            pool.frame_maintenance(1);
+            let mut called_destroy = false;
+            pool.frame_maintenance(1, |_| called_destroy = true);
 
+            assert!(!called_destroy);
             assert_eq!(drop_counter_before, drop_counter.load(Ordering::Acquire),);
         }
 
@@ -212,8 +220,11 @@ mod tests {
         // Doing frame maintenance twice will drop all resources
         {
             let drop_counter_before = drop_counter.load(Ordering::Acquire);
-            pool.frame_maintenance(2);
-            pool.frame_maintenance(3);
+            let mut called_destroy = false;
+            pool.frame_maintenance(2, |_| called_destroy = true);
+            assert!(!called_destroy);
+            pool.frame_maintenance(3, |_| called_destroy = true);
+            assert!(called_destroy);
             let drop_counter_now = drop_counter.load(Ordering::Acquire);
             assert_eq!(
                 drop_counter_before + initial_resource_descs.len() * 2,
@@ -236,9 +247,12 @@ mod tests {
             assert_ne!(handle0, handle1);
             drop(handle1);
 
-            pool.frame_maintenance(4);
+            let mut called_destroy = false;
+            pool.frame_maintenance(4, |_| called_destroy = true);
+            assert!(called_destroy);
             assert_eq!(drop_counter_before, drop_counter.load(Ordering::Acquire),);
-            pool.frame_maintenance(5);
+            pool.frame_maintenance(5, |_| called_destroy = true);
+            assert!(called_destroy);
             assert_eq!(
                 drop_counter_before + 1,
                 drop_counter.load(Ordering::Acquire),
@@ -314,8 +328,8 @@ mod tests {
         // Query with invalid handle
         let inner_handle = *handle;
         drop(handle);
-        pool.frame_maintenance(0);
-        pool.frame_maintenance(1);
+        pool.frame_maintenance(0, |_| {});
+        pool.frame_maintenance(1, |_| {});
         assert!(matches!(
             pool.get_resource(inner_handle),
             Err(PoolError::ResourceNotAvailable)
diff --git a/crates/re_renderer/src/wgpu_resources/texture_pool.rs b/crates/re_renderer/src/wgpu_resources/texture_pool.rs
index 96f697d4067f..7ad410ae4762 100644
--- a/crates/re_renderer/src/wgpu_resources/texture_pool.rs
+++ b/crates/re_renderer/src/wgpu_resources/texture_pool.rs
@@ -99,7 +99,8 @@ impl GpuTexturePool {
 
     /// Called by `RenderContext` every frame. Updates statistics and may free unused textures.
     pub fn frame_maintenance(&mut self, frame_index: u64) {
-        self.pool.frame_maintenance(frame_index);
+        self.pool
+            .frame_maintenance(frame_index, |res| res.texture.destroy());
     }
 
     /// Takes strong texture handle to ensure the user is still holding on to the texture.