From 286306d7e8051e5e7a289608ff2e1338bf633354 Mon Sep 17 00:00:00 2001
From: Nicolas Silva <nical@fastmail.com>
Date: Thu, 15 Feb 2024 21:42:28 +0100
Subject: [PATCH] Store the device's queue via a weak ref instead of an ID
 (#5230)

---
 wgpu-core/src/device/global.rs   |  5 ++---
 wgpu-core/src/device/resource.rs | 24 ++++++++++++++++--------
 wgpu-core/src/instance.rs        |  8 ++++----
 wgpu-core/src/present.rs         |  3 +--
 4 files changed, 23 insertions(+), 17 deletions(-)

diff --git a/wgpu-core/src/device/global.rs b/wgpu-core/src/device/global.rs
index fe2a5de041..e22b6e8392 100644
--- a/wgpu-core/src/device/global.rs
+++ b/wgpu-core/src/device/global.rs
@@ -1330,9 +1330,8 @@ impl Global {
             if !device.is_valid() {
                 break DeviceError::Lost;
             }
-            let queue = match hub.queues.get(device.queue_id.read().unwrap()) {
-                Ok(queue) => queue,
-                Err(_) => break DeviceError::InvalidQueueId,
+            let Some(queue) = device.get_queue() else {
+                break DeviceError::InvalidQueueId;
             };
             let encoder = match device
                 .command_allocator
diff --git a/wgpu-core/src/device/resource.rs b/wgpu-core/src/device/resource.rs
index d4432a177a..ca9f6dd3cf 100644
--- a/wgpu-core/src/device/resource.rs
+++ b/wgpu-core/src/device/resource.rs
@@ -13,7 +13,6 @@ use crate::{
     hal_api::HalApi,
     hal_label,
     hub::Hub,
-    id::QueueId,
     init_tracker::{
         BufferInitTracker, BufferInitTrackerAction, MemoryInitKind, TextureInitRange,
         TextureInitTracker, TextureInitTrackerAction,
@@ -38,6 +37,7 @@ use crate::{
 
 use arrayvec::ArrayVec;
 use hal::{CommandEncoder as _, Device as _};
+use once_cell::sync::OnceCell;
 use parking_lot::{Mutex, MutexGuard, RwLock};
 
 use smallvec::SmallVec;
@@ -56,7 +56,7 @@ use std::{
 
 use super::{
     life::{self, ResourceMaps},
-    queue::{self},
+    queue::{self, Queue},
     DeviceDescriptor, DeviceError, ImplicitPipelineContext, UserClosures, ENTRYPOINT_FAILURE_ERROR,
     IMPLICIT_BIND_GROUP_LAYOUT_ERROR_LABEL, ZERO_BUFFER_SIZE,
 };
@@ -89,8 +89,8 @@ use super::{
 pub struct Device<A: HalApi> {
     raw: Option<A::Device>,
     pub(crate) adapter: Arc<Adapter<A>>,
-    pub(crate) queue_id: RwLock<Option<QueueId>>,
-    queue_to_drop: RwLock<Option<A::Queue>>,
+    pub(crate) queue: OnceCell<Weak<Queue<A>>>,
+    queue_to_drop: OnceCell<A::Queue>,
     pub(crate) zero_buffer: Option<A::Buffer>,
     pub(crate) info: ResourceInfo<Device<A>>,
 
@@ -162,7 +162,7 @@ impl<A: HalApi> Drop for Device<A> {
         unsafe {
             raw.destroy_buffer(self.zero_buffer.take().unwrap());
             raw.destroy_fence(self.fence.write().take().unwrap());
-            let queue = self.queue_to_drop.write().take().unwrap();
+            let queue = self.queue_to_drop.take().unwrap();
             raw.exit(queue);
         }
     }
@@ -260,8 +260,8 @@ impl<A: HalApi> Device<A> {
         Ok(Self {
             raw: Some(raw_device),
             adapter: adapter.clone(),
-            queue_id: RwLock::new(None),
-            queue_to_drop: RwLock::new(None),
+            queue: OnceCell::new(),
+            queue_to_drop: OnceCell::new(),
             zero_buffer: Some(zero_buffer),
             info: ResourceInfo::new("<device>"),
             command_allocator: Mutex::new(Some(com_alloc)),
@@ -302,7 +302,7 @@ impl<A: HalApi> Device<A> {
     }
 
     pub(crate) fn release_queue(&self, queue: A::Queue) {
-        self.queue_to_drop.write().replace(queue);
+        assert!(self.queue_to_drop.set(queue).is_ok());
     }
 
     pub(crate) fn lock_life<'a>(&'a self) -> MutexGuard<'a, LifetimeTracker<A>> {
@@ -359,6 +359,14 @@ impl<A: HalApi> Device<A> {
         }
     }
 
+    pub fn get_queue(&self) -> Option<Arc<Queue<A>>> {
+        self.queue.get().as_ref()?.upgrade()
+    }
+
+    pub fn set_queue(&self, queue: Arc<Queue<A>>) {
+        assert!(self.queue.set(Arc::downgrade(&queue)).is_ok());
+    }
+
     /// Check this device for completed commands.
     ///
     /// The `maintain` argument tells how the maintence function should behave, either
diff --git a/wgpu-core/src/instance.rs b/wgpu-core/src/instance.rs
index 582571c2b8..95d56ad1ff 100644
--- a/wgpu-core/src/instance.rs
+++ b/wgpu-core/src/instance.rs
@@ -1072,10 +1072,10 @@ impl Global {
             let device = hub.devices.get(device_id).unwrap();
             queue.device = Some(device.clone());
 
-            let (queue_id, _) = queue_fid.assign(queue);
+            let (queue_id, queue) = queue_fid.assign(queue);
             resource_log!("Created Queue {:?}", queue_id);
 
-            device.queue_id.write().replace(queue_id);
+            device.set_queue(queue);
 
             return (device_id, queue_id, None);
         };
@@ -1124,10 +1124,10 @@ impl Global {
             let device = hub.devices.get(device_id).unwrap();
             queue.device = Some(device.clone());
 
-            let (queue_id, _) = queues_fid.assign(queue);
+            let (queue_id, queue) = queues_fid.assign(queue);
             resource_log!("Created Queue {:?}", queue_id);
 
-            device.queue_id.write().replace(queue_id);
+            device.set_queue(queue);
 
             return (device_id, queue_id, None);
         };
diff --git a/wgpu-core/src/present.rs b/wgpu-core/src/present.rs
index ed1810a653..ce270384a4 100644
--- a/wgpu-core/src/present.rs
+++ b/wgpu-core/src/present.rs
@@ -294,8 +294,7 @@ impl Global {
         if !device.is_valid() {
             return Err(DeviceError::Lost.into());
         }
-        let queue_id = device.queue_id.read().unwrap();
-        let queue = hub.queues.get(queue_id).unwrap();
+        let queue = device.get_queue().unwrap();
 
         #[cfg(feature = "trace")]
         if let Some(ref mut trace) = *device.trace.lock() {