Narrow locks in queue::write_buffer (#8146)

This commit is contained in:
Connor Fitzgerald 2025-08-25 11:54:30 -04:00 committed by GitHub
parent 1301f80688
commit 884cea5ad9
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -542,13 +542,10 @@ impl Queue {
return Ok(());
};
let snatch_guard = self.device.snatchable_lock.read();
// Platform validation requires that the staging buffer always be
// freed, even if an error occurs. All paths from here must call
// `device.pending_writes.consume`.
let mut staging_buffer = StagingBuffer::new(&self.device, data_size)?;
let mut pending_writes = self.pending_writes.lock();
let staging_buffer = {
profiling::scope!("copy");
@ -556,6 +553,9 @@ impl Queue {
staging_buffer.flush()
};
let snatch_guard = self.device.snatchable_lock.read();
let mut pending_writes = self.pending_writes.lock();
let result = self.write_staging_buffer_impl(
&snatch_guard,
&mut pending_writes,
@ -564,7 +564,12 @@ impl Queue {
buffer_offset,
);
drop(snatch_guard);
pending_writes.consume(staging_buffer);
drop(pending_writes);
result
}
@ -595,15 +600,15 @@ impl Queue {
let buffer = buffer.get()?;
let snatch_guard = self.device.snatchable_lock.read();
let mut pending_writes = self.pending_writes.lock();
// At this point, we have taken ownership of the staging_buffer from the
// user. Platform validation requires that the staging buffer always
// be freed, even if an error occurs. All paths from here must call
// `device.pending_writes.consume`.
let staging_buffer = staging_buffer.flush();
let snatch_guard = self.device.snatchable_lock.read();
let mut pending_writes = self.pending_writes.lock();
let result = self.write_staging_buffer_impl(
&snatch_guard,
&mut pending_writes,
@ -612,7 +617,12 @@ impl Queue {
buffer_offset,
);
drop(snatch_guard);
pending_writes.consume(staging_buffer);
drop(pending_writes);
result
}