From a420e453c3d9c93dfb1a8526bf11c000d895c916 Mon Sep 17 00:00:00 2001
From: Kevin Reid <kpreid@google.com>
Date: Tue, 26 Jul 2022 07:59:27 -0700
Subject: [PATCH] StagingBelt: check for free chunks in the `receiver` as well
 as `free_chunks`. (#2906)

* StagingBelt: check for free chunks in the `receiver` as well as `free_chunks`.

Previously, chunks would only be taken from the GPU callback receiver
and put on `free_chunks` during `finish()`. So, there might be chunks
that are actually available for use but wouldn't be used.

Now, we consult the receiver whenever we're about to consult the
`free_chunks`, so the reuse of chunks will be as good as possible (given
the application's uses of operations that trigger `map_async` callbacks).

* Changelog entry.
---
 CHANGELOG.md          |  4 ++++
 wgpu/src/util/belt.rs | 48 ++++++++++++++++++++++++++-----------------
 2 files changed, 33 insertions(+), 19 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 4843baf10e..f91d440ad5 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -72,6 +72,10 @@ the same every time it is rendered, we now warn if it is missing.
 #### General
 - Added downlevel restriction error message for `InvalidFormatUsages` error by @Seamooo in [#2886](https://github.com/gfx-rs/wgpu/pull/2886)
 
+### Performance
+
+- Made `StagingBelt::write_buffer()` check more thoroughly for reusable memory; by @kpreid in [#2906](https://github.com/gfx-rs/wgpu/pull/2906)
+
 ### Documentation
 
 - Expanded `StagingBelt` documentation by @kpreid in [#2905](https://github.com/gfx-rs/wgpu/pull/2905)
diff --git a/wgpu/src/util/belt.rs b/wgpu/src/util/belt.rs
index 7f4df9191c..9800718e82 100644
--- a/wgpu/src/util/belt.rs
+++ b/wgpu/src/util/belt.rs
@@ -88,23 +88,27 @@ impl StagingBelt {
             .position(|chunk| chunk.offset + size.get() <= chunk.size)
         {
             self.active_chunks.swap_remove(index)
-        } else if let Some(index) = self
-            .free_chunks
-            .iter()
-            .position(|chunk| size.get() <= chunk.size)
-        {
-            self.free_chunks.swap_remove(index)
         } else {
-            let size = self.chunk_size.max(size.get());
-            Chunk {
-                buffer: Arc::new(device.create_buffer(&BufferDescriptor {
-                    label: Some("(wgpu internal) StagingBelt staging buffer"),
+            self.receive_chunks(); // ensure self.free_chunks is up to date
+
+            if let Some(index) = self
+                .free_chunks
+                .iter()
+                .position(|chunk| size.get() <= chunk.size)
+            {
+                self.free_chunks.swap_remove(index)
+            } else {
+                let size = self.chunk_size.max(size.get());
+                Chunk {
+                    buffer: Arc::new(device.create_buffer(&BufferDescriptor {
+                        label: Some("(wgpu internal) StagingBelt staging buffer"),
+                        size,
+                        usage: BufferUsages::MAP_WRITE | BufferUsages::COPY_SRC,
+                        mapped_at_creation: true,
+                    })),
                     size,
-                    usage: BufferUsages::MAP_WRITE | BufferUsages::COPY_SRC,
-                    mapped_at_creation: true,
-                })),
-                size,
-                offset: 0,
+                    offset: 0,
+                }
             }
         };
 
@@ -142,10 +146,7 @@ impl StagingBelt {
     /// [`StagingBelt::write_buffer()`] are submitted. Additional calls are harmless.
     /// Not calling this as soon as possible may result in increased buffer memory usage.
     pub fn recall(&mut self) {
-        while let Ok(mut chunk) = self.receiver.try_recv() {
-            chunk.offset = 0;
-            self.free_chunks.push(chunk);
-        }
+        self.receive_chunks();
 
         let sender = &self.sender;
         for chunk in self.closed_chunks.drain(..) {
@@ -159,6 +160,15 @@ impl StagingBelt {
                 });
         }
     }
+
+    /// Move all chunks that the GPU is done with (and are now mapped again)
+    /// from `self.receiver` to `self.free_chunks`.
+    fn receive_chunks(&mut self) {
+        while let Ok(mut chunk) = self.receiver.try_recv() {
+            chunk.offset = 0;
+            self.free_chunks.push(chunk);
+        }
+    }
 }
 
 impl fmt::Debug for StagingBelt {