From bf99cf7f30e9c7ea0e879068773155ce18dfd0c0 Mon Sep 17 00:00:00 2001
From: Samson <16504129+sagudev@users.noreply.github.com>
Date: Thu, 20 Jun 2024 07:56:59 +0200
Subject: Proper GPUDevice cleanup (#32520)

* Make device cleanup right

* Use weakref for GPUDevice in globalscope

* No need to destroy device on drop

* DeviceReason early return

* make remove_gpu_device to be the only way to remove device
---
 .../script/dom/bindings/codegen/Bindings.conf      |  1 +
 components/script/dom/globalscope.rs               | 32 +++++++++++++++++-----
 components/script/dom/gpudevice.rs                 |  7 -----
 components/script/script_thread.rs                 | 11 ++++++--
 4 files changed, 35 insertions(+), 16 deletions(-)

(limited to 'components/script')

diff --git a/components/script/dom/bindings/codegen/Bindings.conf b/components/script/dom/bindings/codegen/Bindings.conf
index c7b4cbace88..6e10ffc6a3e 100644
--- a/components/script/dom/bindings/codegen/Bindings.conf
+++ b/components/script/dom/bindings/codegen/Bindings.conf
@@ -169,6 +169,7 @@ DOMInterfaces = {
 },
 
 'GPUDevice': {
+    'weakReferenceable': True, # for usage in GlobalScope https://github.com/servo/servo/issues/32519
     'inRealms': ['PopErrorScope', 'CreateComputePipelineAsync', 'CreateRenderPipelineAsync'],
 }
 
diff --git a/components/script/dom/globalscope.rs b/components/script/dom/globalscope.rs
index ca4f5e7cd76..27138a15d84 100644
--- a/components/script/dom/globalscope.rs
+++ b/components/script/dom/globalscope.rs
@@ -322,7 +322,7 @@ pub struct GlobalScope {
     gpu_id_hub: Arc<Mutex<Identities>>,
 
     /// WebGPU devices
-    gpu_devices: DomRefCell<HashMapTracedValues<WebGPUDevice, Dom<GPUDevice>>>,
+    gpu_devices: DomRefCell<HashMapTracedValues<WebGPUDevice, WeakRef<GPUDevice>>>,
 
     // https://w3c.github.io/performance-timeline/#supportedentrytypes-attribute
     #[ignore_malloc_size_of = "mozjs"]
@@ -3091,7 +3091,16 @@ impl GlobalScope {
     pub fn add_gpu_device(&self, device: &GPUDevice) {
         self.gpu_devices
             .borrow_mut()
-            .insert(device.id(), Dom::from_ref(device));
+            .insert(device.id(), WeakRef::new(device));
+    }
+
+    pub fn remove_gpu_device(&self, device: WebGPUDevice) {
+        let device = self
+            .gpu_devices
+            .borrow_mut()
+            .remove(&device)
+            .expect("GPUDevice should still be in devices hashmap");
+        assert!(device.root().is_none())
     }
 
     pub fn gpu_device_lost(&self, device: WebGPUDevice, reason: DeviceLostReason, msg: String) {
@@ -3100,15 +3109,24 @@ impl GlobalScope {
             DeviceLostReason::Destroyed => GPUDeviceLostReason::Destroyed,
         };
         let _ac = enter_realm(&*self);
-        self.gpu_devices
+        if let Some(device) = self
+            .gpu_devices
             .borrow_mut()
-            .remove(&device)
-            .expect("GPUDevice should still exists")
-            .lose(reason, msg);
+            .get_mut(&device)
+            .expect("GPUDevice should still be in devices hashmap")
+            .root()
+        {
+            device.lose(reason, msg);
+        }
     }
 
     pub fn handle_uncaptured_gpu_error(&self, device: WebGPUDevice, error: webgpu::Error) {
-        if let Some(gpu_device) = self.gpu_devices.borrow().get(&device) {
+        if let Some(gpu_device) = self
+            .gpu_devices
+            .borrow()
+            .get(&device)
+            .and_then(|device| device.root())
+        {
             gpu_device.fire_uncaptured_error(error);
         } else {
             warn!("Recived error for lost GPUDevice!")
diff --git a/components/script/dom/gpudevice.rs b/components/script/dom/gpudevice.rs
index 9e7c8c76b09..da1432e99b6 100644
--- a/components/script/dom/gpudevice.rs
+++ b/components/script/dom/gpudevice.rs
@@ -1031,13 +1031,6 @@ impl AsyncWGPUListener for GPUDevice {
 
 impl Drop for GPUDevice {
     fn drop(&mut self) {
-        if let Err(e) = self
-            .channel
-            .0
-            .send(WebGPURequest::DestroyDevice(self.device.0))
-        {
-            warn!("Failed to send DestroyDevice ({:?}) ({})", self.device.0, e);
-        }
         if let Err(e) = self
             .channel
             .0
diff --git a/components/script/script_thread.rs b/components/script/script_thread.rs
index 1d3fb58766a..6fe8f8f4ead 100644
--- a/components/script/script_thread.rs
+++ b/components/script/script_thread.rs
@@ -92,7 +92,7 @@ use style::dom::OpaqueNode;
 use style::thread_state::{self, ThreadState};
 use time::precise_time_ns;
 use url::Position;
-use webgpu::WebGPUMsg;
+use webgpu::{WebGPUDevice, WebGPUMsg};
 use webrender_api::DocumentId;
 use webrender_traits::WebRenderScriptApi;
 
@@ -2423,7 +2423,14 @@ impl ScriptThread {
     fn handle_msg_from_webgpu_server(&self, msg: WebGPUMsg) {
         match msg {
             WebGPUMsg::FreeAdapter(id) => self.gpu_id_hub.lock().kill_adapter_id(id),
-            WebGPUMsg::FreeDevice(id) => self.gpu_id_hub.lock().kill_device_id(id),
+            WebGPUMsg::FreeDevice {
+                device_id,
+                pipeline_id,
+            } => {
+                self.gpu_id_hub.lock().kill_device_id(device_id);
+                let global = self.documents.borrow().find_global(pipeline_id).unwrap();
+                global.remove_gpu_device(WebGPUDevice(device_id));
+            },
             WebGPUMsg::FreeBuffer(id) => self.gpu_id_hub.lock().kill_buffer_id(id),
             WebGPUMsg::FreePipelineLayout(id) => self.gpu_id_hub.lock().kill_pipeline_layout_id(id),
             WebGPUMsg::FreeComputePipeline(id) => {
-- 
cgit v1.2.3