From 70b4b981dff15d7535308846bbb9beb86dcc09b5 Mon Sep 17 00:00:00 2001
From: IndecisiveTurtle <47210458+raphaelthegreat@users.noreply.github.com>
Date: Thu, 3 Jul 2025 19:35:58 +0300
Subject: [PATCH] memory_tracker: Improve locking more on invalidation

---
 src/video_core/buffer_cache/buffer_cache.cpp |  8 ++---
 src/video_core/buffer_cache/memory_tracker.h | 34 ++++++++++++++------
 src/video_core/page_manager.cpp              |  3 +-
 3 files changed, 29 insertions(+), 16 deletions(-)

diff --git a/src/video_core/buffer_cache/buffer_cache.cpp b/src/video_core/buffer_cache/buffer_cache.cpp
index 289eca278..d55e05d1e 100644
--- a/src/video_core/buffer_cache/buffer_cache.cpp
+++ b/src/video_core/buffer_cache/buffer_cache.cpp
@@ -136,11 +136,9 @@ void BufferCache::InvalidateMemory(VAddr device_addr, u64 size) {
     if (!IsRegionRegistered(device_addr, size)) {
         return;
     }
-    if (Config::readbacks() && memory_tracker->IsRegionGpuModified(device_addr, size)) {
-        ReadMemory(device_addr, size, true);
-    } else {
-        memory_tracker->MarkRegionAsCpuModified(device_addr, size);
-    }
+    memory_tracker->InvalidateRegion(
+        device_addr, size, Config::readbacks(),
+        [this, device_addr, size] { ReadMemory(device_addr, size, true); });
 }
 
 void BufferCache::ReadMemory(VAddr device_addr, u64 size, bool is_write) {
diff --git a/src/video_core/buffer_cache/memory_tracker.h b/src/video_core/buffer_cache/memory_tracker.h
index 601bb4c4b..ca87c7df0 100644
--- a/src/video_core/buffer_cache/memory_tracker.h
+++ b/src/video_core/buffer_cache/memory_tracker.h
@@ -51,16 +51,6 @@ public:
                             });
     }
 
-    /// Mark region as modified from the host GPU
-    void MarkRegionAsGpuModified(VAddr dirty_cpu_addr, u64 query_size) noexcept {
-        IteratePages<false>(dirty_cpu_addr, query_size,
-                            [](RegionManager* manager, u64 offset, size_t size) {
-                                std::scoped_lock lk{manager->lock};
-                                manager->template ChangeRegionState<Type::GPU, true>(
-                                    manager->GetCpuAddr() + offset, size);
-                            });
-    }
-
     /// Unmark region as modified from the host GPU
     void UnmarkRegionAsGpuModified(VAddr dirty_cpu_addr, u64 query_size) noexcept {
         IteratePages<false>(dirty_cpu_addr, query_size,
@@ -71,6 +61,30 @@ public:
                             });
     }
 
+    /// Removes all protection from a page and ensures GPU data has been flushed if requested
+    void InvalidateRegion(VAddr cpu_addr, u64 size, bool try_flush, auto&& on_flush) noexcept {
+        IteratePages<false>(
+            cpu_addr, size,
+            [try_flush, &on_flush](RegionManager* manager, u64 offset, size_t size) {
+                const bool should_flush = [&] {
+                    // Perform both the GPU modification check and CPU state change with the lock
+                    // in case we are racing with GPU thread trying to mark the page as GPU
+                    // modified. If we need to flush the flush function is going to perform CPU
+                    // state change.
+                    std::scoped_lock lk{manager->lock};
+                    if (try_flush && manager->template IsRegionModified<Type::GPU>(offset, size)) {
+                        return true;
+                    }
+                    manager->template ChangeRegionState<Type::CPU, true>(
+                        manager->GetCpuAddr() + offset, size);
+                    return false;
+                }();
+                if (should_flush) {
+                    on_flush();
+                }
+            });
+    }
+
     /// Call 'func' for each CPU modified range and unmark those pages as CPU modified
     void ForEachUploadRange(VAddr query_cpu_range, u64 query_size, bool is_written, auto&& func) {
         IteratePages<true>(query_cpu_range, query_size,
diff --git a/src/video_core/page_manager.cpp b/src/video_core/page_manager.cpp
index 87434b043..63297bfdc 100644
--- a/src/video_core/page_manager.cpp
+++ b/src/video_core/page_manager.cpp
@@ -201,7 +201,8 @@ struct PageManager::Impl {
         RENDERER_TRACE;
         auto* memory = Core::Memory::Instance();
         auto& impl = memory->GetAddressSpace();
-        // ASSERT(perms != Core::MemoryPermission::Write);
+        ASSERT_MSG(perms != Core::MemoryPermission::Write,
+                   "Attempted to protect region as write-only which is not a valid permission");
         impl.Protect(address, size, perms);
     }