From 50659696e8a9bbef34381e2ee9e074d5643233fe Mon Sep 17 00:00:00 2001
From: Kevin Reid <kpreid@google.com>
Date: Thu, 21 Jul 2022 08:12:35 -0700
Subject: [PATCH] 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 8cea3ba6358..25567a6f163 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 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 {