6 files changed, 152 insertions, 26 deletions

diff --git a/src/core/CMakeLists.txt b/src/core/CMakeLists.txt
index 52151e788..1662ec63d 100644
--- a/src/core/CMakeLists.txt
+++ b/src/core/CMakeLists.txt
@@ -174,6 +174,7 @@ add_library(core STATIC
     hle/kernel/k_scheduler.h
     hle/kernel/k_scheduler_lock.h
     hle/kernel/k_scoped_lock.h
+    hle/kernel/k_scoped_resource_reservation.h
     hle/kernel/k_scoped_scheduler_lock_and_sleep.h
     hle/kernel/k_synchronization_object.cpp
     hle/kernel/k_synchronization_object.h
diff --git a/src/core/hle/kernel/k_scoped_resource_reservation.h b/src/core/hle/kernel/k_scoped_resource_reservation.h
new file mode 100644
index 000000000..c5deca00b
--- /dev/null
+++ b/src/core/hle/kernel/k_scoped_resource_reservation.h
@@ -0,0 +1,67 @@
+// Copyright 2021 yuzu Emulator Project
+// Licensed under GPLv2 or any later version
+// Refer to the license.txt file included.
+
+// This file references various implementation details from Atmosphere, an open-source firmware for
+// the Nintendo Switch. Copyright 2018-2020 Atmosphere-NX.
+
+#pragma once
+
+#include "common/common_types.h"
+#include "core/hle/kernel/k_resource_limit.h"
+#include "core/hle/kernel/process.h"
+
+namespace Kernel {
+
+class KScopedResourceReservation {
+public:
+    explicit KScopedResourceReservation(std::shared_ptr<KResourceLimit> l, LimitableResource r,
+                                        s64 v, s64 timeout)
+        : resource_limit(std::move(l)), value(v), resource(r) {
+        if (resource_limit && value) {
+            success = resource_limit->Reserve(resource, value, timeout);
+        } else {
+            success = true;
+        }
+    }
+
+    explicit KScopedResourceReservation(std::shared_ptr<KResourceLimit> l, LimitableResource r,
+                                        s64 v = 1)
+        : resource_limit(std::move(l)), value(v), resource(r) {
+        if (resource_limit && value) {
+            success = resource_limit->Reserve(resource, value);
+        } else {
+            success = true;
+        }
+    }
+
+    explicit KScopedResourceReservation(const Process* p, LimitableResource r, s64 v, s64 t)
+        : KScopedResourceReservation(p->GetResourceLimit(), r, v, t) {}
+
+    explicit KScopedResourceReservation(const Process* p, LimitableResource r, s64 v = 1)
+        : KScopedResourceReservation(p->GetResourceLimit(), r, v) {}
+
+    ~KScopedResourceReservation() noexcept {
+        if (resource_limit && value && success) {
+            // resource was not committed, release the reservation.
+            resource_limit->Release(resource, value);
+        }
+    }
+
+    /// Commit the resource reservation, destruction of this object does not release the resource
+    void Commit() {
+        resource_limit = nullptr;
+    }
+
+    [[nodiscard]] bool Succeeded() const {
+        return success;
+    }
+
+private:
+    std::shared_ptr<KResourceLimit> resource_limit;
+    s64 value;
+    LimitableResource resource;
+    bool success;
+};
+
+} // namespace Kernel
diff --git a/src/core/hle/kernel/memory/page_table.cpp b/src/core/hle/kernel/memory/page_table.cpp
index 5947fc748..00ed9b881 100644
--- a/src/core/hle/kernel/memory/page_table.cpp
+++ b/src/core/hle/kernel/memory/page_table.cpp
@@ -6,7 +6,7 @@
 #include "common/assert.h"
 #include "common/scope_exit.h"
 #include "core/core.h"
-#include "core/hle/kernel/k_resource_limit.h"
+#include "core/hle/kernel/k_scoped_resource_reservation.h"
 #include "core/hle/kernel/kernel.h"
 #include "core/hle/kernel/memory/address_space_info.h"
 #include "core/hle/kernel/memory/memory_block.h"
@@ -409,27 +409,25 @@ ResultCode PageTable::MapPhysicalMemory(VAddr addr, std::size_t size) {
         return RESULT_SUCCESS;
     }
 
-    auto process{system.Kernel().CurrentProcess()};
     const std::size_t remaining_size{size - mapped_size};
     const std::size_t remaining_pages{remaining_size / PageSize};
 
-    if (process->GetResourceLimit() &&
-        !process->GetResourceLimit()->Reserve(LimitableResource::PhysicalMemory, remaining_size)) {
+    // Reserve the memory from the process resource limit.
+    KScopedResourceReservation memory_reservation(
+        system.Kernel().CurrentProcess()->GetResourceLimit(), LimitableResource::PhysicalMemory,
+        remaining_size);
+    if (!memory_reservation.Succeeded()) {
+        LOG_ERROR(Kernel, "Could not reserve remaining {:X} bytes", remaining_size);
         return ResultResourceLimitedExceeded;
     }
 
     PageLinkedList page_linked_list;
-    {
-        auto block_guard = detail::ScopeExit([&] {
-            system.Kernel().MemoryManager().Free(page_linked_list, remaining_pages, memory_pool);
-            process->GetResourceLimit()->Release(LimitableResource::PhysicalMemory, remaining_size);
-        });
 
-        CASCADE_CODE(system.Kernel().MemoryManager().Allocate(page_linked_list, remaining_pages,
-                                                              memory_pool));
+    CASCADE_CODE(
+        system.Kernel().MemoryManager().Allocate(page_linked_list, remaining_pages, memory_pool));
 
-        block_guard.Cancel();
-    }
+    // We succeeded, so commit the memory reservation.
+    memory_reservation.Commit();
 
     MapPhysicalMemory(page_linked_list, addr, end_addr);
 
@@ -781,9 +779,13 @@ ResultVal<VAddr> PageTable::SetHeapSize(std::size_t size) {
 
         const u64 delta{size - previous_heap_size};
 
-        auto process{system.Kernel().CurrentProcess()};
-        if (process->GetResourceLimit() && delta != 0 &&
-            !process->GetResourceLimit()->Reserve(LimitableResource::PhysicalMemory, delta)) {
+        // Reserve memory for the heap extension.
+        KScopedResourceReservation memory_reservation(
+            system.Kernel().CurrentProcess()->GetResourceLimit(), LimitableResource::PhysicalMemory,
+            delta);
+
+        if (!memory_reservation.Succeeded()) {
+            LOG_ERROR(Kernel, "Could not reserve heap extension of size {:X} bytes", delta);
             return ResultResourceLimitedExceeded;
         }
 
@@ -800,6 +802,9 @@ ResultVal<VAddr> PageTable::SetHeapSize(std::size_t size) {
         CASCADE_CODE(
             Operate(current_heap_addr, num_pages, page_linked_list, OperationType::MapGroup));
 
+        // Succeeded in allocation, commit the resource reservation
+        memory_reservation.Commit();
+
         block_manager->Update(current_heap_addr, num_pages, MemoryState::Normal,
                               MemoryPermission::ReadAndWrite);
 
diff --git a/src/core/hle/kernel/process.cpp b/src/core/hle/kernel/process.cpp
index 39dc3898a..05e21830c 100644
--- a/src/core/hle/kernel/process.cpp
+++ b/src/core/hle/kernel/process.cpp
@@ -16,6 +16,7 @@
 #include "core/hle/kernel/code_set.h"
 #include "core/hle/kernel/k_resource_limit.h"
 #include "core/hle/kernel/k_scheduler.h"
+#include "core/hle/kernel/k_scoped_resource_reservation.h"
 #include "core/hle/kernel/k_thread.h"
 #include "core/hle/kernel/kernel.h"
 #include "core/hle/kernel/memory/memory_block_manager.h"
@@ -116,6 +117,9 @@ std::shared_ptr<Process> Process::Create(Core::System& system, std::string name,
 
     std::shared_ptr<Process> process = std::make_shared<Process>(system);
     process->name = std::move(name);
+
+    // TODO: This is inaccurate
+    // The process should hold a reference to the kernel-wide resource limit.
     process->resource_limit = std::make_shared<KResourceLimit>(kernel, system);
     process->status = ProcessStatus::Created;
     process->program_id = 0;
@@ -154,6 +158,9 @@ void Process::DecrementThreadCount() {
 }
 
 u64 Process::GetTotalPhysicalMemoryAvailable() const {
+    // TODO: This is expected to always return the application memory pool size after accurately
+    // reserving kernel resources. The current workaround uses a process-local resource limit of
+    // application memory pool size, which is inaccurate.
     const u64 capacity{resource_limit->GetFreeValue(LimitableResource::PhysicalMemory) +
                        page_table->GetTotalHeapSize() + GetSystemResourceSize() + image_size +
                        main_thread_stack_size};
@@ -263,6 +270,17 @@ ResultCode Process::LoadFromMetadata(const FileSys::ProgramMetadata& metadata,
     system_resource_size = metadata.GetSystemResourceSize();
     image_size = code_size;
 
+    // Set initial resource limits
+    resource_limit->SetLimitValue(
+        LimitableResource::PhysicalMemory,
+        kernel.MemoryManager().GetSize(Memory::MemoryManager::Pool::Application));
+    KScopedResourceReservation memory_reservation(resource_limit, LimitableResource::PhysicalMemory,
+                                                  code_size + system_resource_size);
+    if (!memory_reservation.Succeeded()) {
+        LOG_ERROR(Kernel, "Could not reserve process memory requirements of size {:X} bytes",
+                  code_size + system_resource_size);
+        return ERR_RESOURCE_LIMIT_EXCEEDED;
+    }
     // Initialize proces address space
     if (const ResultCode result{
             page_table->InitializeForProcess(metadata.GetAddressSpaceType(), false, 0x8000000,
@@ -304,24 +322,22 @@ ResultCode Process::LoadFromMetadata(const FileSys::ProgramMetadata& metadata,
         UNREACHABLE();
     }
 
-    // Set initial resource limits
-    resource_limit->SetLimitValue(
-        LimitableResource::PhysicalMemory,
-        kernel.MemoryManager().GetSize(Memory::MemoryManager::Pool::Application));
     resource_limit->SetLimitValue(LimitableResource::Threads, 608);
     resource_limit->SetLimitValue(LimitableResource::Events, 700);
     resource_limit->SetLimitValue(LimitableResource::TransferMemory, 128);
     resource_limit->SetLimitValue(LimitableResource::Sessions, 894);
-    ASSERT(resource_limit->Reserve(LimitableResource::PhysicalMemory, code_size));
 
     // Create TLS region
     tls_region_address = CreateTLSRegion();
+    memory_reservation.Commit();
 
     return handle_table.SetSize(capabilities.GetHandleTableSize());
 }
 
 void Process::Run(s32 main_thread_priority, u64 stack_size) {
     AllocateMainThreadStack(stack_size);
+    resource_limit->Reserve(LimitableResource::Threads, 1);
+    resource_limit->Reserve(LimitableResource::PhysicalMemory, main_thread_stack_size);
 
     const std::size_t heap_capacity{memory_usage_capacity - main_thread_stack_size - image_size};
     ASSERT(!page_table->SetHeapCapacity(heap_capacity).IsError());
@@ -329,8 +345,6 @@ void Process::Run(s32 main_thread_priority, u64 stack_size) {
     ChangeStatus(ProcessStatus::Running);
 
     SetupMainThread(system, *this, main_thread_priority, main_thread_stack_top);
-    resource_limit->Reserve(LimitableResource::Threads, 1);
-    resource_limit->Reserve(LimitableResource::PhysicalMemory, main_thread_stack_size);
 }
 
 void Process::PrepareForTermination() {
@@ -357,6 +371,11 @@ void Process::PrepareForTermination() {
     FreeTLSRegion(tls_region_address);
     tls_region_address = 0;
 
+    if (resource_limit) {
+        resource_limit->Release(LimitableResource::PhysicalMemory,
+                                main_thread_stack_size + image_size);
+    }
+
     ChangeStatus(ProcessStatus::Exited);
 }
 
diff --git a/src/core/hle/kernel/shared_memory.cpp b/src/core/hle/kernel/shared_memory.cpp
index 0cd467110..67d748561 100644
--- a/src/core/hle/kernel/shared_memory.cpp
+++ b/src/core/hle/kernel/shared_memory.cpp
@@ -4,6 +4,7 @@
 
 #include "common/assert.h"
 #include "core/core.h"
+#include "core/hle/kernel/k_scoped_resource_reservation.h"
 #include "core/hle/kernel/kernel.h"
 #include "core/hle/kernel/memory/page_table.h"
 #include "core/hle/kernel/shared_memory.h"
@@ -21,6 +22,11 @@ std::shared_ptr<SharedMemory> SharedMemory::Create(
     Memory::MemoryPermission user_permission, PAddr physical_address, std::size_t size,
     std::string name) {
 
+    const auto resource_limit = kernel.GetSystemResourceLimit();
+    KScopedResourceReservation memory_reservation(resource_limit, LimitableResource::PhysicalMemory,
+                                                  size);
+    ASSERT(memory_reservation.Succeeded());
+
     std::shared_ptr<SharedMemory> shared_memory{
         std::make_shared<SharedMemory>(kernel, device_memory)};
 
@@ -32,6 +38,7 @@ std::shared_ptr<SharedMemory> SharedMemory::Create(
     shared_memory->size = size;
     shared_memory->name = name;
 
+    memory_reservation.Commit();
     return shared_memory;
 }
 
diff --git a/src/core/hle/kernel/svc.cpp b/src/core/hle/kernel/svc.cpp
index 4ef3c7ac5..1d377ffe6 100644
--- a/src/core/hle/kernel/svc.cpp
+++ b/src/core/hle/kernel/svc.cpp
@@ -30,6 +30,7 @@
 #include "core/hle/kernel/k_readable_event.h"
 #include "core/hle/kernel/k_resource_limit.h"
 #include "core/hle/kernel/k_scheduler.h"
+#include "core/hle/kernel/k_scoped_resource_reservation.h"
 #include "core/hle/kernel/k_scoped_scheduler_lock_and_sleep.h"
 #include "core/hle/kernel/k_synchronization_object.h"
 #include "core/hle/kernel/k_thread.h"
@@ -1516,8 +1517,13 @@ static ResultCode CreateThread(Core::System& system, Handle* out_handle, VAddr e
         return ResultInvalidPriority;
     }
 
-    ASSERT(process.GetResourceLimit()->Reserve(
-        LimitableResource::Threads, 1, system.CoreTiming().GetGlobalTimeNs().count() + 100000000));
+    KScopedResourceReservation thread_reservation(
+        kernel.CurrentProcess(), LimitableResource::Threads, 1,
+        system.CoreTiming().GetGlobalTimeNs().count() + 100000000);
+    if (!thread_reservation.Succeeded()) {
+        LOG_ERROR(Kernel_SVC, "Could not reserve a new thread");
+        return ERR_RESOURCE_LIMIT_EXCEEDED;
+    }
 
     std::shared_ptr<KThread> thread;
     {
@@ -1537,6 +1543,7 @@ static ResultCode CreateThread(Core::System& system, Handle* out_handle, VAddr e
     // Set the thread name for debugging purposes.
     thread->SetName(
         fmt::format("thread[entry_point={:X}, handle={:X}]", entry_point, *new_thread_handle));
+    thread_reservation.Commit();
 
     return RESULT_SUCCESS;
 }
@@ -1884,6 +1891,13 @@ static ResultCode CreateTransferMemory(Core::System& system, Handle* handle, VAd
     }
 
     auto& kernel = system.Kernel();
+    // Reserve a new transfer memory from the process resource limit.
+    KScopedResourceReservation trmem_reservation(kernel.CurrentProcess(),
+                                                 LimitableResource::TransferMemory);
+    if (!trmem_reservation.Succeeded()) {
+        LOG_ERROR(Kernel_SVC, "Could not reserve a new transfer memory");
+        return ERR_RESOURCE_LIMIT_EXCEEDED;
+    }
     auto transfer_mem_handle = TransferMemory::Create(kernel, system.Memory(), addr, size, perms);
 
     if (const auto reserve_result{transfer_mem_handle->Reserve()}; reserve_result.IsError()) {
@@ -1895,6 +1909,7 @@ static ResultCode CreateTransferMemory(Core::System& system, Handle* handle, VAd
     if (result.Failed()) {
         return result.Code();
     }
+    trmem_reservation.Commit();
 
     *handle = *result;
     return RESULT_SUCCESS;
@@ -2002,8 +2017,17 @@ static ResultCode SetThreadCoreMask32(Core::System& system, Handle thread_handle
 static ResultCode SignalEvent(Core::System& system, Handle event_handle) {
     LOG_DEBUG(Kernel_SVC, "called, event_handle=0x{:08X}", event_handle);
 
+    auto& kernel = system.Kernel();
     // Get the current handle table.
-    const HandleTable& handle_table = system.Kernel().CurrentProcess()->GetHandleTable();
+    const HandleTable& handle_table = kernel.CurrentProcess()->GetHandleTable();
+
+    // Reserve a new event from the process resource limit.
+    KScopedResourceReservation event_reservation(kernel.CurrentProcess(),
+                                                 LimitableResource::Events);
+    if (!event_reservation.Succeeded()) {
+        LOG_ERROR(Kernel, "Could not reserve a new event");
+        return ERR_RESOURCE_LIMIT_EXCEEDED;
+    }
 
     // Get the writable event.
     auto writable_event = handle_table.Get<KWritableEvent>(event_handle);
@@ -2012,6 +2036,9 @@ static ResultCode SignalEvent(Core::System& system, Handle event_handle) {
         return ResultInvalidHandle;
     }
 
+    // Commit the successfuly reservation.
+    event_reservation.Commit();
+
     return writable_event->Signal();
 }