mock: Fix memory leak issue

3 files changed, 27 insertions, 0 deletions

diff --git a/icd/generated/function_definitions.h b/icd/generated/function_definitions.h
index 5eb09156..d5a7a343 100644
--- a/icd/generated/function_definitions.h
+++ b/icd/generated/function_definitions.h
@@ -250,6 +250,14 @@ static VKAPI_ATTR void VKAPI_CALL DestroyDevice(
         }
     }
 
+    for (auto& cp : command_pool_map[device]) {
+        for (auto& cb : command_pool_buffer_map[cp]) {
+            DestroyDispObjHandle((void*) cb);
+        }
+        command_pool_buffer_map.erase(cp);
+    }
+    command_pool_map[device].clear();
+
     queue_map.erase(device);
     buffer_map.erase(device);
     image_memory_size_map.erase(device);
@@ -1095,6 +1103,7 @@ static VKAPI_ATTR VkResult VKAPI_CALL CreateCommandPool(
 {
     unique_lock_t lock(global_lock);
     *pCommandPool = (VkCommandPool)global_unique_handle++;
+    command_pool_map[device].insert(*pCommandPool);
     return VK_SUCCESS;
 }
 
@@ -1113,6 +1122,7 @@ static VKAPI_ATTR void VKAPI_CALL DestroyCommandPool(
         }
         command_pool_buffer_map.erase(it);
     }
+    command_pool_map[device].erase(commandPool);
 }
 
 static VKAPI_ATTR VkResult VKAPI_CALL ResetCommandPool(
diff --git a/icd/mock_icd.h b/icd/mock_icd.h
index dd712bcd..4210ed30 100644
--- a/icd/mock_icd.h
+++ b/icd/mock_icd.h
@@ -22,6 +22,7 @@
 #include <algorithm>
 #include <array>
 #include <mutex>
+#include <unordered_set>
 #include <unordered_map>
 #include <string>
 #include <vector>
@@ -64,6 +65,7 @@ struct BufferState {
 };
 static std::unordered_map<VkDevice, std::unordered_map<VkBuffer, BufferState>> buffer_map;
 static std::unordered_map<VkDevice, std::unordered_map<VkImage, VkDeviceSize>> image_memory_size_map;
+static std::unordered_map<VkDevice, std::unordered_set<VkCommandPool>> command_pool_map;
 static std::unordered_map<VkCommandPool, std::vector<VkCommandBuffer>> command_pool_buffer_map;
 
 static constexpr uint32_t icd_swapchain_image_count = 1;
diff --git a/scripts/mock_icd_generator.py b/scripts/mock_icd_generator.py
index 95166616..b09abb8c 100644
--- a/scripts/mock_icd_generator.py
+++ b/scripts/mock_icd_generator.py
@@ -77,6 +77,12 @@ CUSTOM_C_INTERCEPTS = {
         DestroyDispObjHandle((void*) pCommandBuffers[i]);
     }
 ''',
+'vkCreateCommandPool': '''
+    unique_lock_t lock(global_lock);
+    *pCommandPool = (VkCommandPool)global_unique_handle++;
+    command_pool_map[device].insert(*pCommandPool);
+    return VK_SUCCESS;
+''',
 'vkDestroyCommandPool': '''
     // destroy command buffers for this pool
     unique_lock_t lock(global_lock);
@@ -87,6 +93,7 @@ CUSTOM_C_INTERCEPTS = {
         }
         command_pool_buffer_map.erase(it);
     }
+    command_pool_map[device].erase(commandPool);
 ''',
 'vkEnumeratePhysicalDevices': '''
     VkResult result_code = VK_SUCCESS;
@@ -115,6 +122,14 @@ CUSTOM_C_INTERCEPTS = {
         }
     }
 
+    for (auto& cp : command_pool_map[device]) {
+        for (auto& cb : command_pool_buffer_map[cp]) {
+            DestroyDispObjHandle((void*) cb);
+        }
+        command_pool_buffer_map.erase(cp);
+    }
+    command_pool_map[device].clear();
+
     queue_map.erase(device);
     buffer_map.erase(device);
     image_memory_size_map.erase(device);