Fix leaks via owner class wrappers on some vulkan objects

6 files changed, 131 insertions, 65 deletions

diff --git a/src/device_context.cc b/src/device_context.cc
index 99a4979..cea0540 100644
--- a/src/device_context.cc
+++ b/src/device_context.cc
@@ -38,7 +38,6 @@ DeviceContext::Clock::Clock(const DeviceContext& context) : device(context) {
 DeviceContext::Clock::~Clock() {}
 
 DeviceContext::Clock::time_point_t DeviceContext::Clock::now() {
-
     auto ts = timespec{};
     if (clock_gettime(CLOCK_MONOTONIC, &ts)) {
         throw errno;
diff --git a/src/physical_device_context.cc b/src/physical_device_context.cc
index 66940f5..9c4ad8e 100644
--- a/src/physical_device_context.cc
+++ b/src/physical_device_context.cc
@@ -40,6 +40,7 @@ PhysicalDeviceContext::PhysicalDeviceContext(
         auto count = std::uint32_t{};
         THROW_NON_VKSUCCESS(vtable.EnumerateDeviceExtensionProperties(
             physical_device, nullptr, &count, nullptr));
+
         auto supported_extensions = std::vector<VkExtensionProperties>(count);
         THROW_NON_VKSUCCESS(vtable.EnumerateDeviceExtensionProperties(
             physical_device, nullptr, &count,
diff --git a/src/queue_context.cc b/src/queue_context.cc
index 69bcb13..29dcbfb 100644
--- a/src/queue_context.cc
+++ b/src/queue_context.cc
@@ -13,26 +13,34 @@
 
 namespace low_latency {
 
+QueueContext::CommandPoolOwner::CommandPoolOwner(
+    const QueueContext& queue_context)
+    : queue_context(queue_context) {
+
+    const auto& device_context = this->queue_context.device_context;
+
+    const auto cpci = VkCommandPoolCreateInfo{
+        .sType = VK_STRUCTURE_TYPE_COMMAND_POOL_CREATE_INFO,
+        .flags = VK_COMMAND_POOL_CREATE_TRANSIENT_BIT |
+                 VK_COMMAND_POOL_CREATE_RESET_COMMAND_BUFFER_BIT,
+        .queueFamilyIndex = queue_context.queue_family_index,
+    };
+
+    THROW_NON_VKSUCCESS(device_context.vtable.CreateCommandPool(
+        device_context.device, &cpci, nullptr, &this->command_pool));
+}
+
+QueueContext::CommandPoolOwner::~CommandPoolOwner() {
+    const auto& device_context = this->queue_context.device_context;
+    device_context.vtable.DestroyCommandPool(device_context.device,
+                                             this->command_pool, nullptr);
+}
+
 QueueContext::QueueContext(DeviceContext& device_context, const VkQueue& queue,
                            const std::uint32_t& queue_family_index)
     : device_context(device_context), queue(queue),
-      queue_family_index(queue_family_index) {
-
-    // Important we make the command pool before the timestamp pool, because
-    // it's a dependency.
-    this->command_pool = [&]() {
-        const auto cpci = VkCommandPoolCreateInfo{
-            .sType = VK_STRUCTURE_TYPE_COMMAND_POOL_CREATE_INFO,
-            .flags = VK_COMMAND_POOL_CREATE_TRANSIENT_BIT |
-                     VK_COMMAND_POOL_CREATE_RESET_COMMAND_BUFFER_BIT,
-            .queueFamilyIndex = queue_family_index,
-        };
-
-        auto command_pool = VkCommandPool{};
-        THROW_NON_VKSUCCESS(device_context.vtable.CreateCommandPool(
-            device_context.device, &cpci, nullptr, &command_pool));
-        return command_pool;
-    }();
+      queue_family_index(queue_family_index),
+      command_pool(std::make_unique<CommandPoolOwner>(*this)) {
 
     // Only construct a timestamp pool if we support it!
     if (device_context.physical_device.supports_required_extensions) {
@@ -44,10 +52,6 @@ QueueContext::~QueueContext() {
     this->in_flight_frames.clear();
     this->submissions.clear();
     this->timestamp_pool.reset();
-
-    const auto& vtable = this->device_context.vtable;
-    vtable.DestroyCommandPool(this->device_context.device, this->command_pool,
-                              nullptr);
 }
 
 void QueueContext::notify_submit(
diff --git a/src/queue_context.hh b/src/queue_context.hh
index 0e6441e..701fc0d 100644
--- a/src/queue_context.hh
+++ b/src/queue_context.hh
@@ -24,7 +24,7 @@ class QueueContext final : public Context {
     // The amount of queue submissions we allow tracked per queue before
     // we give up tracking them. For a queue that is presented to,
     // these submissions will be constantly moved to Frame structs so
-    // it's not an issue that we only track so many - unless it just 
+    // it's not an issue that we only track so many - unless it just
     // happens that an application makes an unexpectedly large
     // amount of vkQueueSubmit's per frame. For queues which don't
     // present, this limit stops them from growing limitlessly in memory
@@ -37,7 +37,23 @@ class QueueContext final : public Context {
     const VkQueue queue;
     const std::uint32_t queue_family_index;
 
-    VkCommandPool command_pool;
+    struct CommandPoolOwner final {
+      private:
+        const QueueContext& queue_context;
+        VkCommandPool command_pool;
+
+      public:
+        CommandPoolOwner(const QueueContext& queue_context);
+        CommandPoolOwner(const CommandPoolOwner&) = delete;
+        CommandPoolOwner(CommandPoolOwner&&) = delete;
+        CommandPoolOwner operator=(const CommandPoolOwner&) = delete;
+        CommandPoolOwner operator=(CommandPoolOwner&&) = delete;
+        ~CommandPoolOwner();
+
+      public:
+        operator const VkCommandPool&() const { return this->command_pool; }
+    };
+    const std::unique_ptr<CommandPoolOwner> command_pool;
 
     std::unique_ptr<TimestampPool> timestamp_pool;
 
diff --git a/src/timestamp_pool.cc b/src/timestamp_pool.cc
index e02a4fc..247d411 100644
--- a/src/timestamp_pool.cc
+++ b/src/timestamp_pool.cc
@@ -2,7 +2,6 @@
 #include "device_context.hh"
 #include "queue_context.hh"
 
-#include <chrono>
 #include <ranges>
 #include <span>
 #include <vulkan/utility/vk_dispatch_table.h>
@@ -10,39 +9,59 @@
 
 namespace low_latency {
 
-TimestampPool::QueryChunk::QueryChunk(const QueueContext& queue_context) {
-    const auto& device_context = queue_context.device_context;
-    const auto& vtable = device_context.vtable;
+TimestampPool::QueryChunk::QueryPoolOwner::QueryPoolOwner(
+    const QueueContext& queue_context)
+    : queue_context(queue_context) {
 
-    this->query_pool = [&]() {
-        const auto qpci = VkQueryPoolCreateInfo{
-            .sType = VK_STRUCTURE_TYPE_QUERY_POOL_CREATE_INFO,
-            .queryType = VK_QUERY_TYPE_TIMESTAMP,
-            .queryCount = QueryChunk::CHUNK_SIZE};
+    const auto& device_context = this->queue_context.device_context;
+    const auto qpci =
+        VkQueryPoolCreateInfo{.sType = VK_STRUCTURE_TYPE_QUERY_POOL_CREATE_INFO,
+                              .queryType = VK_QUERY_TYPE_TIMESTAMP,
+                              .queryCount = QueryChunk::CHUNK_SIZE};
 
-        auto qp = VkQueryPool{};
-        THROW_NON_VKSUCCESS(
-            vtable.CreateQueryPool(device_context.device, &qpci, nullptr, &qp));
-        return qp;
-    }();
+    THROW_NON_VKSUCCESS(device_context.vtable.CreateQueryPool(
+        device_context.device, &qpci, nullptr, &this->query_pool));
+}
+
+TimestampPool::QueryChunk::QueryPoolOwner::~QueryPoolOwner() {
+    const auto& device_context = this->queue_context.device_context;
+    device_context.vtable.DestroyQueryPool(device_context.device,
+                                           this->query_pool, nullptr);
+}
+
+TimestampPool::QueryChunk::QueryChunk(const QueueContext& queue_context)
+    : query_pool(std::make_unique<QueryPoolOwner>(queue_context)),
+      command_buffers(std::make_unique<CommandBuffersOwner>(queue_context)) {
 
     this->free_indices = []() {
         constexpr auto KEYS = std::views::iota(0u, QueryChunk::CHUNK_SIZE);
         return std::make_unique<free_indices_t>(std::from_range, KEYS);
     }();
+}
 
-    this->command_buffers = [&]() -> auto {
-        auto cbs = std::make_unique<std::vector<VkCommandBuffer>>(CHUNK_SIZE);
-        const auto cbai = VkCommandBufferAllocateInfo{
-            .sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_ALLOCATE_INFO,
-            .commandPool = queue_context.command_pool,
-            .level = VK_COMMAND_BUFFER_LEVEL_PRIMARY,
-            .commandBufferCount = static_cast<std::uint32_t>(std::size(*cbs)),
-        };
-        THROW_NON_VKSUCCESS(vtable.AllocateCommandBuffers(
-            device_context.device, &cbai, std::data(*cbs)));
-        return cbs;
-    }();
+TimestampPool::QueryChunk::CommandBuffersOwner::CommandBuffersOwner(
+    const QueueContext& queue_context)
+    : queue_context(queue_context), command_buffers(CHUNK_SIZE) {
+
+    const auto& device_context = queue_context.device_context;
+
+    const auto cbai = VkCommandBufferAllocateInfo{
+        .sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_ALLOCATE_INFO,
+        .commandPool = *queue_context.command_pool,
+        .level = VK_COMMAND_BUFFER_LEVEL_PRIMARY,
+        .commandBufferCount = CHUNK_SIZE,
+    };
+    THROW_NON_VKSUCCESS(device_context.vtable.AllocateCommandBuffers(
+        device_context.device, &cbai, std::data(this->command_buffers)));
+}
+
+TimestampPool::QueryChunk::CommandBuffersOwner::~CommandBuffersOwner() {
+    const auto& device_context = this->queue_context.device_context;
+
+    device_context.vtable.FreeCommandBuffers(
+        device_context.device, *this->queue_context.command_pool,
+        static_cast<std::uint32_t>(std::size(this->command_buffers)),
+        std::data(this->command_buffers));
 }
 
 TimestampPool::QueryChunk::~QueryChunk() {}
@@ -87,7 +106,7 @@ TimestampPool::Handle::Handle(const TimestampPool& timestamp_pool,
                               const std::shared_ptr<QueryChunk>& origin_chunk,
                               const std::uint64_t& query_index)
     : timestamp_pool(timestamp_pool), origin_chunk(origin_chunk),
-      query_pool(origin_chunk->query_pool), query_index(query_index),
+      query_pool(*origin_chunk->query_pool), query_index(query_index),
       command_buffer((*origin_chunk->command_buffers)[query_index]) {}
 
 TimestampPool::Handle::~Handle() {
@@ -187,16 +206,6 @@ DeviceContext::Clock::time_point_t TimestampPool::Handle::get_time_required() {
     return *time;
 }
 
-TimestampPool::~TimestampPool() {
-    const auto& device = this->queue_context.device_context.device;
-    const auto& vtable = this->queue_context.device_context.vtable;
-    for (const auto& query_chunk : this->query_chunks) {
-        vtable.FreeCommandBuffers(device, this->queue_context.command_pool,
-                                  static_cast<std::uint32_t>(
-                                      std::size(*query_chunk->command_buffers)),
-                                  std::data(*query_chunk->command_buffers));
-        vtable.DestroyQueryPool(device, query_chunk->query_pool, nullptr);
-    }
-}
+TimestampPool::~TimestampPool() {}
 
 } // namespace low_latency
\ No newline at end of file
diff --git a/src/timestamp_pool.hh b/src/timestamp_pool.hh
index 8e336d5..67b34de 100644
--- a/src/timestamp_pool.hh
+++ b/src/timestamp_pool.hh
@@ -33,13 +33,50 @@ class TimestampPool final {
     // We reuse these when they're released.
     struct QueryChunk final {
       private:
-        using free_indices_t = std::unordered_set<std::uint64_t>;
         static constexpr auto CHUNK_SIZE = 512u;
 
       public:
-        VkQueryPool query_pool;
+        struct QueryPoolOwner final {
+          private:
+            const QueueContext& queue_context;
+            VkQueryPool query_pool;
+
+          public:
+            QueryPoolOwner(const QueueContext& queue_context);
+            QueryPoolOwner(const QueryPoolOwner&) = delete;
+            QueryPoolOwner(QueryPoolOwner&&) = delete;
+            QueryPoolOwner operator=(const QueryPoolOwner&) = delete;
+            QueryPoolOwner operator=(QueryPoolOwner&&) = delete;
+            ~QueryPoolOwner();
+
+          public:
+            operator const VkQueryPool&() const { return this->query_pool; }
+        };
+        std::unique_ptr<QueryPoolOwner> query_pool;
+
+        using free_indices_t = std::unordered_set<std::uint64_t>;
         std::unique_ptr<free_indices_t> free_indices;
-        std::unique_ptr<std::vector<VkCommandBuffer>> command_buffers;
+
+        struct CommandBuffersOwner final {
+          private:
+            const QueueContext& queue_context;
+            std::vector<VkCommandBuffer> command_buffers;
+
+          public:
+            CommandBuffersOwner(const QueueContext& queue_context);
+            CommandBuffersOwner(const CommandBuffersOwner&) = delete;
+            CommandBuffersOwner(CommandBuffersOwner&&) = delete;
+            CommandBuffersOwner operator=(const CommandBuffersOwner&) = delete;
+            CommandBuffersOwner operator=(CommandBuffersOwner&&) = delete;
+            ~CommandBuffersOwner();
+
+          public:
+            VkCommandBuffer operator[](const std::size_t& i) {
+                assert(i < CHUNK_SIZE);
+                return this->command_buffers[i];
+            }
+        };
+        std::unique_ptr<CommandBuffersOwner> command_buffers;
 
       public:
         QueryChunk(const QueueContext& queue_context);