From 4a835bb2b9130b91d35968b2f7800084cf286de4 Mon Sep 17 00:00:00 2001
From: gdkchan <gab.dark.100@gmail.com>
Date: Mon, 25 Sep 2023 20:50:06 -0300
Subject: [PATCH] Make Vulkan memory allocator actually thread safe (#5575)

* Make Vulkan memory allocator actually thread safe

* Make free thread safe too

* PR feedback
---
 .../MemoryAllocator.cs                        | 33 ++++++++--
 .../MemoryAllocatorBlockList.cs               | 65 ++++++++++++++-----
 2 files changed, 74 insertions(+), 24 deletions(-)

diff --git a/src/Ryujinx.Graphics.Vulkan/MemoryAllocator.cs b/src/Ryujinx.Graphics.Vulkan/MemoryAllocator.cs
index ab70062703..aa0b410c40 100644
--- a/src/Ryujinx.Graphics.Vulkan/MemoryAllocator.cs
+++ b/src/Ryujinx.Graphics.Vulkan/MemoryAllocator.cs
@@ -1,6 +1,7 @@
 using Silk.NET.Vulkan;
 using System;
 using System.Collections.Generic;
+using System.Threading;
 
 namespace Ryujinx.Graphics.Vulkan
 {
@@ -13,6 +14,7 @@ namespace Ryujinx.Graphics.Vulkan
         private readonly Device _device;
         private readonly List<MemoryAllocatorBlockList> _blockLists;
         private readonly int _blockAlignment;
+        private readonly ReaderWriterLockSlim _lock;
 
         public MemoryAllocator(Vk api, VulkanPhysicalDevice physicalDevice, Device device)
         {
@@ -21,6 +23,7 @@ namespace Ryujinx.Graphics.Vulkan
             _device = device;
             _blockLists = new List<MemoryAllocatorBlockList>();
             _blockAlignment = (int)Math.Min(int.MaxValue, MaxDeviceMemoryUsageEstimate / _physicalDevice.PhysicalDeviceProperties.Limits.MaxMemoryAllocationCount);
+            _lock = new(LockRecursionPolicy.NoRecursion);
         }
 
         public MemoryAllocation AllocateDeviceMemory(
@@ -40,21 +43,37 @@ namespace Ryujinx.Graphics.Vulkan
 
         private MemoryAllocation Allocate(int memoryTypeIndex, ulong size, ulong alignment, bool map, bool isBuffer)
         {
-            for (int i = 0; i < _blockLists.Count; i++)
+            _lock.EnterReadLock();
+
+            try
             {
-                var bl = _blockLists[i];
-                if (bl.MemoryTypeIndex == memoryTypeIndex && bl.ForBuffer == isBuffer)
+                for (int i = 0; i < _blockLists.Count; i++)
                 {
-                    lock (bl)
+                    var bl = _blockLists[i];
+                    if (bl.MemoryTypeIndex == memoryTypeIndex && bl.ForBuffer == isBuffer)
                     {
                         return bl.Allocate(size, alignment, map);
                     }
                 }
             }
+            finally
+            {
+                _lock.ExitReadLock();
+            }
 
-            var newBl = new MemoryAllocatorBlockList(_api, _device, memoryTypeIndex, _blockAlignment, isBuffer);
-            _blockLists.Add(newBl);
-            return newBl.Allocate(size, alignment, map);
+            _lock.EnterWriteLock();
+
+            try
+            {
+                var newBl = new MemoryAllocatorBlockList(_api, _device, memoryTypeIndex, _blockAlignment, isBuffer);
+                _blockLists.Add(newBl);
+
+                return newBl.Allocate(size, alignment, map);
+            }
+            finally
+            {
+                _lock.ExitWriteLock();
+            }
         }
 
         internal int FindSuitableMemoryTypeIndex(
diff --git a/src/Ryujinx.Graphics.Vulkan/MemoryAllocatorBlockList.cs b/src/Ryujinx.Graphics.Vulkan/MemoryAllocatorBlockList.cs
index bd57e778b4..6a40b16e3f 100644
--- a/src/Ryujinx.Graphics.Vulkan/MemoryAllocatorBlockList.cs
+++ b/src/Ryujinx.Graphics.Vulkan/MemoryAllocatorBlockList.cs
@@ -3,6 +3,7 @@ using Silk.NET.Vulkan;
 using System;
 using System.Collections.Generic;
 using System.Diagnostics;
+using System.Threading;
 
 namespace Ryujinx.Graphics.Vulkan
 {
@@ -166,6 +167,8 @@ namespace Ryujinx.Graphics.Vulkan
 
         private readonly int _blockAlignment;
 
+        private readonly ReaderWriterLockSlim _lock;
+
         public MemoryAllocatorBlockList(Vk api, Device device, int memoryTypeIndex, int blockAlignment, bool forBuffer)
         {
             _blocks = new List<Block>();
@@ -174,6 +177,7 @@ namespace Ryujinx.Graphics.Vulkan
             MemoryTypeIndex = memoryTypeIndex;
             ForBuffer = forBuffer;
             _blockAlignment = blockAlignment;
+            _lock = new(LockRecursionPolicy.NoRecursion);
         }
 
         public unsafe MemoryAllocation Allocate(ulong size, ulong alignment, bool map)
@@ -184,19 +188,28 @@ namespace Ryujinx.Graphics.Vulkan
                 throw new ArgumentOutOfRangeException(nameof(alignment), $"Invalid alignment 0x{alignment:X}.");
             }
 
-            for (int i = 0; i < _blocks.Count; i++)
-            {
-                var block = _blocks[i];
+            _lock.EnterReadLock();
 
-                if (block.Mapped == map && block.Size >= size)
+            try
+            {
+                for (int i = 0; i < _blocks.Count; i++)
                 {
-                    ulong offset = block.Allocate(size, alignment);
-                    if (offset != InvalidOffset)
+                    var block = _blocks[i];
+
+                    if (block.Mapped == map && block.Size >= size)
                     {
-                        return new MemoryAllocation(this, block, block.Memory, GetHostPointer(block, offset), offset, size);
+                        ulong offset = block.Allocate(size, alignment);
+                        if (offset != InvalidOffset)
+                        {
+                            return new MemoryAllocation(this, block, block.Memory, GetHostPointer(block, offset), offset, size);
+                        }
                     }
                 }
             }
+            finally
+            {
+                _lock.ExitReadLock();
+            }
 
             ulong blockAlignedSize = BitUtils.AlignUp(size, (ulong)_blockAlignment);
 
@@ -244,14 +257,23 @@ namespace Ryujinx.Graphics.Vulkan
 
             if (block.IsTotallyFree())
             {
-                for (int i = 0; i < _blocks.Count; i++)
+                _lock.EnterWriteLock();
+
+                try
                 {
-                    if (_blocks[i] == block)
+                    for (int i = 0; i < _blocks.Count; i++)
                     {
-                        _blocks.RemoveAt(i);
-                        break;
+                        if (_blocks[i] == block)
+                        {
+                            _blocks.RemoveAt(i);
+                            break;
+                        }
                     }
                 }
+                finally
+                {
+                    _lock.ExitWriteLock();
+                }
 
                 block.Destroy(_api, _device);
             }
@@ -259,13 +281,22 @@ namespace Ryujinx.Graphics.Vulkan
 
         private void InsertBlock(Block block)
         {
-            int index = _blocks.BinarySearch(block);
-            if (index < 0)
-            {
-                index = ~index;
-            }
+            _lock.EnterWriteLock();
 
-            _blocks.Insert(index, block);
+            try
+            {
+                int index = _blocks.BinarySearch(block);
+                if (index < 0)
+                {
+                    index = ~index;
+                }
+
+                _blocks.Insert(index, block);
+            }
+            finally
+            {
+                _lock.ExitWriteLock();
+            }
         }
 
         public void Dispose()