From e249898ed9834b88f54202fcbf1c0cbe0835f76d Mon Sep 17 00:00:00 2001
From: Kyle Charbonneau <kylechar@google.com>
Date: Thu, 31 Oct 2024 17:48:45 +0000
Subject: [PATCH] Give DescriptorSetAllocator a mutex

There are memory corruption crashes inside DescriptorSetAllocator with
Android Graphite/Dawn/Vulkan. The class is accessed with MutexProtected
from BindGroupLayoutVk but not from DeviceVk so concurrent access to
member variables on different threads seems possible. Since the class is
reference counted move the mutex inside the class so that calling a
method on any reference is thread safe.

Bug: 368362677
Change-Id: I7ca400555261c75995431d4a4023a571747961ee
Reviewed-on: https://dawn-review.googlesource.com/c/dawn/+/212774
Reviewed-by: Corentin Wallez <cwallez@chromium.org>
Reviewed-by: Geoff Lang <geofflang@chromium.org>
Commit-Queue: Kyle Charbonneau <kylechar@google.com>
---
 src/dawn/native/vulkan/BindGroupLayoutVk.h        | 2 +-
 src/dawn/native/vulkan/DescriptorSetAllocator.cpp | 6 ++++++
 src/dawn/native/vulkan/DescriptorSetAllocator.h   | 4 ++++
 3 files changed, 11 insertions(+), 1 deletion(-)

diff --git a/src/dawn/native/vulkan/BindGroupLayoutVk.h b/src/dawn/native/vulkan/BindGroupLayoutVk.h
index 0ac799c73a1..c1624e7af8e 100644
--- a/src/dawn/native/vulkan/BindGroupLayoutVk.h
+++ b/src/dawn/native/vulkan/BindGroupLayoutVk.h
@@ -94,7 +94,7 @@ class BindGroupLayout final : public BindGroupLayoutInternalBase {
     VkDescriptorSetLayout mHandle = VK_NULL_HANDLE;
 
     MutexProtected<SlabAllocator<BindGroup>> mBindGroupAllocator;
-    MutexProtected<Ref<DescriptorSetAllocator>> mDescriptorSetAllocator;
+    Ref<DescriptorSetAllocator> mDescriptorSetAllocator;
 };
 
 }  // namespace dawn::native::vulkan
diff --git a/src/dawn/native/vulkan/DescriptorSetAllocator.cpp b/src/dawn/native/vulkan/DescriptorSetAllocator.cpp
index e74adf3f0d4..123718d7ca3 100644
--- a/src/dawn/native/vulkan/DescriptorSetAllocator.cpp
+++ b/src/dawn/native/vulkan/DescriptorSetAllocator.cpp
@@ -94,6 +94,8 @@ DescriptorSetAllocator::~DescriptorSetAllocator() {
 }
 
 ResultOrError<DescriptorSetAllocation> DescriptorSetAllocator::Allocate(BindGroupLayout* layout) {
+    Mutex::AutoLock lock(&mMutex);
+
     if (mAvailableDescriptorPoolIndices.empty()) {
         DAWN_TRY(AllocateDescriptorPool(layout));
     }
@@ -116,6 +118,8 @@ ResultOrError<DescriptorSetAllocation> DescriptorSetAllocator::Allocate(BindGrou
 }
 
 void DescriptorSetAllocator::Deallocate(DescriptorSetAllocation* allocationInfo) {
+    Mutex::AutoLock lock(&mMutex);
+
     DAWN_ASSERT(allocationInfo != nullptr);
     DAWN_ASSERT(allocationInfo->set != VK_NULL_HANDLE);
 
@@ -136,6 +140,8 @@ void DescriptorSetAllocator::Deallocate(DescriptorSetAllocation* allocationInfo)
 }
 
 void DescriptorSetAllocator::FinishDeallocation(ExecutionSerial completedSerial) {
+    Mutex::AutoLock lock(&mMutex);
+
     for (const Deallocation& dealloc : mPendingDeallocations.IterateUpTo(completedSerial)) {
         DAWN_ASSERT(dealloc.poolIndex < mDescriptorPools.size());
 
diff --git a/src/dawn/native/vulkan/DescriptorSetAllocator.h b/src/dawn/native/vulkan/DescriptorSetAllocator.h
index 4da275a3dcc..c1b9cd805cd 100644
--- a/src/dawn/native/vulkan/DescriptorSetAllocator.h
+++ b/src/dawn/native/vulkan/DescriptorSetAllocator.h
@@ -31,6 +31,7 @@
 #include <vector>
 
 #include "absl/container/flat_hash_map.h"
+#include "dawn/common/Mutex.h"
 #include "dawn/common/SerialQueue.h"
 #include "dawn/common/vulkan_platform.h"
 #include "dawn/native/Error.h"
@@ -81,6 +82,9 @@ class DescriptorSetAllocator : public ObjectBase {
     };
     SerialQueue<ExecutionSerial, Deallocation> mPendingDeallocations;
     ExecutionSerial mLastDeallocationSerial = ExecutionSerial(0);
+
+    // Used to guard all public member functions.
+    Mutex mMutex;
 };
 
 }  // namespace dawn::native::vulkan