Fix DestroyCommandPool command buffer leak

1 files changed, 38 insertions, 3 deletions

diff --git a/scripts/mock_icd_generator.py b/scripts/mock_icd_generator.py
index 310e81e9..5973c005 100644
--- a/scripts/mock_icd_generator.py
+++ b/scripts/mock_icd_generator.py
@@ -67,6 +67,7 @@ static unordered_map<VkDeviceMemory, VkDeviceSize> allocated_memory_size_map;
 static unordered_map<VkDevice, unordered_map<uint32_t, unordered_map<uint32_t, VkQueue>>> queue_map;
 static unordered_map<VkDevice, unordered_map<VkBuffer, VkBufferCreateInfo>> buffer_map;
 static unordered_map<VkDevice, unordered_map<VkImage, VkDeviceSize>> image_memory_size_map;
+static unordered_map<VkCommandPool, std::vector<VkCommandBuffer>> command_pool_buffer_map;
 
 static constexpr uint32_t icd_swapchain_image_count = 1;
 static unordered_map<VkSwapchainKHR, VkImage[icd_swapchain_image_count]> swapchain_image_map;
@@ -444,10 +445,42 @@ CUSTOM_C_INTERCEPTS = {
         DestroyDispObjHandle((void*)instance);
     }
 ''',
+'vkAllocateCommandBuffers': '''
+    unique_lock_t lock(global_lock);
+    for (uint32_t i = 0; i < pAllocateInfo->commandBufferCount; ++i) {
+        pCommandBuffers[i] = (VkCommandBuffer)CreateDispObjHandle();
+        command_pool_buffer_map[pAllocateInfo->commandPool].push_back(pCommandBuffers[i]);
+    }
+    return VK_SUCCESS;
+''',
 'vkFreeCommandBuffers': '''
-    for (auto i = 0u; i < commandBufferCount; ++i)
-        if (pCommandBuffers[i])
-            DestroyDispObjHandle((void*) pCommandBuffers[i]);
+    unique_lock_t lock(global_lock);
+    for (auto i = 0u; i < commandBufferCount; ++i) {
+        if (!pCommandBuffers[i]) {
+            continue;
+        }
+
+        for (auto& pair : command_pool_buffer_map) {
+            auto& cbs = pair.second;
+            auto it = std::find(cbs.begin(), cbs.end(), pCommandBuffers[i]);
+            if (it != cbs.end()) {
+                cbs.erase(it);
+            }
+        }
+            
+        DestroyDispObjHandle((void*) pCommandBuffers[i]);
+    }
+''',
+'vkDestroyCommandPool': '''
+    // destroy command buffers for this pool
+    unique_lock_t lock(global_lock);
+    auto it = command_pool_buffer_map.find(commandPool);
+    if (it != command_pool_buffer_map.end()) {
+        for (auto& cb : it->second) {
+            DestroyDispObjHandle((void*) cb);
+        }
+        command_pool_buffer_map.erase(it);
+    }
 ''',
 'vkEnumeratePhysicalDevices': '''
     VkResult result_code = VK_SUCCESS;
@@ -1366,6 +1399,8 @@ class MockICDOutputGenerator(OutputGenerator):
             'vkCreateInstance',
             'vkDestroyInstance',
             'vkFreeCommandBuffers',
+            'vkAllocateCommandBuffers',
+            'vkDestroyCommandPool',
             #'vkCreateDebugReportCallbackEXT',
             #'vkDestroyDebugReportCallbackEXT',
             'vkEnumerateInstanceLayerProperties',