Fix memory leak (#1991)

Co-authored-by: comrademarc <40955683+comrademarc@users.noreply.github.com>
vulkano-rs · Sep 18, 2022 · 873fa78 · 873fa78
1 parent afe0b31
commit 873fa78
Show file tree

Hide file tree

Showing 7 changed files with 20 additions and 13 deletions.
diff --git a/vulkano/Cargo.toml b/vulkano/Cargo.toml
@@ -22,7 +22,6 @@ crossbeam-queue = "0.3"
 half = "1.8"
 libloading = "0.7"
 nalgebra = { version = "0.31.0", optional = true }
-once_cell = { version = "1.13", features = ["parking_lot"] }
 parking_lot = { version = "0.12", features = ["send_guard"] }
 smallvec = "1.8"
 

diff --git a/vulkano/src/buffer/cpu_access.rs b/vulkano/src/buffer/cpu_access.rs
@@ -235,7 +235,7 @@ where
         let mem_reqs = buffer.memory_requirements();
 
         let memory = MemoryPool::alloc_from_requirements(
-            device.standard_memory_pool(),
+            &device.standard_memory_pool(),
             &mem_reqs,
             AllocLayout::Linear,
             MappingRequirement::Map,

diff --git a/vulkano/src/buffer/device_local.rs b/vulkano/src/buffer/device_local.rs
@@ -359,7 +359,7 @@ where
         let (buffer, mem_reqs) = Self::build_buffer(&device, size, usage, &queue_family_indices)?;
 
         let memory = MemoryPool::alloc_from_requirements(
-            device.standard_memory_pool(),
+            &device.standard_memory_pool(),
             &mem_reqs,
             AllocLayout::Linear,
             MappingRequirement::DoNotMap,

diff --git a/vulkano/src/device/mod.rs b/vulkano/src/device/mod.rs
@@ -123,7 +123,6 @@ pub use crate::{
     fns::DeviceFunctions,
 };
 use ash::vk::Handle;
-use once_cell::sync::OnceCell;
 use parking_lot::{Mutex, MutexGuard};
 use smallvec::SmallVec;
 use std::{
@@ -137,7 +136,7 @@ use std::{
     mem::MaybeUninit,
     ops::Deref,
     ptr,
-    sync::Arc,
+    sync::{Arc, Weak},
 };
 
 pub(crate) mod extensions;
@@ -156,7 +155,7 @@ pub struct Device {
     api_version: Version,
 
     fns: DeviceFunctions,
-    standard_memory_pool: OnceCell<Arc<StandardMemoryPool>>,
+    standard_memory_pool: Mutex<Weak<StandardMemoryPool>>,
     enabled_extensions: DeviceExtensions,
     enabled_features: Features,
     active_queue_family_indices: SmallVec<[u32; 2]>,
@@ -416,7 +415,7 @@ impl Device {
             physical_device,
             api_version,
             fns,
-            standard_memory_pool: OnceCell::new(),
+            standard_memory_pool: Mutex::new(Weak::new()),
             enabled_extensions,
             enabled_features,
             active_queue_family_indices,
@@ -503,9 +502,18 @@ impl Device {
     }
 
     /// Returns the standard memory pool used by default if you don't provide any other pool.
-    pub fn standard_memory_pool<'a>(self: &'a Arc<Self>) -> &'a Arc<StandardMemoryPool> {
-        self.standard_memory_pool
-            .get_or_init(|| StandardMemoryPool::new(self.clone()))
+    pub fn standard_memory_pool(self: &Arc<Self>) -> Arc<StandardMemoryPool> {
+        let mut pool = self.standard_memory_pool.lock();
+
+        if let Some(p) = pool.upgrade() {
+            return p;
+        }
+
+        // The weak pointer is empty, so we create the pool.
+        let new_pool = StandardMemoryPool::new(self.clone());
+        *pool = Arc::downgrade(&new_pool);
+
+        new_pool
     }
 
     /// Gives you access to the standard descriptor pool that is used by default if you don't

diff --git a/vulkano/src/image/attachment.rs b/vulkano/src/image/attachment.rs
@@ -437,7 +437,7 @@ impl AttachmentImage {
 
         let mem_reqs = image.memory_requirements();
         let memory = MemoryPool::alloc_from_requirements(
-            device.standard_memory_pool(),
+            &device.standard_memory_pool(),
             &mem_reqs,
             AllocLayout::Optimal,
             MappingRequirement::DoNotMap,

diff --git a/vulkano/src/image/immutable.rs b/vulkano/src/image/immutable.rs
@@ -190,7 +190,7 @@ impl ImmutableImage {
 
         let mem_reqs = image.memory_requirements();
         let memory = MemoryPool::alloc_from_requirements(
-            device.standard_memory_pool(),
+            &device.standard_memory_pool(),
             &mem_reqs,
             AllocLayout::Optimal,
             MappingRequirement::DoNotMap,

diff --git a/vulkano/src/image/storage.rs b/vulkano/src/image/storage.rs
@@ -120,7 +120,7 @@ impl StorageImage {
 
         let mem_reqs = image.memory_requirements();
         let memory = MemoryPool::alloc_from_requirements(
-            device.standard_memory_pool(),
+            &device.standard_memory_pool(),
             &mem_reqs,
             AllocLayout::Optimal,
             MappingRequirement::DoNotMap,