mock: Fix memory leak issue

1 files changed, 15 insertions, 0 deletions

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);