Shorten Lock Lifetimes (#4976)

wgpu-core/src/command/render.rs

@@ -2370,12 +2370,12 @@ impl<G: GlobalIdentityHandlerFactory> Global<G> {
             (trackers, pending_discard_init_fixups)
         };
 
-        let query_set_guard = hub.query_sets.read();
-
         let cmd_buf = hub.command_buffers.get(encoder_id).unwrap();
         let mut cmd_buf_data = cmd_buf.data.lock();
         let cmd_buf_data = cmd_buf_data.as_mut().unwrap();
 
+        let query_set_guard = hub.query_sets.read();
+
         let encoder = &mut cmd_buf_data.encoder;
         let status = &mut cmd_buf_data.status;
         let tracker = &mut cmd_buf_data.trackers;

wgpu-core/src/command/transfer.rs

@@ -15,7 +15,6 @@ use crate::{
         TextureInitTrackerAction,
     },
     resource::{Resource, Texture, TextureErrorDimension},
-    storage::Storage,
     track::{TextureSelector, Tracker},
 };
 
@@ -24,7 +23,7 @@ use hal::CommandEncoder as _;
 use thiserror::Error;
 use wgt::{BufferAddress, BufferUsages, Extent3d, TextureUsages};
 
-use std::iter;
+use std::{iter, sync::Arc};
 
 use super::{memory_init::CommandBufferTextureMemoryActions, CommandEncoder};
 
