From 06fd7f2012a6087c4b1fc76842faad44cd097bde Mon Sep 17 00:00:00 2001
From: Liam <byteslice@airmail.cc>
Date: Sun, 28 Jan 2024 21:58:18 -0500
Subject: [PATCH] nvservices: unmap only on last container free

---
 src/core/hle/service/nvdrv/core/container.cpp | 7 ++++++-
 src/core/hle/service/nvdrv/core/container.h   | 1 +
 src/core/hle/service/nvdrv/core/nvmap.cpp     | 8 ++++++--
 3 files changed, 13 insertions(+), 3 deletions(-)

diff --git a/src/core/hle/service/nvdrv/core/container.cpp b/src/core/hle/service/nvdrv/core/container.cpp
index e89cca6f2..9edce03f6 100644
--- a/src/core/hle/service/nvdrv/core/container.cpp
+++ b/src/core/hle/service/nvdrv/core/container.cpp
@@ -49,6 +49,7 @@ SessionId Container::OpenSession(Kernel::KProcess* process) {
             continue;
         }
         if (session.process == process) {
+            session.ref_count++;
             return session.id;
         }
     }
@@ -66,6 +67,7 @@ SessionId Container::OpenSession(Kernel::KProcess* process) {
     }
     auto& session = impl->sessions[new_id];
     session.is_active = true;
+    session.ref_count = 1;
     // Optimization
     if (process->IsApplication()) {
         auto& page_table = process->GetPageTable().GetBasePageTable();
@@ -114,8 +116,11 @@ SessionId Container::OpenSession(Kernel::KProcess* process) {
 
 void Container::CloseSession(SessionId session_id) {
     std::scoped_lock lk(impl->session_guard);
-    impl->file.UnmapAllHandles(session_id);
     auto& session = impl->sessions[session_id.id];
+    if (--session.ref_count > 0) {
+        return;
+    }
+    impl->file.UnmapAllHandles(session_id);
     auto& smmu = impl->host1x.MemoryManager();
     if (session.has_preallocated_area) {
         const DAddr region_start = session.mapper->GetRegionStart();
diff --git a/src/core/hle/service/nvdrv/core/container.h b/src/core/hle/service/nvdrv/core/container.h
index b4d3938a8..f159ced09 100644
--- a/src/core/hle/service/nvdrv/core/container.h
+++ b/src/core/hle/service/nvdrv/core/container.h
@@ -46,6 +46,7 @@ struct Session {
     bool has_preallocated_area{};
     std::unique_ptr<HeapMapper> mapper{};
     bool is_active{};
+    s32 ref_count{};
 };
 
 class Container {
diff --git a/src/core/hle/service/nvdrv/core/nvmap.cpp b/src/core/hle/service/nvdrv/core/nvmap.cpp
index bc1c033c6..453cb5831 100644
--- a/src/core/hle/service/nvdrv/core/nvmap.cpp
+++ b/src/core/hle/service/nvdrv/core/nvmap.cpp
@@ -333,9 +333,13 @@ void NvMap::UnmapAllHandles(NvCore::SessionId session_id) {
     }();
 
     for (auto& [id, handle] : handles_copy) {
-        if (handle->session_id.id == session_id.id) {
-            FreeHandle(id, false);
+        {
+            std::scoped_lock lk{handle->mutex};
+            if (handle->session_id.id != session_id.id || handle->dupes <= 0) {
+                continue;
+            }
         }
+        FreeHandle(id, false);
     }
 }