From cb6479fcdc23c9a8faefe9a23fa890a646df5ad7 Mon Sep 17 00:00:00 2001
From: Kevin Reid <kpreid@google.com>
Date: Thu, 21 Jul 2022 08:12:35 -0700
Subject: [PATCH 1/2] 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).
---
 wgpu/src/util/belt.rs | 48 ++++++++++++++++++++++++++-----------------
 1 file changed, 29 insertions(+), 19 deletions(-)

diff --git a/wgpu/src/util/belt.rs b/wgpu/src/util/belt.rs
index 8cea3ba635..4cb6d2a105 100644
--- a/wgpu/src/util/belt.rs
+++ b/wgpu/src/util/belt.rs
@@ -73,23 +73,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,
+                }
             }
         };
 
@@ -121,10 +125,7 @@ impl StagingBelt {
     ///
     /// This has to be called after the command encoders written to `write_buffer` are submitted!
     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(..) {
@@ -138,6 +139,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 {

From 652bf087bccb42e985f32b4911133dba26f681cf Mon Sep 17 00:00:00 2001
From: Kevin Reid <kpreid@google.com>
Date: Tue, 26 Jul 2022 07:49:07 -0700
Subject: [PATCH 2/2] Changelog entry.

---
 CHANGELOG.md | 4 ++++
 1 file changed, 4 insertions(+)

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)