From c144f2a6971306b4a84cc597fd7d6ace2a37ff5a Mon Sep 17 00:00:00 2001
From: teoxoy <28601907+teoxoy@users.noreply.github.com>
Date: Mon, 31 Mar 2025 17:29:57 +0200
Subject: [PATCH] invalidate `Device` on OOM errors with the exception of
 buffer, texture, query set and acceleration structure creation

---
 wgpu-core/src/device/ray_tracing.rs |  8 ++++----
 wgpu-core/src/device/resource.rs    | 19 +++++++++++++------
 2 files changed, 17 insertions(+), 10 deletions(-)

diff --git a/wgpu-core/src/device/ray_tracing.rs b/wgpu-core/src/device/ray_tracing.rs
index 0146741254..76c94f948a 100644
--- a/wgpu-core/src/device/ray_tracing.rs
+++ b/wgpu-core/src/device/ray_tracing.rs
@@ -8,7 +8,7 @@ use crate::lock::rank;
 use crate::resource::{Fallible, TrackingData};
 use crate::snatch::Snatchable;
 use crate::{
-    device::{Device, DeviceError},
+    device::Device,
     global::Global,
     id::{self, BlasId, TlasId},
     lock::RwLock,
@@ -109,7 +109,7 @@ impl Device {
                     allow_compaction: false,
                 })
         }
-        .map_err(DeviceError::from_hal)?;
+        .map_err(|e| self.handle_hal_error_with_nonfatal_oom(e))?;
 
         let handle = unsafe {
             self.raw()
@@ -177,7 +177,7 @@ impl Device {
                     allow_compaction: false,
                 })
         }
-        .map_err(DeviceError::from_hal)?;
+        .map_err(|e| self.handle_hal_error_with_nonfatal_oom(e))?;
 
         let instance_buffer_size =
             self.alignments.raw_tlas_instance_size * desc.max_instances.max(1) as usize;
@@ -190,7 +190,7 @@ impl Device {
                 memory_flags: hal::MemoryFlags::PREFER_COHERENT,
             })
         }
-        .map_err(DeviceError::from_hal)?;
+        .map_err(|e| self.handle_hal_error_with_nonfatal_oom(e))?;
 
         Ok(Arc::new(resource::Tlas {
             raw: Snatchable::new(raw),
diff --git a/wgpu-core/src/device/resource.rs b/wgpu-core/src/device/resource.rs
index 7b54a78809..b715fa828d 100644
--- a/wgpu-core/src/device/resource.rs
+++ b/wgpu-core/src/device/resource.rs
@@ -363,8 +363,8 @@ impl Device {
 
     pub fn handle_hal_error(&self, error: hal::DeviceError) -> DeviceError {
         match error {
-            hal::DeviceError::OutOfMemory => {}
-            hal::DeviceError::Lost
+            hal::DeviceError::OutOfMemory
+            | hal::DeviceError::Lost
             | hal::DeviceError::ResourceCreationFailed
             | hal::DeviceError::Unexpected => {
                 self.lose(&error.to_string());
@@ -373,6 +373,13 @@ impl Device {
         DeviceError::from_hal(error)
     }
 
+    pub fn handle_hal_error_with_nonfatal_oom(&self, error: hal::DeviceError) -> DeviceError {
+        match error {
+            hal::DeviceError::OutOfMemory => DeviceError::from_hal(error),
+            error => self.handle_hal_error(error),
+        }
+    }
+
     /// Run some destroy operations that were deferred.
     ///
     /// Destroying the resources requires taking a write lock on the device's snatch lock,
@@ -679,8 +686,8 @@ impl Device {
             usage,
             memory_flags: hal::MemoryFlags::empty(),
         };
-        let buffer =
-            unsafe { self.raw().create_buffer(&hal_desc) }.map_err(|e| self.handle_hal_error(e))?;
+        let buffer = unsafe { self.raw().create_buffer(&hal_desc) }
+            .map_err(|e| self.handle_hal_error_with_nonfatal_oom(e))?;
 
         let timestamp_normalization_bind_group = Snatchable::new(
             self.timestamp_normalizer
@@ -1100,7 +1107,7 @@ impl Device {
         };
 
         let raw_texture = unsafe { self.raw().create_texture(&hal_desc) }
-            .map_err(|e| self.handle_hal_error(e))?;
+            .map_err(|e| self.handle_hal_error_with_nonfatal_oom(e))?;
 
         let clear_mode = if hal_usage
             .intersects(wgt::TextureUses::DEPTH_STENCIL_WRITE | wgt::TextureUses::COLOR_TARGET)
@@ -3862,7 +3869,7 @@ impl Device {
         let hal_desc = desc.map_label(|label| label.to_hal(self.instance_flags));
 
         let raw = unsafe { self.raw().create_query_set(&hal_desc) }
-            .map_err(|e| self.handle_hal_error(e))?;
+            .map_err(|e| self.handle_hal_error_with_nonfatal_oom(e))?;
 
         let query_set = QuerySet {
             raw: ManuallyDrop::new(raw),