kern: refactor priority inheritance to represent locks as C++ objects

libraries/libmesosphere/include/mesosphere/kern_k_thread.hpp

@@ -201,6 +201,28 @@ namespace ams::kern {
             };
             static_assert(ams::util::HasRedBlackKeyType<ConditionVariableComparator>);
             static_assert(std::same_as<ams::util::RedBlackKeyType<ConditionVariableComparator, void>, ConditionVariableComparator::RedBlackKeyType>);
+
+            struct LockWithPriorityInheritanceComparator {
+                struct RedBlackKeyType {
+                    s32 m_priority;
+
+                    constexpr ALWAYS_INLINE s32 GetPriority() const {
+                        return m_priority;
+                    }
+                };
+
+                template<typename T> requires (std::same_as<T, KThread> || std::same_as<T, RedBlackKeyType>)
+                static constexpr ALWAYS_INLINE int Compare(const T &lhs, const KThread &rhs) {
+                    if (lhs.GetPriority() < rhs.GetPriority()) {
+                        /* And then by priority. */
+                        return -1;
+                    } else {
+                        return 1;
+                    }
+                }
+            };
+            static_assert(ams::util::HasRedBlackKeyType<LockWithPriorityInheritanceComparator>);
+            static_assert(std::same_as<ams::util::RedBlackKeyType<LockWithPriorityInheritanceComparator, void>, LockWithPriorityInheritanceComparator::RedBlackKeyType>);
         private:
             util::IntrusiveListNode         m_process_list_node;
             util::IntrusiveRedBlackTreeNode m_condvar_arbiter_tree_node;
