Invoke a DeviceLostClosure immediately if set on an invalid device. (#5358)

Invoke a DeviceLostClosure immediately if set on an invalid device.

To make the device invalid, this defines an explicit, test-only method
make_invalid. It also modifies calls that expect to always retrieve a
valid device.

Co-authored-by: Erich Gubler <erichdongubler@gmail.com>

CHANGELOG.md

@@ -133,6 +133,7 @@ Bottom level categories:
 - Fix behavior of integer `clamp` when `min` argument > `max` argument. By @cwfitzgerald in [#5300](https://github.com/gfx-rs/wgpu/pull/5300).
 - Fix linking when targeting android. By @ashdnazg in [#5326](https://github.com/gfx-rs/wgpu/pull/5326).
 - fix resource leak for buffers/textures dropped while having pending writes. By @robtfm in [#5413](https://github.com/gfx-rs/wgpu/pull/5413)
+- Failing to set the device lost closure will call the closure before returning. By @bradwerth in [#5358](https://github.com/gfx-rs/wgpu/pull/5358).
 
 #### Naga
 - In spv-in, remove unnecessary "gl_PerVertex" name check so unused builtins will always be skipped. By @Imberflur in [#5227](https://github.com/gfx-rs/wgpu/pull/5227).

tests/tests/device.rs

@@ -1,3 +1,5 @@
+use std::sync::atomic::AtomicBool;
+
 use wgpu_test::{fail, gpu_test, FailureCase, GpuTestConfiguration, TestParameters};
 
 #[gpu_test]
@@ -518,12 +520,11 @@ static DEVICE_DESTROY_THEN_LOST: GpuTestConfiguration = GpuTestConfiguration::ne
     .run_async(|ctx| async move {
         // This test checks that when device.destroy is called, the provided
         // DeviceLostClosure is called with reason DeviceLostReason::Destroyed.
-        let was_called = std::sync::Arc::<std::sync::atomic::AtomicBool>::new(false.into());
+        static WAS_CALLED: AtomicBool = AtomicBool::new(false);
 
         // Set a LoseDeviceCallback on the device.
-        let was_called_clone = was_called.clone();
-        let callback = Box::new(move |reason, _m| {
-            was_called_clone.store(true, std::sync::atomic::Ordering::SeqCst);
+        let callback = Box::new(|reason, _m| {
+            WAS_CALLED.store(true, std::sync::atomic::Ordering::SeqCst);
             assert!(
                 matches!(reason, wgt::DeviceLostReason::Destroyed),
                 "Device lost info reason should match DeviceLostReason::Destroyed."
@@ -542,7 +543,7 @@ static DEVICE_DESTROY_THEN_LOST: GpuTestConfiguration = GpuTestConfiguration::ne
             .is_queue_empty());
 
         assert!(
-            was_called.load(std::sync::atomic::Ordering::SeqCst),
+            WAS_CALLED.load(std::sync::atomic::Ordering::SeqCst),
             "Device lost callback should have been called."
         );
     });
@@ -554,20 +555,13 @@ static DEVICE_DROP_THEN_LOST: GpuTestConfiguration = GpuTestConfiguration::new()
         // This test checks that when the device is dropped (such as in a GC),
         // the provided DeviceLostClosure is called with reason DeviceLostReason::Unknown.
         // Fails on webgl because webgl doesn't implement drop.
-        let was_called = std::sync::Arc::<std::sync::atomic::AtomicBool>::new(false.into());
+        static WAS_CALLED: std::sync::atomic::AtomicBool = AtomicBool::new(false);
 
         // Set a LoseDeviceCallback on the device.
-        let was_called_clone = was_called.clone();
-        let callback = Box::new(move |reason, message| {
-            was_called_clone.store(true, std::sync::atomic::Ordering::SeqCst);
-            assert!(
-                matches!(reason, wgt::DeviceLostReason::Dropped),
-                "Device lost info reason should match DeviceLostReason::Dropped."
-            );
-            assert!(
-                message == "Device dropped.",
-                "Device lost info message should be \"Device dropped.\"."
-            );
+        let callback = Box::new(|reason, message| {
+            WAS_CALLED.store(true, std::sync::atomic::Ordering::SeqCst);
+            assert_eq!(reason, wgt::DeviceLostReason::Dropped);
+            assert_eq!(message, "Device dropped.");
         });
         ctx.device.set_device_lost_callback(callback);
 
@@ -575,7 +569,34 @@ static DEVICE_DROP_THEN_LOST: GpuTestConfiguration = GpuTestConfiguration::new()
         drop(ctx.device);
 
         assert!(
-            was_called.load(std::sync::atomic::Ordering::SeqCst),
+            WAS_CALLED.load(std::sync::atomic::Ordering::SeqCst),
+            "Device lost callback should have been called."
+        );
+    });
+
+#[gpu_test]
+static DEVICE_INVALID_THEN_SET_LOST_CALLBACK: GpuTestConfiguration = GpuTestConfiguration::new()
+    .parameters(TestParameters::default().expect_fail(FailureCase::webgl2()))
+    .run_sync(|ctx| {
+        // This test checks that when the device is invalid, a subsequent call
+        // to set the device lost callback will immediately call the callback.
+        // Invalidating the device is done via a testing-only method. Fails on
+        // webgl because webgl doesn't implement make_invalid.
+
+        // Make the device invalid.
+        ctx.device.make_invalid();
+
+        static WAS_CALLED: AtomicBool = AtomicBool::new(false);
+
+        // Set a LoseDeviceCallback on the device.
+        let callback = Box::new(|reason, _m| {
+            WAS_CALLED.store(true, std::sync::atomic::Ordering::SeqCst);
+            assert_eq!(reason, wgt::DeviceLostReason::DeviceInvalid);
+        });
+        ctx.device.set_device_lost_callback(callback);
+
+        assert!(
+            WAS_CALLED.load(std::sync::atomic::Ordering::SeqCst),
             "Device lost callback should have been called."
         );
     });
@@ -586,16 +607,12 @@ static DEVICE_LOST_REPLACED_CALLBACK: GpuTestConfiguration = GpuTestConfiguratio
     .run_sync(|ctx| {
         // This test checks that a device_lost_callback is called when it is
         // replaced by another callback.
-        let was_called = std::sync::Arc::<std::sync::atomic::AtomicBool>::new(false.into());
+        static WAS_CALLED: AtomicBool = AtomicBool::new(false);
 
         // Set a LoseDeviceCallback on the device.
-        let was_called_clone = was_called.clone();
-        let callback = Box::new(move |reason, _m| {
-            was_called_clone.store(true, std::sync::atomic::Ordering::SeqCst);
-            assert!(
-                matches!(reason, wgt::DeviceLostReason::ReplacedCallback),
-                "Device lost info reason should match DeviceLostReason::ReplacedCallback."
-            );
+        let callback = Box::new(|reason, _m| {
+            WAS_CALLED.store(true, std::sync::atomic::Ordering::SeqCst);
+            assert_eq!(reason, wgt::DeviceLostReason::ReplacedCallback);
         });
         ctx.device.set_device_lost_callback(callback);
 
@@ -604,7 +621,7 @@ static DEVICE_LOST_REPLACED_CALLBACK: GpuTestConfiguration = GpuTestConfiguratio
         ctx.device.set_device_lost_callback(replacement_callback);
 
         assert!(
-            was_called.load(std::sync::atomic::Ordering::SeqCst),
+            WAS_CALLED.load(std::sync::atomic::Ordering::SeqCst),
             "Device lost callback should have been called."
         );
     });
@@ -619,21 +636,13 @@ static DROPPED_GLOBAL_THEN_DEVICE_LOST: GpuTestConfiguration = GpuTestConfigurat
         // wgpu without providing a more orderly shutdown. In such a case, the
         // device lost callback should be invoked with the message "Device is
         // dying."
-        let was_called = std::sync::Arc::<std::sync::atomic::AtomicBool>::new(false.into());
+        static WAS_CALLED: AtomicBool = AtomicBool::new(false);
 
         // Set a LoseDeviceCallback on the device.
-        let was_called_clone = was_called.clone();
-        let callback = Box::new(move |reason, message| {
-            was_called_clone.store(true, std::sync::atomic::Ordering::SeqCst);
-            assert!(
-                matches!(reason, wgt::DeviceLostReason::Dropped),
-                "Device lost info reason should match DeviceLostReason::Dropped."
-            );
-            assert!(
-                message == "Device is dying.",
-                "Device lost info message is \"{}\" and it should be \"Device is dying.\".",
-                message
-            );
+        let callback = Box::new(|reason, message| {
+            WAS_CALLED.store(true, std::sync::atomic::Ordering::SeqCst);
+            assert_eq!(reason, wgt::DeviceLostReason::Dropped);
+            assert_eq!(message, "Device is dying.");
         });
         ctx.device.set_device_lost_callback(callback);
 
@@ -641,7 +650,7 @@ static DROPPED_GLOBAL_THEN_DEVICE_LOST: GpuTestConfiguration = GpuTestConfigurat
 
         // Confirm that the callback was invoked.
         assert!(
-            was_called.load(std::sync::atomic::Ordering::SeqCst),
+            WAS_CALLED.load(std::sync::atomic::Ordering::SeqCst),
             "Device lost callback should have been called."
         );
     });

wgpu-core/src/device/global.rs

@@ -2240,6 +2240,15 @@ impl Global {
         }
     }
 
+    // This is a test-only function to force the device into an
+    // invalid state by inserting an error value in its place in
+    // the registry.
+    pub fn device_make_invalid<A: HalApi>(&self, device_id: DeviceId) {
+        let hub = A::hub(self);
+        hub.devices
+            .force_replace_with_error(device_id, "Made invalid.");
+    }
+
     pub fn device_drop<A: HalApi>(&self, device_id: DeviceId) {
         profiling::scope!("Device::drop");
         api_log!("Device::drop {device_id:?}");
@@ -2275,7 +2284,7 @@ impl Global {
     ) {
         let hub = A::hub(self);
 
-        if let Ok(device) = hub.devices.get(device_id) {
+        if let Ok(Some(device)) = hub.devices.try_get(device_id) {
             let mut life_tracker = device.lock_life();
             if let Some(existing_closure) = life_tracker.device_lost_closure.take() {
                 // It's important to not hold the lock while calling the closure.
@@ -2284,6 +2293,12 @@ impl Global {
                 life_tracker = device.lock_life();
             }
             life_tracker.device_lost_closure = Some(device_lost_closure);
+        } else {
+            // No device? Okay. Just like we have to call any existing closure
+            // before we drop it, we need to call this closure before we exit
+            // this function, because there's no device that is ever going to
+            // call it.
+            device_lost_closure.call(DeviceLostReason::DeviceInvalid, "".to_string());
         }
     }
 

wgpu-types/src/lib.rs

@@ -7190,7 +7190,7 @@ mod send_sync {
 ///
 /// Corresponds to [WebGPU `GPUDeviceLostReason`](https://gpuweb.github.io/gpuweb/#enumdef-gpudevicelostreason).
 #[repr(u8)]
-#[derive(Debug, Copy, Clone)]
+#[derive(Debug, Copy, Clone, Eq, PartialEq)]
 pub enum DeviceLostReason {
     /// Triggered by driver
     Unknown = 0,
@@ -7210,4 +7210,10 @@ pub enum DeviceLostReason {
     /// exactly once before it is dropped, which helps with managing the
     /// memory owned by the callback.
     ReplacedCallback = 3,
+    /// When setting the callback, but the device is already invalid
+    ///
+    /// As above, when the callback is provided, wgpu guarantees that it
+    /// will eventually be called. If the device is already invalid, wgpu
+    /// will call the callback immediately, with this reason.
+    DeviceInvalid = 4,
 }

wgpu/src/backend/webgpu.rs

@@ -1948,6 +1948,11 @@ impl crate::context::Context for ContextWebGpu {
         create_identified(device_data.0.create_render_bundle_encoder(&mapped_desc))
     }
 
+    #[doc(hidden)]
+    fn device_make_invalid(&self, _device: &Self::DeviceId, _device_data: &Self::DeviceData) {
+        // Unimplemented
+    }
+
     fn device_drop(&self, _device: &Self::DeviceId, _device_data: &Self::DeviceData) {
         // Device is dropped automatically
     }

wgpu/src/backend/wgpu_core.rs

@@ -1346,14 +1346,17 @@ impl crate::Context for ContextWgpuCore {
             Err(e) => panic!("Error in Device::create_render_bundle_encoder: {e}"),
         }
     }
+    #[doc(hidden)]
+    fn device_make_invalid(&self, device: &Self::DeviceId, _device_data: &Self::DeviceData) {
+        wgc::gfx_select!(device => self.0.device_make_invalid(*device));
+    }
     #[cfg_attr(not(any(native, Emscripten)), allow(unused))]
     fn device_drop(&self, device: &Self::DeviceId, _device_data: &Self::DeviceData) {
         #[cfg(any(native, Emscripten))]
         {
-            match wgc::gfx_select!(device => self.0.device_poll(*device, wgt::Maintain::wait())) {
-                Ok(_) => {}
-                Err(err) => self.handle_error_fatal(err, "Device::drop"),
-            }
+            // Call device_poll, but don't check for errors. We have to use its
+            // return value, but we just drop it.
+            let _ = wgc::gfx_select!(device => self.0.device_poll(*device, wgt::Maintain::wait()));
             wgc::gfx_select!(device => self.0.device_drop(*device));
         }
     }

wgpu/src/context.rs

@@ -267,6 +267,8 @@ pub trait Context: Debug + WasmNotSendSync + Sized {
         device_data: &Self::DeviceData,
         desc: &RenderBundleEncoderDescriptor<'_>,
     ) -> (Self::RenderBundleEncoderId, Self::RenderBundleEncoderData);
+    #[doc(hidden)]
+    fn device_make_invalid(&self, device: &Self::DeviceId, device_data: &Self::DeviceData);
     fn device_drop(&self, device: &Self::DeviceId, device_data: &Self::DeviceData);
     fn device_set_device_lost_callback(
         &self,
@@ -1293,6 +1295,8 @@ pub(crate) trait DynContext: Debug + WasmNotSendSync {
         device_data: &crate::Data,
         desc: &RenderBundleEncoderDescriptor<'_>,
     ) -> (ObjectId, Box<crate::Data>);
+    #[doc(hidden)]
+    fn device_make_invalid(&self, device: &ObjectId, device_data: &crate::Data);
     fn device_drop(&self, device: &ObjectId, device_data: &crate::Data);
     fn device_set_device_lost_callback(
         &self,
@@ -2350,6 +2354,13 @@ where
         (render_bundle_encoder.into(), Box::new(data) as _)
     }
 
+    #[doc(hidden)]
+    fn device_make_invalid(&self, device: &ObjectId, device_data: &crate::Data) {
+        let device = <T::DeviceId>::from(*device);
+        let device_data = downcast_ref(device_data);
+        Context::device_make_invalid(self, &device, device_data)
+    }
+
     fn device_drop(&self, device: &ObjectId, device_data: &crate::Data) {
         let device = <T::DeviceId>::from(*device);
         let device_data = downcast_ref(device_data);

wgpu/src/lib.rs

@@ -2703,6 +2703,12 @@ impl Device {
             Box::new(callback),
         )
     }
+
+    /// Test-only function to make this device invalid.
+    #[doc(hidden)]
+    pub fn make_invalid(&self) {
+        DynContext::device_make_invalid(&*self.context, &self.id, self.data.as_ref())
+    }
 }
 
 impl Drop for Device {