Vulkan: Fix dangling pointer in DescriptorSetAllocator

Bug: dawn:2349
Change-Id: I1cedde78a7768158d1b2f05721e2c1063e47766e
Reviewed-on: https://dawn-review.googlesource.com/c/dawn/+/177926
Kokoro: Kokoro <noreply+kokoro@google.com>
Commit-Queue: Jiawei Shao <jiawei.shao@intel.com>
Reviewed-by: Corentin Wallez <cwallez@chromium.org>
diff --git a/src/dawn/native/vulkan/BindGroupLayoutVk.cpp b/src/dawn/native/vulkan/BindGroupLayoutVk.cpp
index 3f78a37..70e90bc 100644
--- a/src/dawn/native/vulkan/BindGroupLayoutVk.cpp
+++ b/src/dawn/native/vulkan/BindGroupLayoutVk.cpp
@@ -146,7 +146,7 @@
     // TODO(enga): Consider deduping allocators for layouts with the same descriptor type
     // counts.
     mDescriptorSetAllocator =
-        DescriptorSetAllocator::Create(this, std::move(descriptorCountPerType));
+        DescriptorSetAllocator::Create(device, std::move(descriptorCountPerType));
 
     SetLabelImpl();
 
@@ -185,7 +185,7 @@
     Device* device,
     const BindGroupDescriptor* descriptor) {
     DescriptorSetAllocation descriptorSetAllocation;
-    DAWN_TRY_ASSIGN(descriptorSetAllocation, mDescriptorSetAllocator->Allocate());
+    DAWN_TRY_ASSIGN(descriptorSetAllocation, mDescriptorSetAllocator->Allocate(this));
 
     return AcquireRef(mBindGroupAllocator->Allocate(device, descriptor, descriptorSetAllocation));
 }
diff --git a/src/dawn/native/vulkan/DescriptorSetAllocator.cpp b/src/dawn/native/vulkan/DescriptorSetAllocator.cpp
index ff3d6e8..e74adf3 100644
--- a/src/dawn/native/vulkan/DescriptorSetAllocator.cpp
+++ b/src/dawn/native/vulkan/DescriptorSetAllocator.cpp
@@ -42,17 +42,15 @@
 
 // static
 Ref<DescriptorSetAllocator> DescriptorSetAllocator::Create(
-    BindGroupLayout* layout,
+    DeviceBase* device,
     absl::flat_hash_map<VkDescriptorType, uint32_t> descriptorCountPerType) {
-    return AcquireRef(new DescriptorSetAllocator(layout, descriptorCountPerType));
+    return AcquireRef(new DescriptorSetAllocator(device, descriptorCountPerType));
 }
 
 DescriptorSetAllocator::DescriptorSetAllocator(
-    BindGroupLayout* layout,
+    DeviceBase* device,
     absl::flat_hash_map<VkDescriptorType, uint32_t> descriptorCountPerType)
-    : ObjectBase(layout->GetDevice()), mLayout(layout) {
-    DAWN_ASSERT(layout != nullptr);
-
+    : ObjectBase(device) {
     // Compute the total number of descriptors for this layout.
     uint32_t totalDescriptorCount = 0;
     mPoolSizes.reserve(descriptorCountPerType.size());
@@ -95,9 +93,9 @@
     }
 }
 
-ResultOrError<DescriptorSetAllocation> DescriptorSetAllocator::Allocate() {
+ResultOrError<DescriptorSetAllocation> DescriptorSetAllocator::Allocate(BindGroupLayout* layout) {
     if (mAvailableDescriptorPoolIndices.empty()) {
-        DAWN_TRY(AllocateDescriptorPool());
+        DAWN_TRY(AllocateDescriptorPool(layout));
     }
 
     DAWN_ASSERT(!mAvailableDescriptorPoolIndices.empty());
@@ -150,7 +148,7 @@
     mPendingDeallocations.ClearUpTo(completedSerial);
 }
 
-MaybeError DescriptorSetAllocator::AllocateDescriptorPool() {
+MaybeError DescriptorSetAllocator::AllocateDescriptorPool(BindGroupLayout* layout) {
     VkDescriptorPoolCreateInfo createInfo;
     createInfo.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_POOL_CREATE_INFO;
     createInfo.pNext = nullptr;
@@ -166,7 +164,7 @@
                                                             nullptr, &*descriptorPool),
                             "CreateDescriptorPool"));
 
-    std::vector<VkDescriptorSetLayout> layouts(mMaxSets, mLayout->GetHandle());
+    std::vector<VkDescriptorSetLayout> layouts(mMaxSets, layout->GetHandle());
 
     VkDescriptorSetAllocateInfo allocateInfo;
     allocateInfo.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_ALLOCATE_INFO;
diff --git a/src/dawn/native/vulkan/DescriptorSetAllocator.h b/src/dawn/native/vulkan/DescriptorSetAllocator.h
index c5aaf4e..4da275a 100644
--- a/src/dawn/native/vulkan/DescriptorSetAllocator.h
+++ b/src/dawn/native/vulkan/DescriptorSetAllocator.h
@@ -49,22 +49,19 @@
 
   public:
     static Ref<DescriptorSetAllocator> Create(
-        BindGroupLayout* layout,
+        DeviceBase* device,
         absl::flat_hash_map<VkDescriptorType, uint32_t> descriptorCountPerType);
 
-    ResultOrError<DescriptorSetAllocation> Allocate();
+    ResultOrError<DescriptorSetAllocation> Allocate(BindGroupLayout* layout);
     void Deallocate(DescriptorSetAllocation* allocationInfo);
     void FinishDeallocation(ExecutionSerial completedSerial);
 
   private:
-    DescriptorSetAllocator(BindGroupLayout* layout,
+    DescriptorSetAllocator(DeviceBase* device,
                            absl::flat_hash_map<VkDescriptorType, uint32_t> descriptorCountPerType);
     ~DescriptorSetAllocator() override;
 
-    MaybeError AllocateDescriptorPool();
-
-    // TODO(https://crbug.com/dawn/2349): Investigate DanglingUntriaged in dawn/native.
-    raw_ptr<const BindGroupLayout, DanglingUntriaged> mLayout;
+    MaybeError AllocateDescriptorPool(BindGroupLayout* layout);
 
     std::vector<VkDescriptorPoolSize> mPoolSizes;
     SetIndex mMaxSets;