@@ -209,6 +231,67 @@ namespace ams::kern {
             using ConditionVariableThreadTreeTraits = util::IntrusiveRedBlackTreeMemberTraitsDeferredAssert<&KThread::m_condvar_arbiter_tree_node>;
             using ConditionVariableThreadTree       = ConditionVariableThreadTreeTraits::TreeType<ConditionVariableComparator>;
 
+            using LockWithPriorityInheritanceThreadTreeTraits = util::IntrusiveRedBlackTreeMemberTraitsDeferredAssert<&KThread::m_condvar_arbiter_tree_node>;
+            using LockWithPriorityInheritanceThreadTree       = ConditionVariableThreadTreeTraits::TreeType<LockWithPriorityInheritanceComparator>;
+        public:
+            class LockWithPriorityInheritanceInfo : public KSlabAllocated<LockWithPriorityInheritanceInfo>, public util::IntrusiveListBaseNode<LockWithPriorityInheritanceInfo> {
+                private:
+                    LockWithPriorityInheritanceThreadTree m_tree;
+                    KProcessAddress m_address_key;
+                    KThread *m_owner;
+                    u32 m_waiter_count;
+                public:
+                    constexpr LockWithPriorityInheritanceInfo() : m_tree(), m_address_key(Null<KProcessAddress>), m_owner(nullptr), m_waiter_count() {
+                        /* ... */
+                    }
+
+                    static LockWithPriorityInheritanceInfo *Create(KProcessAddress address_key) {
+                        /* Create a new lock info. */
+                        auto *new_lock = LockWithPriorityInheritanceInfo::Allocate();
+                        MESOSPHERE_ABORT_UNLESS(new_lock != nullptr);
+
+                        /* Set the new lock's address key. */
+                        new_lock->m_address_key = address_key;
+
+                        return new_lock;
+                    }
+
+                    void SetOwner(KThread *new_owner) {
+                        /* Set new owner. */
+                        m_owner = new_owner;
+                    }
+
+                    void AddWaiter(KThread *waiter) {
+                        /* Insert the waiter. */
+                        m_tree.insert(*waiter);
+                        m_waiter_count++;
+
+                        waiter->SetWaitingLockInfo(this);
+                    }
+
+                    [[nodiscard]] bool RemoveWaiter(KThread *waiter) {
+                        m_tree.erase(m_tree.iterator_to(*waiter));
+
+                        waiter->SetWaitingLockInfo(nullptr);
+
+                        return (--m_waiter_count) == 0;
+                    }
+
+                    KThread *GetHighestPriorityWaiter() { return std::addressof(m_tree.front()); }
+                    const KThread *GetHighestPriorityWaiter() const { return std::addressof(m_tree.front()); }
+
+                    LockWithPriorityInheritanceThreadTree &GetThreadTree() { return m_tree; }
+                    const LockWithPriorityInheritanceThreadTree &GetThreadTree() const { return m_tree; }
+
+                    constexpr KProcessAddress GetAddressKey() const { return m_address_key; }
+
+                    constexpr KThread *GetOwner() const { return m_owner; }
+
+                    constexpr u32 GetWaiterCount() const { return m_waiter_count; }
+            };
+        private:
+            using LockWithPriorityInheritanceInfoList = util::IntrusiveListBaseTraits<LockWithPriorityInheritanceInfo>::ListType;
+
             ConditionVariableThreadTree                        *m_condvar_tree;
             uintptr_t                                           m_condvar_key;
             alignas(16) KThreadContext::CallerSaveFpuRegisters  m_caller_save_fpu_registers;
@@ -228,9 +311,9 @@ namespace ams::kern {
             s64                                                 m_last_scheduled_tick;
             QueueEntry                                          m_per_core_priority_queue_entry[cpu::NumCores];
             KThreadQueue                                       *m_wait_queue;
-            WaiterList                                          m_waiter_list;
+            LockWithPriorityInheritanceInfoList                 m_held_lock_info_list;
+            LockWithPriorityInheritanceInfo                    *m_waiting_lock_info;
             WaiterList                                          m_pinned_waiter_list;
-            KThread                                            *m_lock_owner;
             uintptr_t                                           m_debug_params[3];
             KAutoObject                                        *m_closed_object;
             u32                                                 m_address_key_value;
@@ -264,8 +347,8 @@ namespace ams::kern {
                   m_process_list_node{}, m_condvar_arbiter_tree_node{util::ConstantInitialize}, m_priority{-1}, m_condvar_tree{}, m_condvar_key{},
                   m_caller_save_fpu_registers{}, m_virtual_affinity_mask{}, m_physical_affinity_mask{}, m_thread_id{}, m_cpu_time{0}, m_address_key{Null<KProcessAddress>}, m_parent{},
                   m_kernel_stack_top{}, m_light_ipc_data{}, m_tls_address{Null<KProcessAddress>}, m_tls_heap_address{}, m_activity_pause_lock{}, m_sync_object_buffer{util::ConstantInitialize},
-                  m_schedule_count{}, m_last_scheduled_tick{}, m_per_core_priority_queue_entry{}, m_wait_queue{}, m_waiter_list{}, m_pinned_waiter_list{},
+                  m_schedule_count{}, m_last_scheduled_tick{}, m_per_core_priority_queue_entry{}, m_wait_queue{}, m_held_lock_info_list{}, m_pinned_waiter_list{},
-                  m_lock_owner{}, m_debug_params{}, m_closed_object{}, m_address_key_value{}, m_suspend_request_flags{}, m_suspend_allowed_flags{}, m_synced_index{},
+                  m_waiting_lock_info{}, m_debug_params{}, m_closed_object{}, m_address_key_value{}, m_suspend_request_flags{}, m_suspend_allowed_flags{}, m_synced_index{},
                   m_wait_result{svc::ResultNoSynchronizationObject()}, m_debug_exception_result{ResultSuccess()}, m_base_priority{}, m_base_priority_on_unpin{},
                   m_physical_ideal_core_id{}, m_virtual_ideal_core_id{}, m_num_kernel_waiters{}, m_current_core_id{}, m_core_id{}, m_original_physical_affinity_mask{},
                   m_original_physical_ideal_core_id{}, m_num_core_migration_disables{}, m_thread_state{}, m_termination_requested{false}, m_wait_cancelled{},
@@ -411,6 +494,10 @@ namespace ams::kern {
             void ClearUsermodeExceptionSvcPermissions();
         private:
             void UpdateState();
+
+            ALWAYS_INLINE void AddHeldLock(LockWithPriorityInheritanceInfo *lock_info);
+            ALWAYS_INLINE LockWithPriorityInheritanceInfo *FindHeldLock(KProcessAddress address_key);
+
             ALWAYS_INLINE void AddWaiterImpl(KThread *thread);
             ALWAYS_INLINE void RemoveWaiterImpl(KThread *thread);
             ALWAYS_INLINE static void RestorePriority(KThread *thread);
@@ -445,6 +532,8 @@ namespace ams::kern {
             constexpr uintptr_t GetAddressArbiterKey() const { return m_condvar_key; }
 
             constexpr void SetConditionVariable(ConditionVariableThreadTree *tree, KProcessAddress address, uintptr_t cv_key, u32 value) {
+                MESOSPHERE_ASSERT(m_waiting_lock_info == nullptr);
+
                 m_condvar_tree      = tree;
                 m_condvar_key       = cv_key;
                 m_address_key       = address;
@@ -460,6 +549,8 @@ namespace ams::kern {
             }
 
             constexpr void SetAddressArbiter(ConditionVariableThreadTree *tree, uintptr_t address) {
+                MESOSPHERE_ASSERT(m_waiting_lock_info == nullptr);
+
                 m_condvar_tree = tree;
                 m_condvar_key  = address;
             }
@@ -495,15 +586,17 @@ namespace ams::kern {
 
             void AddWaiter(KThread *thread);
             void RemoveWaiter(KThread *thread);
-            KThread *RemoveWaiterByKey(s32 *out_num_waiters, KProcessAddress key);
+            KThread *RemoveWaiterByKey(bool *out_has_waiters, KProcessAddress key);
 
             constexpr KProcessAddress GetAddressKey() const { return m_address_key; }
             constexpr u32 GetAddressKeyValue() const { return m_address_key_value; }
-            constexpr void SetAddressKey(KProcessAddress key) { m_address_key = key; }
+            constexpr void SetAddressKey(KProcessAddress key) { MESOSPHERE_ASSERT(m_waiting_lock_info == nullptr); m_address_key = key; }
-            constexpr void SetAddressKey(KProcessAddress key, u32 val) { m_address_key = key; m_address_key_value = val; }
+            constexpr void SetAddressKey(KProcessAddress key, u32 val) { MESOSPHERE_ASSERT(m_waiting_lock_info == nullptr); m_address_key = key; m_address_key_value = val; }
 
-            constexpr void SetLockOwner(KThread *owner) { m_lock_owner = owner; }
+            constexpr void SetWaitingLockInfo(LockWithPriorityInheritanceInfo *lock) { m_waiting_lock_info = lock; }
-            constexpr KThread *GetLockOwner() const { return m_lock_owner; }
+            constexpr LockWithPriorityInheritanceInfo *GetWaitingLockInfo() { return m_waiting_lock_info; }
+
+            constexpr KThread *GetLockOwner() const { return m_waiting_lock_info != nullptr ? m_waiting_lock_info->GetOwner() : nullptr; }
 
             constexpr void ClearWaitQueue() { m_wait_queue = nullptr; }
 
@@ -533,8 +626,6 @@ namespace ams::kern {
             constexpr u32 *GetLightSessionData() const { return m_light_ipc_data; }
             constexpr void SetLightSessionData(u32 *data) { m_light_ipc_data = data; }
 
-            bool HasWaiters() const { return !m_waiter_list.empty(); }
-
             constexpr s64 GetLastScheduledTick() const { return m_last_scheduled_tick; }
             constexpr void SetLastScheduledTick(s64 tick) { m_last_scheduled_tick = tick; }
 

libraries/libmesosphere/source/init/kern_init_slab_setup.cpp

@@ -19,6 +19,7 @@ namespace ams::kern::init {
 
     /* For macro convenience. */
     using KSessionRequestMappings = KSessionRequest::SessionMappings::DynamicMappings;
+    using KThreadLockInfo         = KThread::LockWithPriorityInheritanceInfo;
 
     #define SLAB_COUNT(CLASS) g_slab_resource_counts.num_##CLASS
 
@@ -43,8 +44,9 @@ namespace ams::kern::init {
         HANDLER(KDebug,                  (SLAB_COUNT(KDebug)),                                                      ## __VA_ARGS__) \
         HANDLER(KIoPool,                 (SLAB_COUNT(KIoPool)),                                                     ## __VA_ARGS__) \
         HANDLER(KIoRegion,               (SLAB_COUNT(KIoRegion)),                                                   ## __VA_ARGS__) \
+        HANDLER(KSessionRequestMappings, (SLAB_COUNT(KSessionRequestMappings)),                                     ## __VA_ARGS__) \
         HANDLER(KSecureSystemResource,   (SLAB_COUNT(KProcess)),                                                    ## __VA_ARGS__) \
-        HANDLER(KSessionRequestMappings, (SLAB_COUNT(KSessionRequestMappings)),                                     ## __VA_ARGS__)
+        HANDLER(KThreadLockInfo,         (SLAB_COUNT(KThread)),                                                     ## __VA_ARGS__)
 
     namespace {
 

libraries/libmesosphere/source/kern_k_condition_variable.cpp

@@ -77,14 +77,14 @@ namespace ams::kern {
             KScopedSchedulerLock sl;
 
             /* Remove waiter thread. */
-            s32 num_waiters;
+            bool has_waiters;
-            KThread * const next_owner_thread = owner_thread->RemoveWaiterByKey(std::addressof(num_waiters), addr);
+            KThread * const next_owner_thread = owner_thread->RemoveWaiterByKey(std::addressof(has_waiters), addr);
 
             /* Determine the next tag. */
             u32 next_value = 0;
             if (next_owner_thread != nullptr) {
                 next_value = next_owner_thread->GetAddressKeyValue();
-                if (num_waiters > 1) {
+                if (has_waiters) {
                     next_value |= ams::svc::HandleWaitMask;
                 }
             }
@@ -200,9 +200,11 @@ namespace ams::kern {
             while ((it != m_tree.end()) && (count <= 0 || num_waiters < count) && (it->GetConditionVariableKey() == cv_key)) {
                 KThread *target_thread = std::addressof(*it);
 
-                this->SignalImpl(target_thread);
                 it = m_tree.erase(it);
                 target_thread->ClearConditionVariable();
+
+                this->SignalImpl(target_thread);
+
                 ++num_waiters;
             }
 
@@ -232,15 +234,15 @@ namespace ams::kern {
             /* Update the value and process for the next owner. */
             {
                 /* Remove waiter thread. */
-                s32 num_waiters;
+                bool has_waiters;
-                KThread *next_owner_thread = cur_thread->RemoveWaiterByKey(std::addressof(num_waiters), GetInteger(addr));
+                KThread *next_owner_thread = cur_thread->RemoveWaiterByKey(std::addressof(has_waiters), GetInteger(addr));
 
                 /* Update for the next owner thread. */
                 u32 next_value = 0;
                 if (next_owner_thread != nullptr) {
                     /* Get the next tag value. */
                     next_value = next_owner_thread->GetAddressKeyValue();
-                    if (num_waiters > 1) {
+                    if (has_waiters) {
                         next_value |= ams::svc::HandleWaitMask;
                     }
 

libraries/libmesosphere/source/kern_k_light_lock.cpp

@@ -68,13 +68,13 @@ namespace ams::kern {
             KScopedSchedulerLock sl;
 
             /* Get the next owner. */
-            s32 num_waiters;
+            bool has_waiters;
-            KThread *next_owner = owner_thread->RemoveWaiterByKey(std::addressof(num_waiters), reinterpret_cast<uintptr_t>(std::addressof(m_tag)));
+            KThread *next_owner = owner_thread->RemoveWaiterByKey(std::addressof(has_waiters), reinterpret_cast<uintptr_t>(std::addressof(m_tag)));
 
             /* Pass the lock to the next owner. */
             uintptr_t next_tag = 0;
             if (next_owner != nullptr) {
-                next_tag = reinterpret_cast<uintptr_t>(next_owner) | static_cast<uintptr_t>(num_waiters > 1);
+                next_tag = reinterpret_cast<uintptr_t>(next_owner) | static_cast<uintptr_t>(has_waiters);
 
                 next_owner->EndWait(ResultSuccess());
 

libraries/libmesosphere/source/kern_k_process.cpp

@@ -817,8 +817,8 @@ namespace ams::kern {
             m_exception_thread = nullptr;
 
             /* Remove waiter thread. */
-            s32 num_waiters;
+            bool has_waiters;
-            if (KThread *next = thread->RemoveWaiterByKey(std::addressof(num_waiters), reinterpret_cast<uintptr_t>(std::addressof(m_exception_thread)) | 1); next != nullptr) {
+            if (KThread *next = thread->RemoveWaiterByKey(std::addressof(has_waiters), reinterpret_cast<uintptr_t>(std::addressof(m_exception_thread)) | 1); next != nullptr) {
                 next->EndWait(ResultSuccess());
             }
 

libraries/libmesosphere/source/kern_k_scheduler.cpp

@@ -243,7 +243,7 @@ namespace ams::kern {
             if (AMS_LIKELY(!cur_thread->IsTerminationRequested()) && AMS_LIKELY(cur_thread->GetActiveCore() == m_core_id)) {
                 m_state.prev_thread = cur_thread;
             } else {
-                m_state.prev_thread =nullptr;
+                m_state.prev_thread = nullptr;
             }
         }
 

libraries/libmesosphere/source/kern_k_thread.cpp

@@ -254,7 +254,7 @@ namespace ams::kern {
         m_light_ipc_data                = nullptr;
 
         /* We're not waiting for a lock, and we haven't disabled migration. */
-        m_lock_owner                    = nullptr;
+        m_waiting_lock_info             = nullptr;
         m_num_core_migration_disables   = 0;
 
         /* We have no waiters, and no closed objects. */
@@ -397,25 +397,39 @@ namespace ams::kern {
 
         /* Release any waiters. */
         {
-            MESOSPHERE_ASSERT(m_lock_owner == nullptr);
+            MESOSPHERE_ASSERT(m_waiting_lock_info == nullptr);
             KScopedSchedulerLock sl;
 
-            auto it = m_waiter_list.begin();
+            /* Check that we have no kernel waiters. */
-            while (it != m_waiter_list.end()) {
+            MESOSPHERE_ABORT_UNLESS(m_num_kernel_waiters == 0);
-                /* Get the thread. */
-                KThread * const waiter = std::addressof(*it);
 
-                /* The thread shouldn't be a kernel waiter. */
+            auto it = m_held_lock_info_list.begin();
-                MESOSPHERE_ASSERT(!IsKernelAddressKey(waiter->GetAddressKey()));
+            while (it != m_held_lock_info_list.end()) {
+                /* Get the lock info. */
+                auto * const lock_info = std::addressof(*it);
 
-                /* Clear the lock owner. */
+                /* The lock shouldn't have a kernel waiter. */
-                waiter->SetLockOwner(nullptr);
+                MESOSPHERE_ASSERT(!IsKernelAddressKey(lock_info->GetAddressKey()));
 
-                /* Erase the waiter from our list. */
+                /* Remove all waiters. */
-                it = m_waiter_list.erase(it);
+                while (lock_info->GetWaiterCount() != 0) {
+                    /* Get the front waiter. */
+                    KThread * const waiter = lock_info->GetHighestPriorityWaiter();
 
-                /* Cancel the thread's wait. */
+                    /* Remove it from the lock. */
-                waiter->CancelWait(svc::ResultInvalidState(), true);
+                    if (lock_info->RemoveWaiter(waiter)) {
+                        MESOSPHERE_ASSERT(lock_info->GetWaiterCount() == 0);
+                    }
+
+                    /* Cancel the thread's wait. */
+                    waiter->CancelWait(svc::ResultInvalidState(), true);
+                }
+
+                /* Remove the held lock from our list. */
+                it = m_held_lock_info_list.erase(it);
+
+                /* Free the lock info. */
+                LockWithPriorityInheritanceInfo::Free(lock_info);
             }
         }
 
@@ -823,11 +837,8 @@ namespace ams::kern {
     void KThread::IncreaseBasePriority(s32 priority) {
         MESOSPHERE_ASSERT_THIS();
         MESOSPHERE_ASSERT(ams::svc::HighestThreadPriority <= priority && priority <= ams::svc::LowestThreadPriority);
-
+        MESOSPHERE_ASSERT(KScheduler::IsSchedulerLockedByCurrentThread());
-        /* Set our unpin base priority, if we're pinned. */
+        MESOSPHERE_ASSERT(!this->GetStackParameters().is_pinned);
-        if (this->GetStackParameters().is_pinned && m_base_priority_on_unpin > priority) {
-            m_base_priority_on_unpin = priority;
-        }
 
         /* Set our base priority. */
         if (m_base_priority > priority) {
@@ -1044,28 +1055,58 @@ namespace ams::kern {
         R_SUCCEED();
     }
 
-    void KThread::AddWaiterImpl(KThread *thread) {
+    void KThread::AddHeldLock(LockWithPriorityInheritanceInfo *lock_info) {
         MESOSPHERE_ASSERT_THIS();
         MESOSPHERE_ASSERT(KScheduler::IsSchedulerLockedByCurrentThread());
 
-        /* Find the right spot to insert the waiter. */
+        /* Set ourselves as the lock's owner. */
-        auto it = m_waiter_list.begin();
+        lock_info->SetOwner(this);
-        while (it != m_waiter_list.end()) {
+
-            if (it->GetPriority() > thread->GetPriority()) {
+        /* Add the lock to our held list. */
-                break;
+        m_held_lock_info_list.push_front(*lock_info);
+    }
+
+    KThread::LockWithPriorityInheritanceInfo *KThread::FindHeldLock(KProcessAddress address_key) {
+        MESOSPHERE_ASSERT_THIS();
+        MESOSPHERE_ASSERT(KScheduler::IsSchedulerLockedByCurrentThread());
+
+        /* Try to find an existing held lock. */
+        for (auto &held_lock : m_held_lock_info_list) {
+            if (held_lock.GetAddressKey() == address_key) {
+                return std::addressof(held_lock);
             }
-            it++;
         }
 
+        return nullptr;
+    }
+
+
+    void KThread::AddWaiterImpl(KThread *thread) {
+        MESOSPHERE_ASSERT_THIS();
+        MESOSPHERE_ASSERT(KScheduler::IsSchedulerLockedByCurrentThread());
+        MESOSPHERE_ASSERT(thread->GetConditionVariableTree() == nullptr);
+
+        /* Get the thread's address key. */
+        const auto address_key = thread->GetAddressKey();
+
         /* Keep track of how many kernel waiters we have. */
-        if (IsKernelAddressKey(thread->GetAddressKey())) {
+        if (IsKernelAddressKey(address_key)) {
             MESOSPHERE_ABORT_UNLESS((m_num_kernel_waiters++) >= 0);
             KScheduler::SetSchedulerUpdateNeeded();
         }
 
-        /* Insert the waiter. */
+        /* Get the relevant lock info. */
-        m_waiter_list.insert(it, *thread);
+        auto *lock_info = this->FindHeldLock(address_key);
-        thread->SetLockOwner(this);
+        if (lock_info == nullptr) {
+            /* Create a new lock for the address key. */
+            lock_info = LockWithPriorityInheritanceInfo::Create(address_key);
+
+            /* Add the new lock to our list. */
+            this->AddHeldLock(lock_info);
+        }
+
+        /* Add the thread as waiter to the lock info. */
+        lock_info->AddWaiter(thread);
     }
 
     void KThread::RemoveWaiterImpl(KThread *thread) {
@@ -1078,19 +1119,25 @@ namespace ams::kern {
             KScheduler::SetSchedulerUpdateNeeded();
         }
 
+        /* Get the info for the lock the thread is waiting on. */
+        auto *lock_info = thread->GetWaitingLockInfo();
+        MESOSPHERE_ASSERT(lock_info->GetOwner() == this);
+
         /* Remove the waiter. */
-        m_waiter_list.erase(m_waiter_list.iterator_to(*thread));
+        if (lock_info->RemoveWaiter(thread)) {
-        thread->SetLockOwner(nullptr);
+            m_held_lock_info_list.erase(m_held_lock_info_list.iterator_to(*lock_info));
+            LockWithPriorityInheritanceInfo::Free(lock_info);
+        }
     }
 
     void KThread::RestorePriority(KThread *thread) {
         MESOSPHERE_ASSERT(KScheduler::IsSchedulerLockedByCurrentThread());
 
-        while (true) {
+        while (thread != nullptr) {
             /* We want to inherit priority where possible. */
             s32 new_priority = thread->GetBasePriority();
-            if (thread->HasWaiters()) {
+            for (const auto &held_lock : thread->m_held_lock_info_list) {
-                new_priority = std::min(new_priority, thread->m_waiter_list.front().GetPriority());
+                new_priority = std::min(new_priority, held_lock.GetHighestPriorityWaiter()->GetPriority());
             }
 
             /* If the priority we would inherit is not different from ours, don't do anything. */
@@ -1098,6 +1145,14 @@ namespace ams::kern {
                 return;
             }
 
+            /* Get the owner of whatever lock this thread is waiting on. */
+            KThread * const lock_owner = thread->GetLockOwner();
+
+            /* If the thread is waiting on some lock, remove it as a waiter to prevent violating red black tree invariants. */
+            if (lock_owner != nullptr) {
+                lock_owner->RemoveWaiterImpl(thread);
+            }
+
             /* Ensure we don't violate condition variable red black tree invariants. */
             if (auto *cv_tree = thread->GetConditionVariableTree(); cv_tree != nullptr) {
                 BeforeUpdatePriority(cv_tree, thread);
@@ -1112,73 +1167,94 @@ namespace ams::kern {
                 AfterUpdatePriority(cv_tree, thread);
             }
 
+            /* If we removed the thread from some lock's waiting list, add it back. */
+            if (lock_owner != nullptr) {
+                lock_owner->AddWaiterImpl(thread);
+            }
+
             /* Update the scheduler. */
             KScheduler::OnThreadPriorityChanged(thread, old_priority);
 
-            /* Keep the lock owner up to date. */
+            /* Continue inheriting priority. */
-            KThread *lock_owner = thread->GetLockOwner();
-            if (lock_owner == nullptr) {
-                return;
-            }
-
-            /* Update the thread in the lock owner's sorted list, and continue inheriting. */
-            lock_owner->RemoveWaiterImpl(thread);
-            lock_owner->AddWaiterImpl(thread);
             thread = lock_owner;
         }
     }
 
     void KThread::AddWaiter(KThread *thread) {
         MESOSPHERE_ASSERT_THIS();
+
         this->AddWaiterImpl(thread);
-        RestorePriority(this);
+
+        /* If the thread has a higher priority than us, we should inherit. */
+        if (thread->GetPriority() < this->GetPriority()) {
+            RestorePriority(this);
+        }
     }
 
     void KThread::RemoveWaiter(KThread *thread) {
         MESOSPHERE_ASSERT_THIS();
         this->RemoveWaiterImpl(thread);
-        RestorePriority(this);
+
+        /* If our priority is the same as the thread's (and we've inherited), we may need to restore to lower priority. */
+        if (this->GetPriority() == thread->GetPriority() && this->GetPriority() < this->GetBasePriority()) {
+            RestorePriority(this);
+        }
     }
 
-    KThread *KThread::RemoveWaiterByKey(s32 *out_num_waiters, KProcessAddress key) {
+    KThread *KThread::RemoveWaiterByKey(bool *out_has_waiters, KProcessAddress key) {
         MESOSPHERE_ASSERT_THIS();
         MESOSPHERE_ASSERT(KScheduler::IsSchedulerLockedByCurrentThread());
 
-        s32 num_waiters = 0;
+        /* Get the relevant lock info. */
-        KThread *next_lock_owner = nullptr;
+        auto *lock_info = this->FindHeldLock(key);
-        auto it = m_waiter_list.begin();
+        if (lock_info == nullptr) {
-        while (it != m_waiter_list.end()) {
+            *out_has_waiters = false;
-            if (it->GetAddressKey() == key) {
+            return nullptr;
-                KThread *thread = std::addressof(*it);
+        }
 
-                /* Keep track of how many kernel waiters we have. */
+        /* Remove the lock info from our held list. */
-                if (IsKernelAddressKey(thread->GetAddressKey())) {
+        m_held_lock_info_list.erase(m_held_lock_info_list.iterator_to(*lock_info));
-                    MESOSPHERE_ABORT_UNLESS((m_num_kernel_waiters--) > 0);
-                    KScheduler::SetSchedulerUpdateNeeded();
-                }
-                it = m_waiter_list.erase(it);
 
-                /* Update the next lock owner. */
+        /* Keep track of how many kernel waiters we have. */
-                if (next_lock_owner == nullptr) {
+        if (IsKernelAddressKey(lock_info->GetAddressKey())) {
-                    next_lock_owner = thread;
+            m_num_kernel_waiters -= lock_info->GetWaiterCount();
-                    next_lock_owner->SetLockOwner(nullptr);
+            MESOSPHERE_ABORT_UNLESS(m_num_kernel_waiters >= 0);
-                } else {
+            KScheduler::SetSchedulerUpdateNeeded();
-                    next_lock_owner->AddWaiterImpl(thread);
+        }
-                }
+
-                num_waiters++;
+        MESOSPHERE_ASSERT(lock_info->GetWaiterCount() > 0);
-            } else {
+
-                it++;
+        /* Remove the highest priority waiter from the lock to be the next owner. */
+        KThread *next_lock_owner = lock_info->GetHighestPriorityWaiter();
+        if (lock_info->RemoveWaiter(next_lock_owner)) {
+            /* The new owner was the only waiter. */
+            *out_has_waiters = false;
+
+            /* Free the lock info, since it has no waiters. */
+            LockWithPriorityInheritanceInfo::Free(lock_info);
+        } else {
+            /* There are additional waiters on the lock. */
+            *out_has_waiters = true;
+
+            /* Add the lock to the new owner's held list. */
+            next_lock_owner->AddHeldLock(lock_info);
+
+            /* Keep track of any kernel waiters for the new owner. */
+            if (IsKernelAddressKey(lock_info->GetAddressKey())) {
+                next_lock_owner->m_num_kernel_waiters += lock_info->GetWaiterCount();
+                MESOSPHERE_ABORT_UNLESS(next_lock_owner->m_num_kernel_waiters > 0);
+
+                /* NOTE: No need to set scheduler update needed, because we will have already done so when removing earlier. */
             }
         }
 
-        /* Do priority updates, if we have a next owner. */
+        /* If our priority is the same as the next owner's (and we've inherited), we may need to restore to lower priority. */
-        if (next_lock_owner) {
+        if (this->GetPriority() == next_lock_owner->GetPriority() && this->GetPriority() < this->GetBasePriority()) {
             RestorePriority(this);
-            RestorePriority(next_lock_owner);
+            /* NOTE: No need to restore priority on the next lock owner, because it was already the highest priority waiter on the lock. */
         }
 
-        /* Return output. */
+        /* Return the next lock owner. */
-        *out_num_waiters = num_waiters;
         return next_lock_owner;
     }
 
@@ -1309,9 +1385,7 @@ namespace ams::kern {
             }
 
             /* Change the thread's priority to be higher than any system thread's. */
-            if (this->GetBasePriority() >= ams::svc::SystemThreadPriorityHighest) {
+            this->IncreaseBasePriority(TerminatingThreadPriority);
-                this->SetBasePriority(TerminatingThreadPriority);
-            }
 
             /* If the thread is runnable, send a termination interrupt to other cores. */
             if (this->GetState() == ThreadState_Runnable) {