4 files changed, 94 insertions, 10 deletions

diff --git a/src/core/hle/kernel/mutex.cpp b/src/core/hle/kernel/mutex.cpp
index 5cc0bd266..63733ad79 100644
--- a/src/core/hle/kernel/mutex.cpp
+++ b/src/core/hle/kernel/mutex.cpp
@@ -18,13 +18,13 @@ namespace Kernel {
 
 /// Returns the number of threads that are waiting for a mutex, and the highest priority one among
 /// those.
-static std::pair<SharedPtr<Thread>, u32> GetHighestPriorityMutexWaitingThread(VAddr mutex_addr) {
-    auto& thread_list = Core::System::GetInstance().Scheduler().GetThreadList();
+static std::pair<SharedPtr<Thread>, u32> GetHighestPriorityMutexWaitingThread(
+    SharedPtr<Thread> current_thread, VAddr mutex_addr) {
 
     SharedPtr<Thread> highest_priority_thread;
     u32 num_waiters = 0;
 
-    for (auto& thread : thread_list) {
+    for (auto& thread : current_thread->wait_mutex_threads) {
         if (thread->mutex_wait_address != mutex_addr)
             continue;
 
@@ -40,6 +40,21 @@ static std::pair<SharedPtr<Thread>, u32> GetHighestPriorityMutexWaitingThread(VA
     return {highest_priority_thread, num_waiters};
 }
 
+/// Update the mutex owner field of all threads waiting on the mutex to point to the new owner.
+static void TransferMutexOwnership(VAddr mutex_addr, SharedPtr<Thread> current_thread,
+                                   SharedPtr<Thread> new_owner) {
+    auto threads = current_thread->wait_mutex_threads;
+    for (auto& thread : threads) {
+        if (thread->mutex_wait_address != mutex_addr)
+            continue;
+
+        ASSERT(thread->lock_owner == current_thread);
+        current_thread->RemoveMutexWaiter(thread);
+        if (new_owner != thread)
+            new_owner->AddMutexWaiter(thread);
+    }
+}
+
 ResultCode Mutex::TryAcquire(VAddr address, Handle holding_thread_handle,
                              Handle requesting_thread_handle) {
     // The mutex address must be 4-byte aligned
@@ -65,11 +80,14 @@ ResultCode Mutex::TryAcquire(VAddr address, Handle holding_thread_handle,
         return ERR_INVALID_HANDLE;
 
     // Wait until the mutex is released
-    requesting_thread->mutex_wait_address = address;
-    requesting_thread->wait_handle = requesting_thread_handle;
+    GetCurrentThread()->mutex_wait_address = address;
+    GetCurrentThread()->wait_handle = requesting_thread_handle;
 
-    requesting_thread->status = THREADSTATUS_WAIT_MUTEX;
-    requesting_thread->wakeup_callback = nullptr;
+    GetCurrentThread()->status = THREADSTATUS_WAIT_MUTEX;
+    GetCurrentThread()->wakeup_callback = nullptr;
+
+    // Update the lock holder thread's priority to prevent priority inversion.
+    holding_thread->AddMutexWaiter(GetCurrentThread());
 
     Core::System::GetInstance().PrepareReschedule();
 
@@ -82,14 +100,18 @@ ResultCode Mutex::Release(VAddr address) {
         return ResultCode(ErrorModule::Kernel, ErrCodes::MisalignedAddress);
     }
 
-    auto [thread, num_waiters] = GetHighestPriorityMutexWaitingThread(address);
+    auto [thread, num_waiters] = GetHighestPriorityMutexWaitingThread(GetCurrentThread(), address);
 
     // There are no more threads waiting for the mutex, release it completely.
     if (thread == nullptr) {
+        ASSERT(GetCurrentThread()->wait_mutex_threads.empty());
         Memory::Write32(address, 0);
         return RESULT_SUCCESS;
     }
 
+    // Transfer the ownership of the mutex from the previous owner to the new one.
+    TransferMutexOwnership(address, GetCurrentThread(), thread);
+
     u32 mutex_value = thread->wait_handle;
 
     if (num_waiters >= 2) {
@@ -103,6 +125,7 @@ ResultCode Mutex::Release(VAddr address) {
     ASSERT(thread->status == THREADSTATUS_WAIT_MUTEX);
     thread->ResumeFromWait();
 
+    thread->lock_owner = nullptr;
     thread->condvar_wait_address = 0;
     thread->mutex_wait_address = 0;
     thread->wait_handle = 0;
diff --git a/src/core/hle/kernel/svc.cpp b/src/core/hle/kernel/svc.cpp
index 082c36caf..a3015cf7a 100644
--- a/src/core/hle/kernel/svc.cpp
+++ b/src/core/hle/kernel/svc.cpp
@@ -621,6 +621,8 @@ static ResultCode WaitProcessWideKeyAtomic(VAddr mutex_addr, VAddr condition_var
 
     current_thread->WakeAfterDelay(nano_seconds);
 
+    // Note: Deliberately don't attempt to inherit the lock owner's priority.
+
     Core::System::GetInstance().PrepareReschedule();
     return RESULT_SUCCESS;
 }
@@ -651,6 +653,11 @@ static ResultCode SignalProcessWideKey(VAddr condition_variable_addr, s32 target
             ASSERT(thread->status == THREADSTATUS_WAIT_MUTEX);
             thread->ResumeFromWait();
 
+            auto lock_owner = thread->lock_owner;
+            if (lock_owner)
+                lock_owner->RemoveMutexWaiter(thread);
+
+            thread->lock_owner = nullptr;
             thread->mutex_wait_address = 0;
             thread->condvar_wait_address = 0;
             thread->wait_handle = 0;
@@ -666,6 +673,8 @@ static ResultCode SignalProcessWideKey(VAddr condition_variable_addr, s32 target
             // Signal that the mutex now has a waiting thread.
             Memory::Write32(thread->mutex_wait_address, mutex_val | Mutex::MutexHasWaitersFlag);
 
+            owner->AddMutexWaiter(thread);
+
             Core::System::GetInstance().PrepareReschedule();
         }
 
diff --git a/src/core/hle/kernel/thread.cpp b/src/core/hle/kernel/thread.cpp
index 16d9b9e36..36222d45f 100644
--- a/src/core/hle/kernel/thread.cpp
+++ b/src/core/hle/kernel/thread.cpp
@@ -129,6 +129,11 @@ static void ThreadWakeupCallback(u64 thread_handle, int cycles_late) {
         thread->mutex_wait_address = 0;
         thread->condvar_wait_address = 0;
         thread->wait_handle = 0;
+
+        auto lock_owner = thread->lock_owner;
+        // Threads waking up by timeout from WaitProcessWideKey do not perform priority inheritance
+        // and don't have a lock owner.
+        ASSERT(lock_owner == nullptr);
     }
 
     if (resume)
@@ -325,8 +330,8 @@ ResultVal<SharedPtr<Thread>> Thread::Create(std::string name, VAddr entry_point,
 void Thread::SetPriority(u32 priority) {
     ASSERT_MSG(priority <= THREADPRIO_LOWEST && priority >= THREADPRIO_HIGHEST,
                "Invalid priority value.");
-    Core::System::GetInstance().Scheduler().SetThreadPriority(this, priority);
-    nominal_priority = current_priority = priority;
+    nominal_priority = priority;
+    UpdatePriority();
 }
 
 void Thread::BoostPriority(u32 priority) {
@@ -376,6 +381,38 @@ VAddr Thread::GetCommandBufferAddress() const {
     return GetTLSAddress() + CommandHeaderOffset;
 }
 
+void Thread::AddMutexWaiter(SharedPtr<Thread> thread) {
+    thread->lock_owner = this;
+    wait_mutex_threads.emplace_back(std::move(thread));
+    UpdatePriority();
+}
+
+void Thread::RemoveMutexWaiter(SharedPtr<Thread> thread) {
+    boost::remove_erase(wait_mutex_threads, thread);
+    thread->lock_owner = nullptr;
+    UpdatePriority();
+}
+
+void Thread::UpdatePriority() {
+    // Find the highest priority among all the threads that are waiting for this thread's lock
+    u32 new_priority = nominal_priority;
+    for (const auto& thread : wait_mutex_threads) {
+        if (thread->nominal_priority < new_priority)
+            new_priority = thread->nominal_priority;
+    }
+
+    if (new_priority == current_priority)
+        return;
+
+    Core::System::GetInstance().Scheduler().SetThreadPriority(this, new_priority);
+
+    current_priority = new_priority;
+
+    // Recursively update the priority of the thread that depends on the priority of this one.
+    if (lock_owner)
+        lock_owner->UpdatePriority();
+}
+
 ////////////////////////////////////////////////////////////////////////////////////////////////////
 
 /**
diff --git a/src/core/hle/kernel/thread.h b/src/core/hle/kernel/thread.h
index 0e0eae28d..e0a3c0934 100644
--- a/src/core/hle/kernel/thread.h
+++ b/src/core/hle/kernel/thread.h
@@ -109,6 +109,15 @@ public:
      */
     void BoostPriority(u32 priority);
 
+    /// Adds a thread to the list of threads that are waiting for a lock held by this thread.
+    void AddMutexWaiter(SharedPtr<Thread> thread);
+
+    /// Removes a thread from the list of threads that are waiting for a lock held by this thread.
+    void RemoveMutexWaiter(SharedPtr<Thread> thread);
+
+    /// Recalculates the current priority taking into account priority inheritance.
+    void UpdatePriority();
+
     /**
      * Gets the thread's thread ID
      * @return The thread's ID
@@ -205,6 +214,12 @@ public:
     // passed to WaitSynchronization1/N.
     std::vector<SharedPtr<WaitObject>> wait_objects;
 
+    /// List of threads that are waiting for a mutex that is held by this thread.
+    std::vector<SharedPtr<Thread>> wait_mutex_threads;
+
+    /// Thread that owns the lock that this thread is waiting for.
+    SharedPtr<Thread> lock_owner;
+
     // If waiting on a ConditionVariable, this is the ConditionVariable  address
     VAddr condvar_wait_address;
     VAddr mutex_wait_address; ///< If waiting on a Mutex, this is the mutex address