@@ -447,9 +446,8 @@ fn handle_texture_init<A: HalApi>(
     device: &Device<A>,
     copy_texture: &ImageCopyTexture,
     copy_size: &Extent3d,
-    texture_guard: &Storage<Texture<A>, TextureId>,
+    texture: &Arc<Texture<A>>,
 ) {
-    let texture = texture_guard.get(copy_texture.texture).unwrap();
     let init_action = TextureInitTrackerAction {
         texture: texture.clone(),
         range: TextureInitRange {
@@ -494,12 +492,8 @@ fn handle_src_texture_init<A: HalApi>(
     device: &Device<A>,
     source: &ImageCopyTexture,
     copy_size: &Extent3d,
-    texture_guard: &Storage<Texture<A>, TextureId>,
+    texture: &Arc<Texture<A>>,
 ) -> Result<(), TransferError> {
-    let _ = texture_guard
-        .get(source.texture)
-        .map_err(|_| TransferError::InvalidTexture(source.texture))?;
-
     handle_texture_init(
         MemoryInitKind::NeedsInitializedMemory,
         encoder,
@@ -508,7 +502,7 @@ fn handle_src_texture_init<A: HalApi>(
         device,
         source,
         copy_size,
-        texture_guard,
+        texture,
     );
     Ok(())
 }
@@ -524,12 +518,8 @@ fn handle_dst_texture_init<A: HalApi>(
     device: &Device<A>,
     destination: &ImageCopyTexture,
     copy_size: &Extent3d,
-    texture_guard: &Storage<Texture<A>, TextureId>,
+    texture: &Arc<Texture<A>>,
 ) -> Result<(), TransferError> {
-    let texture = texture_guard
-        .get(destination.texture)
-        .map_err(|_| TransferError::InvalidTexture(destination.texture))?;
-
     // Attention: If we don't write full texture subresources, we need to a full
     // clear first since we don't track subrects. This means that in rare cases
     // even a *destination* texture of a transfer may need an immediate texture
@@ -552,7 +542,7 @@ fn handle_dst_texture_init<A: HalApi>(
         device,
         destination,
         copy_size,
-        texture_guard,
+        texture,
     );
     Ok(())
 }
@@ -764,14 +754,13 @@ impl<G: GlobalIdentityHandlerFactory> Global<G> {
         let buffer_memory_init_actions = &mut cmd_buf_data.buffer_memory_init_actions;
         let texture_memory_actions = &mut cmd_buf_data.texture_memory_actions;
 
-        let texture_guard = hub.textures.read();
-
         if copy_size.width == 0 || copy_size.height == 0 || copy_size.depth_or_array_layers == 0 {
             log::trace!("Ignoring copy_buffer_to_texture of size 0");
             return Ok(());
         }
 
-        let dst_texture = texture_guard
+        let dst_texture = hub
+            .textures
             .get(destination.texture)
             .map_err(|_| TransferError::InvalidTexture(destination.texture))?;
 
@@ -782,7 +771,7 @@ impl<G: GlobalIdentityHandlerFactory> Global<G> {
             copy_size,
         )?;
 
-        let (dst_range, dst_base) = extract_texture_selector(destination, copy_size, dst_texture)?;
+        let (dst_range, dst_base) = extract_texture_selector(destination, copy_size, &dst_texture)?;
 
         // Handle texture init *before* dealing with barrier transitions so we
         // have an easier time inserting "immediate-inits" that may be required
@@ -794,7 +783,7 @@ impl<G: GlobalIdentityHandlerFactory> Global<G> {
             device,
             destination,
             copy_size,
-            &texture_guard,
+            &dst_texture,
         )?;
 
         let snatch_guard = device.snatchable_lock.read();
@@ -820,7 +809,7 @@ impl<G: GlobalIdentityHandlerFactory> Global<G> {
 
         let dst_pending = tracker
             .textures
-            .set_single(dst_texture, dst_range, hal::TextureUses::COPY_DST)
+            .set_single(&dst_texture, dst_range, hal::TextureUses::COPY_DST)
             .ok_or(TransferError::InvalidTexture(destination.texture))?;
         let dst_inner = dst_texture.inner();
         let dst_raw = dst_inner
@@ -928,21 +917,20 @@ impl<G: GlobalIdentityHandlerFactory> Global<G> {
         let buffer_memory_init_actions = &mut cmd_buf_data.buffer_memory_init_actions;
         let texture_memory_actions = &mut cmd_buf_data.texture_memory_actions;
 
-        let texture_guard = hub.textures.read();
-
         if copy_size.width == 0 || copy_size.height == 0 || copy_size.depth_or_array_layers == 0 {
             log::trace!("Ignoring copy_texture_to_buffer of size 0");
             return Ok(());
         }
 
-        let src_texture = texture_guard
+        let src_texture = hub
+            .textures
             .get(source.texture)
             .map_err(|_| TransferError::InvalidTexture(source.texture))?;
 
         let (hal_copy_size, array_layer_count) =
             validate_texture_copy_range(source, &src_texture.desc, CopySide::Source, copy_size)?;
 
-        let (src_range, src_base) = extract_texture_selector(source, copy_size, src_texture)?;
+        let (src_range, src_base) = extract_texture_selector(source, copy_size, &src_texture)?;
 
         // Handle texture init *before* dealing with barrier transitions so we
         // have an easier time inserting "immediate-inits" that may be required
@@ -954,14 +942,14 @@ impl<G: GlobalIdentityHandlerFactory> Global<G> {
             device,
             source,
             copy_size,
-            &texture_guard,
+            &src_texture,
         )?;
 
         let snatch_guard = device.snatchable_lock.read();
 
         let src_pending = tracker
             .textures
-            .set_single(src_texture, src_range, hal::TextureUses::COPY_SRC)
+            .set_single(&src_texture, src_range, hal::TextureUses::COPY_SRC)
             .ok_or(TransferError::InvalidTexture(source.texture))?;
         let src_inner = src_texture.inner();
         let src_raw = src_inner
@@ -1104,18 +1092,18 @@ impl<G: GlobalIdentityHandlerFactory> Global<G> {
         let tracker = &mut cmd_buf_data.trackers;
         let texture_memory_actions = &mut cmd_buf_data.texture_memory_actions;
 
-        let texture_guard = hub.textures.read();
-
         if copy_size.width == 0 || copy_size.height == 0 || copy_size.depth_or_array_layers == 0 {
             log::trace!("Ignoring copy_texture_to_texture of size 0");
             return Ok(());
         }
 
-        let src_texture = texture_guard
+        let src_texture = hub
+            .textures
             .get(source.texture)
             .map_err(|_| TransferError::InvalidTexture(source.texture))?;
         let src_inner = src_texture.inner();
-        let dst_texture = texture_guard
+        let dst_texture = hub
+            .textures
             .get(destination.texture)
             .map_err(|_| TransferError::InvalidTexture(source.texture))?;
         let dst_inner = dst_texture.inner();
@@ -1141,9 +1129,9 @@ impl<G: GlobalIdentityHandlerFactory> Global<G> {
             copy_size,
         )?;
 
-        let (src_range, src_tex_base) = extract_texture_selector(source, copy_size, src_texture)?;
+        let (src_range, src_tex_base) = extract_texture_selector(source, copy_size, &src_texture)?;
         let (dst_range, dst_tex_base) =
-            extract_texture_selector(destination, copy_size, dst_texture)?;
+            extract_texture_selector(destination, copy_size, &dst_texture)?;
         let src_texture_aspects = hal::FormatAspects::from(src_texture.desc.format);
         let dst_texture_aspects = hal::FormatAspects::from(dst_texture.desc.format);
         if src_tex_base.aspect != src_texture_aspects {
@@ -1163,7 +1151,7 @@ impl<G: GlobalIdentityHandlerFactory> Global<G> {
             device,
             source,
             copy_size,
-            &texture_guard,
+            &src_texture,
         )?;
         handle_dst_texture_init(
             encoder,
@@ -1172,13 +1160,13 @@ impl<G: GlobalIdentityHandlerFactory> Global<G> {
             device,
             destination,
             copy_size,
-            &texture_guard,
+            &dst_texture,
         )?;
 
         let src_pending = cmd_buf_data
             .trackers
             .textures
-            .set_single(src_texture, src_range, hal::TextureUses::COPY_SRC)
+            .set_single(&src_texture, src_range, hal::TextureUses::COPY_SRC)
             .ok_or(TransferError::InvalidTexture(source.texture))?;
         let src_raw = src_inner
             .as_ref()
@@ -1198,7 +1186,7 @@ impl<G: GlobalIdentityHandlerFactory> Global<G> {
         let dst_pending = cmd_buf_data
             .trackers
             .textures
-            .set_single(dst_texture, dst_range, hal::TextureUses::COPY_DST)
+            .set_single(&dst_texture, dst_range, hal::TextureUses::COPY_DST)
             .ok_or(TransferError::InvalidTexture(destination.texture))?;
         let dst_raw = dst_inner
             .as_ref()

wgpu-core/src/device/global.rs

@@ -487,12 +487,11 @@ impl<G: GlobalIdentityHandlerFactory> Global<G> {
 
         let hub = A::hub(self);
 
-        let buffer = {
-            let mut buffer_guard = hub.buffers.write();
-            buffer_guard
+        let buffer = hub
+            .buffers
+            .write()
             .get_and_mark_destroyed(buffer_id)
-                .map_err(|_| resource::DestroyError::Invalid)?
-        };
+            .map_err(|_| resource::DestroyError::Invalid)?;
 
         let _ = buffer.unmap();
 
@@ -683,8 +682,7 @@ impl<G: GlobalIdentityHandlerFactory> Global<G> {
         let fid = hub.buffers.prepare::<G>(id_in);
 
         let error = loop {
-            let device_guard = hub.devices.read();
-            let device = match device_guard.get(device_id) {
+            let device = match hub.devices.get(device_id) {
                 Ok(device) => device,
                 Err(_) => break DeviceError::Invalid.into(),
             };
@@ -732,8 +730,9 @@ impl<G: GlobalIdentityHandlerFactory> Global<G> {
 
         let hub = A::hub(self);
 
-        let mut texture_guard = hub.textures.write();
-        let texture = texture_guard
+        let texture = hub
+            .textures
+            .write()
             .get_and_mark_destroyed(texture_id)
             .map_err(|_| resource::DestroyError::Invalid)?;
 
@@ -1075,12 +1074,9 @@ impl<G: GlobalIdentityHandlerFactory> Global<G> {
                 trace.add(trace::Action::CreatePipelineLayout(fid.id(), desc.clone()));
             }
 
-            let layout = {
-                let bgl_guard = hub.bind_group_layouts.read();
-                match device.create_pipeline_layout(desc, &*bgl_guard) {
+            let layout = match device.create_pipeline_layout(desc, &hub.bind_group_layouts) {
                 Ok(layout) => layout,
                 Err(e) => break e,
-                }
             };
 
             let (id, _) = fid.assign(layout);

wgpu-core/src/device/queue.rs

@@ -1427,9 +1427,7 @@ impl<G: GlobalIdentityHandlerFactory> Global<G> {
             let pending_writes = pending_writes.as_mut().unwrap();
 
             {
-                let texture_guard = hub.textures.read();
-
-                used_surface_textures.set_size(texture_guard.len());
+                used_surface_textures.set_size(hub.textures.read().len());
                 for (&id, texture) in pending_writes.dst_textures.iter() {
                     match *texture.inner().as_ref().unwrap() {
                         TextureInner::Native { raw: None } => {

wgpu-core/src/device/resource.rs

@@ -2357,7 +2357,7 @@ impl<A: HalApi> Device<A> {
     pub(crate) fn create_pipeline_layout(
         self: &Arc<Self>,
         desc: &binding_model::PipelineLayoutDescriptor,
-        bgl_guard: &Storage<BindGroupLayout<A>, id::BindGroupLayoutId>,
+        bgl_registry: &Registry<id::BindGroupLayoutId, BindGroupLayout<A>>,
     ) -> Result<binding_model::PipelineLayout<A>, binding_model::CreatePipelineLayoutError> {
         use crate::binding_model::CreatePipelineLayoutError as Error;
 
@@ -2410,31 +2410,38 @@ impl<A: HalApi> Device<A> {
 
         let mut count_validator = binding_model::BindingTypeMaxCountValidator::default();
 
-        // validate total resource counts
+        // Collect references to the BGLs
+        let mut bind_group_layouts = ArrayVec::new();
         for &id in desc.bind_group_layouts.iter() {
-            let Ok(bind_group_layout) = bgl_guard.get(id) else {
+            let Ok(bgl) = bgl_registry.get(id) else {
                 return Err(Error::InvalidBindGroupLayout(id));
             };
 
-            if bind_group_layout.device.as_info().id() != self.as_info().id() {
+            bind_group_layouts.push(bgl);
+        }
+
+        // Validate total resource counts and check for a matching device
+        for bgl in &bind_group_layouts {
+            if bgl.device.as_info().id() != self.as_info().id() {
                 return Err(DeviceError::WrongDevice.into());
             }
 
-            count_validator.merge(&bind_group_layout.binding_count_validator);
+            count_validator.merge(&bgl.binding_count_validator);
         }
+
         count_validator
             .validate(&self.limits)
             .map_err(Error::TooManyBindings)?;
 
-        let bgl_vec = desc
-            .bind_group_layouts
+        let raw_bind_group_layouts = bind_group_layouts
             .iter()
-            .map(|&id| bgl_guard.get(id).unwrap().raw())
-            .collect::<Vec<_>>();
+            .map(|bgl| bgl.raw())
+            .collect::<ArrayVec<_, { hal::MAX_BIND_GROUPS }>>();
+
         let hal_desc = hal::PipelineLayoutDescriptor {
             label: desc.label.to_hal(self.instance_flags),
             flags: hal::PipelineLayoutFlags::FIRST_VERTEX_INSTANCE,
-            bind_group_layouts: &bgl_vec,
+            bind_group_layouts: &raw_bind_group_layouts,
             push_constant_ranges: desc.push_constant_ranges.as_ref(),
         };
 
@@ -2446,15 +2453,13 @@ impl<A: HalApi> Device<A> {
                 .map_err(DeviceError::from)?
         };
 
+        drop(raw_bind_group_layouts);
+
         Ok(binding_model::PipelineLayout {
             raw: Some(raw),
             device: self.clone(),
             info: ResourceInfo::new(desc.label.borrow_or_default()),
-            bind_group_layouts: desc
-                .bind_group_layouts
-                .iter()
-                .map(|&id| bgl_guard.get(id).unwrap().clone())
-                .collect(),
+            bind_group_layouts,
             push_constant_ranges: desc.push_constant_ranges.iter().cloned().collect(),
         })
     }
@@ -2495,7 +2500,7 @@ impl<A: HalApi> Device<A> {
             bind_group_layouts: Cow::Borrowed(&ids.group_ids[..group_count]),
             push_constant_ranges: Cow::Borrowed(&[]), //TODO?
         };
-        let layout = self.create_pipeline_layout(&layout_desc, &bgl_registry.read())?;
+        let layout = self.create_pipeline_layout(&layout_desc, bgl_registry)?;
         pipeline_layout_registry.force_replace(ids.root_id, layout);
         Ok(pipeline_layout_registry.get(ids.root_id).unwrap())
     }