From 288ce1b033faf1123b76688d9a02323484f21810 Mon Sep 17 00:00:00 2001
From: Yiwei Zhang <zzyiwei@chromium.org>
Date: Fri, 20 Aug 2021 21:50:14 +0000
Subject: [PATCH 7/8] venus: check descriptor allocations against pool resource

Only kick in when async_set_allocation is enabled.

Signed-off-by: Yiwei Zhang <zzyiwei@chromium.org>
Reviewed-by: Chia-I Wu <olvaffe@gmail.com>
Reviewed-by: Ryan Neph <ryanneph@google.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/12501>
---
 src/virtio/vulkan/vn_descriptor_set.c | 80 +++++++++++++++++++++++++++
 src/virtio/vulkan/vn_descriptor_set.h |  1 +
 2 files changed, 81 insertions(+)

diff --git a/src/virtio/vulkan/vn_descriptor_set.c b/src/virtio/vulkan/vn_descriptor_set.c
index dbf9a879996..b6a1da54a01 100644
--- a/src/virtio/vulkan/vn_descriptor_set.c
+++ b/src/virtio/vulkan/vn_descriptor_set.c
@@ -277,6 +277,73 @@ vn_DestroyDescriptorPool(VkDevice device,
    vk_free(alloc, pool);
 }
 
+static bool
+vn_descriptor_pool_alloc_descriptors(
+   struct vn_descriptor_pool *pool,
+   const struct vn_descriptor_set_layout *layout,
+   uint32_t last_binding_descriptor_count)
+{
+   struct vn_descriptor_pool_state recovery;
+
+   if (!pool->async_set_allocation)
+      return true;
+
+   if (pool->used.set_count == pool->max.set_count)
+      return false;
+
+   /* backup current pool state to recovery */
+   recovery = pool->used;
+
+   ++pool->used.set_count;
+
+   for (uint32_t i = 0; i <= layout->last_binding; i++) {
+      const VkDescriptorType type = layout->bindings[i].type;
+      const uint32_t count = i == layout->last_binding
+                                ? last_binding_descriptor_count
+                                : layout->bindings[i].count;
+
+      pool->used.descriptor_counts[type] += count;
+
+      if (pool->used.descriptor_counts[type] >
+          pool->max.descriptor_counts[type]) {
+         /* restore pool state before this allocation */
+         pool->used = recovery;
+         return false;
+      }
+   }
+
+   return true;
+}
+
+static void
+vn_descriptor_pool_free_descriptors(
+   struct vn_descriptor_pool *pool,
+   const struct vn_descriptor_set_layout *layout,
+   uint32_t last_binding_descriptor_count)
+{
+   if (!pool->async_set_allocation)
+      return;
+
+   for (uint32_t i = 0; i <= layout->last_binding; i++) {
+      const uint32_t count = i == layout->last_binding
+                                ? last_binding_descriptor_count
+                                : layout->bindings[i].count;
+
+      pool->used.descriptor_counts[layout->bindings[i].type] -= count;
+   }
+
+   --pool->used.set_count;
+}
+
+static void
+vn_descriptor_pool_reset_descriptors(struct vn_descriptor_pool *pool)
+{
+   if (!pool->async_set_allocation)
+      return;
+
+   memset(&pool->used, 0, sizeof(pool->used));
+}
+
 VkResult
 vn_ResetDescriptorPool(VkDevice device,
                        VkDescriptorPool descriptorPool,
@@ -298,6 +365,8 @@ vn_ResetDescriptorPool(VkDevice device,
       vk_free(alloc, set);
    }
 
+   vn_descriptor_pool_reset_descriptors(pool);
+
    return VK_SUCCESS;
 }
 
@@ -347,9 +416,18 @@ vn_AllocateDescriptorSets(VkDevice device,
          last_binding_descriptor_count = variable_info->pDescriptorCounts[i];
       }
 
+      if (!vn_descriptor_pool_alloc_descriptors(
+             pool, layout, last_binding_descriptor_count)) {
+         pDescriptorSets[i] = VK_NULL_HANDLE;
+         result = VK_ERROR_OUT_OF_POOL_MEMORY;
+         goto fail;
+      }
+
       set = vk_zalloc(alloc, sizeof(*set), VN_DEFAULT_ALIGN,
                       VK_SYSTEM_ALLOCATION_SCOPE_OBJECT);
       if (!set) {
+         vn_descriptor_pool_free_descriptors(pool, layout,
+                                             last_binding_descriptor_count);
          pDescriptorSets[i] = VK_NULL_HANDLE;
          result = VK_ERROR_OUT_OF_HOST_MEMORY;
          goto fail;
@@ -380,6 +458,8 @@ fail:
       if (!set)
          break;
 
+      vn_descriptor_pool_free_descriptors(pool, set->layout,
+                                          set->last_binding_descriptor_count);
       list_del(&set->head);
       vn_object_base_fini(&set->base);
       vk_free(alloc, set);
diff --git a/src/virtio/vulkan/vn_descriptor_set.h b/src/virtio/vulkan/vn_descriptor_set.h
index 9eb100755aa..5f6afc44fb3 100644
--- a/src/virtio/vulkan/vn_descriptor_set.h
+++ b/src/virtio/vulkan/vn_descriptor_set.h
@@ -50,6 +50,7 @@ struct vn_descriptor_pool {
    VkAllocationCallbacks allocator;
    bool async_set_allocation;
    struct vn_descriptor_pool_state max;
+   struct vn_descriptor_pool_state used;
 
    struct list_head descriptor_sets;
 };
-- 
2.33.0.259.gc128427fd7-goog