Fix allocation leaks on error paths

Author: cwfitzgerald

wgpu-hal/src/dx12/suballocation.rs

@@ -1076,12 +1076,17 @@ impl DeviceAllocationContext<'_> {
             name,
         )?;
 
+        self.counters.buffer_memory.add(allocation.size() as isize);
+
         if let Some(label) = desc.label {
-            resource.set_name(label)?;
+            // If naming fails, free the resource (which also subtracts the counter
+            // added above) so the allocation is not leaked.
+            if let Err(e) = resource.set_name(label) {
+                self.free_resource(resource, allocation);
+                return Err(e);
+            }
         }
 
-        self.counters.buffer_memory.add(allocation.size() as isize);
-
         Ok((resource, allocation))
     }
 
@@ -1110,12 +1115,17 @@ impl DeviceAllocationContext<'_> {
             name,
         )?;
 
+        self.counters.texture_memory.add(allocation.size() as isize);
+
         if let Some(label) = desc.label {
-            resource.set_name(label)?;
+            // If naming fails, free the resource (which also subtracts the counter
+            // added above) so the allocation is not leaked.
+            if let Err(e) = resource.set_name(label) {
+                self.free_resource(resource, allocation);
+                return Err(e);
+            }
         }
 
-        self.counters.texture_memory.add(allocation.size() as isize);
-
         Ok((resource, allocation))
     }
 
@@ -1144,14 +1154,19 @@ impl DeviceAllocationContext<'_> {
             name,
         )?;
 
-        if let Some(label) = desc.label {
-            resource.set_name(label)?;
-        }
-
         self.counters
             .acceleration_structure_memory
             .add(allocation.size() as isize);
 
+        if let Some(label) = desc.label {
+            // If naming fails, free the resource (which also subtracts the counter
+            // added above) so the allocation is not leaked.
+            if let Err(e) = resource.set_name(label) {
+                self.free_resource(resource, allocation);
+                return Err(e);
+            }
+        }
+
         Ok((resource, allocation))
     }
 

wgpu-hal/src/vulkan/device.rs

@@ -1173,20 +1173,21 @@ impl crate::Device for super::Device {
                 unsafe { self.shared.raw.destroy_image(image.raw, None) };
             })?;
 
-        self.counters.texture_memory.add(allocation.size() as isize);
-
         if let Err(e) = unsafe {
             self.shared
                 .raw
                 .bind_image_memory(image.raw, allocation.memory(), allocation.offset())
         } {
             // Bind failed: free the allocation back to the allocator before destroying the
-            // image, so the pooled slot / dedicated memory is not leaked.
+            // image, so the pooled slot / dedicated memory is not leaked. The counter is
+            // only added after a successful bind, so nothing needs to be subtracted here.
             self.mem_allocator.lock().free(&self.shared, allocation);
             unsafe { self.shared.raw.destroy_image(image.raw, None) };
             return Err(super::map_host_device_oom_err(e));
         }
 
+        self.counters.texture_memory.add(allocation.size() as isize);
+
         Ok(unsafe {
             self.texture_from_raw(
                 image.raw,
@@ -2836,13 +2837,20 @@ impl crate::Device for super::Device {
                 .size(desc.size)
                 .ty(conv::map_acceleration_structure_format(desc.format));
 
-            let raw_acceleration_structure = ray_tracing_functions
+            let raw_acceleration_structure = match ray_tracing_functions
                 .acceleration_structure
                 .create_acceleration_structure(&vk_info, None)
-                .map_err(super::map_host_oom_and_ioca_err)
-                .inspect_err(|_| {
+            {
+                Ok(raw) => raw,
+                Err(e) => {
+                    // Creation failed: free the allocation back to the allocator before
+                    // destroying the buffer, so the pooled slot / dedicated memory is not
+                    // leaked.
+                    self.mem_allocator.lock().free(&self.shared, allocation);
                     self.shared.raw.destroy_buffer(raw_buffer, None);
-                })?;
+                    return Err(super::map_host_oom_and_ioca_err(e));
+                }
+            };
 
             if let Some(label) = desc.label {
                 self.shared
@@ -2854,18 +2862,20 @@ impl crate::Device for super::Device {
                     .query_type(vk::QueryType::ACCELERATION_STRUCTURE_COMPACTED_SIZE_KHR)
                     .query_count(1);
 
-                let raw = self
-                    .shared
-                    .raw
-                    .create_query_pool(&vk_info, None)
-                    .map_err(super::map_host_device_oom_err)
-                    .inspect_err(|_| {
+                match self.shared.raw.create_query_pool(&vk_info, None) {
+                    Ok(raw) => Some(raw),
+                    Err(e) => {
+                        // Creation failed: free the allocation back to the allocator before
+                        // destroying the acceleration structure and buffer, so the pooled
+                        // slot / dedicated memory is not leaked.
+                        self.mem_allocator.lock().free(&self.shared, allocation);
                         ray_tracing_functions
                             .acceleration_structure
                             .destroy_acceleration_structure(raw_acceleration_structure, None);
                         self.shared.raw.destroy_buffer(raw_buffer, None);
-                    })?;
-                Some(raw)
+                        return Err(super::map_host_device_oom_err(e));
+                    }
+                }
             } else {
                 None
             };