Buffer snatching (#4896)

Co-authored-by: Connor Fitzgerald <[email protected]>

wgpu-core/src/device/life.rs

@@ -14,7 +14,10 @@ use crate::{
         TextureViewId,
     },
     pipeline::{ComputePipeline, RenderPipeline},
-    resource::{self, Buffer, QuerySet, Resource, Sampler, StagingBuffer, Texture, TextureView},
+    resource::{
+        self, Buffer, DestroyedBuffer, QuerySet, Resource, Sampler, StagingBuffer, Texture,
+        TextureView,
+    },
     track::{ResourceTracker, Tracker},
     FastHashMap, SubmissionIndex,
 };
@@ -39,6 +42,7 @@ pub(crate) struct ResourceMaps<A: HalApi> {
     pub pipeline_layouts: FastHashMap<PipelineLayoutId, Arc<PipelineLayout<A>>>,
     pub render_bundles: FastHashMap<RenderBundleId, Arc<RenderBundle<A>>>,
     pub query_sets: FastHashMap<QuerySetId, Arc<QuerySet<A>>>,
+    pub destroyed_buffers: FastHashMap<BufferId, Arc<DestroyedBuffer<A>>>,
 }
 
 impl<A: HalApi> ResourceMaps<A> {
@@ -56,6 +60,7 @@ impl<A: HalApi> ResourceMaps<A> {
             pipeline_layouts: FastHashMap::default(),
             render_bundles: FastHashMap::default(),
             query_sets: FastHashMap::default(),
+            destroyed_buffers: FastHashMap::default(),
         }
     }
 
@@ -73,6 +78,7 @@ impl<A: HalApi> ResourceMaps<A> {
             pipeline_layouts,
             render_bundles,
             query_sets,
+            destroyed_buffers,
         } = self;
         buffers.clear();
         staging_buffers.clear();
@@ -86,6 +92,7 @@ impl<A: HalApi> ResourceMaps<A> {
         pipeline_layouts.clear();
         render_bundles.clear();
         query_sets.clear();
+        destroyed_buffers.clear();
     }
 
     pub(crate) fn extend(&mut self, mut other: Self) {
@@ -102,6 +109,7 @@ impl<A: HalApi> ResourceMaps<A> {
             pipeline_layouts,
             render_bundles,
             query_sets,
+            destroyed_buffers,
         } = self;
         buffers.extend(other.buffers.drain());
         staging_buffers.extend(other.staging_buffers.drain());
@@ -115,6 +123,7 @@ impl<A: HalApi> ResourceMaps<A> {
         pipeline_layouts.extend(other.pipeline_layouts.drain());
         render_bundles.extend(other.render_bundles.drain());
         query_sets.extend(other.query_sets.drain());
+        destroyed_buffers.extend(other.destroyed_buffers.drain());
     }
 }
 
@@ -281,6 +290,11 @@ impl<A: HalApi> LifetimeTracker<A> {
                         .staging_buffers
                         .insert(raw.as_info().id(), raw);
                 }
+                TempResource::DestroyedBuffer(destroyed) => {
+                    last_resources
+                        .destroyed_buffers
+                        .insert(destroyed.id, destroyed);
+                }
                 TempResource::Texture(raw) => {
                     last_resources.textures.insert(raw.as_info().id(), raw);
                 }
@@ -384,6 +398,9 @@ impl<A: HalApi> LifetimeTracker<A> {
             TempResource::StagingBuffer(raw) => {
                 resources.staging_buffers.insert(raw.as_info().id(), raw);
             }
+            TempResource::DestroyedBuffer(destroyed) => {
+                resources.destroyed_buffers.insert(destroyed.id, destroyed);
+            }
             TempResource::Texture(raw) => {
                 resources.textures.insert(raw.as_info().id(), raw);
             }
@@ -642,6 +659,27 @@ impl<A: HalApi> LifetimeTracker<A> {
         self
     }
 
+    fn triage_suspected_destroyed_buffers(
+        &mut self,
+        #[cfg(feature = "trace")] trace: &mut Option<&mut trace::Trace>,
+    ) {
+        for (id, buffer) in self.suspected_resources.destroyed_buffers.drain() {
+            let submit_index = buffer.submission_index;
+            if let Some(resources) = self.active.iter_mut().find(|a| a.index == submit_index) {
+                resources
+                    .last_resources
+                    .destroyed_buffers
+                    .insert(id, buffer);
+            } else {
+                self.free_resources.destroyed_buffers.insert(id, buffer);
+            }
+            #[cfg(feature = "trace")]
+            if let Some(ref mut t) = *trace {
+                t.add(trace::Action::DestroyBuffer(id));
+            }
+        }
+    }
+
     fn triage_suspected_compute_pipelines(
         &mut self,
         trackers: &Mutex<Tracker<A>>,
@@ -871,6 +909,10 @@ impl<A: HalApi> LifetimeTracker<A> {
             #[cfg(feature = "trace")]
             &mut trace,
         );
+        self.triage_suspected_destroyed_buffers(
+            #[cfg(feature = "trace")]
+            &mut trace,
+        );
     }
 
     /// Determine which buffers are ready to map, and which must wait for the

wgpu-core/src/device/queue.rs

@@ -16,8 +16,8 @@ use crate::{
     identity::{GlobalIdentityHandlerFactory, Input},
     init_tracker::{has_copy_partial_init_tracker_coverage, TextureInitRange},
     resource::{
-        Buffer, BufferAccessError, BufferMapState, Resource, ResourceInfo, ResourceType,
-        StagingBuffer, Texture, TextureInner,
+        Buffer, BufferAccessError, BufferMapState, DestroyedBuffer, Resource, ResourceInfo,
+        ResourceType, StagingBuffer, Texture, TextureInner,
     },
     resource_log, track, FastHashMap, SubmissionIndex,
 };
@@ -163,6 +163,7 @@ pub struct WrappedSubmissionIndex {
 pub enum TempResource<A: HalApi> {
     Buffer(Arc<Buffer<A>>),
     StagingBuffer(Arc<StagingBuffer<A>>),
+    DestroyedBuffer(Arc<DestroyedBuffer<A>>),
     Texture(Arc<Texture<A>>),
 }
 

wgpu-core/src/resource.rs

@@ -413,7 +413,7 @@ pub struct Buffer<A: HalApi> {
 impl<A: HalApi> Drop for Buffer<A> {
     fn drop(&mut self) {
         if let Some(raw) = self.raw.take() {
-            resource_log!("Destroy raw Buffer {:?}", self.info.label());
+            resource_log!("Deallocate raw Buffer (dropped) {:?}", self.info.label());
             unsafe {
                 use hal::Device;
                 self.device.raw().destroy_buffer(raw);
@@ -551,14 +551,25 @@ impl<A: HalApi> Buffer<A> {
         if let Some(ref mut trace) = *device.trace.lock() {
             trace.add(trace::Action::FreeBuffer(buffer_id));
         }
-        // Note: a future commit will replace this with a read guard
-        // and actually snatch the buffer.
-        let snatch_guard = device.snatchable_lock.read();
-        if self.raw.get(&snatch_guard).is_none() {
-            return Err(resource::DestroyError::AlreadyDestroyed);
-        }
 
-        let temp = queue::TempResource::Buffer(self.clone());
+        let temp = {
+            let snatch_guard = device.snatchable_lock.write();
+            let raw = match self.raw.snatch(snatch_guard) {
+                Some(raw) => raw,
+                None => {
+                    return Err(resource::DestroyError::AlreadyDestroyed);
+                }
+            };
+
+            queue::TempResource::DestroyedBuffer(Arc::new(DestroyedBuffer {
+                raw: Some(raw),
+                device: Arc::clone(&self.device),
+                submission_index: self.info.submission_index(),
+                id: self.info.id.unwrap(),
+                label: self.info.label.clone(),
+            }))
+        };
+
         let mut pending_writes = device.pending_writes.lock();
         let pending_writes = pending_writes.as_mut().unwrap();
         if pending_writes.dst_buffers.contains_key(&buffer_id) {
@@ -605,6 +616,38 @@ impl<A: HalApi> Resource<BufferId> for Buffer<A> {
     }
 }
 
+/// A buffer that has been marked as destroyed and is staged for actual deletion soon.
+#[derive(Debug)]
+pub struct DestroyedBuffer<A: HalApi> {
+    raw: Option<A::Buffer>,
+    device: Arc<Device<A>>,
+    label: String,
+    pub(crate) id: BufferId,
+    pub(crate) submission_index: u64,
+}
+
+impl<A: HalApi> DestroyedBuffer<A> {
+    pub fn label(&self) -> &dyn Debug {
+        if !self.label.is_empty() {
+            return &self.label;
+        }
+
+        &self.id
+    }
+}
+
+impl<A: HalApi> Drop for DestroyedBuffer<A> {
+    fn drop(&mut self) {
+        if let Some(raw) = self.raw.take() {
+            resource_log!("Deallocate raw Buffer (destroyed) {:?}", self.label());
+            unsafe {
+                use hal::Device;
+                self.device.raw().destroy_buffer(raw);
+            }
+        }
+    }
+}
+
 /// A temporary buffer, consumed by the command that uses it.
 ///
 /// A [`StagingBuffer`] is designed for one-shot uploads of data to the GPU. It