[RFC PATCH 15/16] locking/mutex: Track locks owned by a task in a per-task counter

[K Prateek Nayak <kprateek.nayak@amd.com>]

With sleeping owner handling, the task activation path will need to be
extra careful when waking up a task that can resolve to a mutex_owner().
Track the lock acquisitions in a per-task counter to selectively put
mutex owner activations onto a slowpath.

The per-task accounting is done locklessly on the current task. This is
only relevant for a blocking owner and barrie in __block_task() ensures
the observer sees the correct "lock_nesting" before p->on_rq is
transitioned to 0.

The counter will be used in subsequent patches to optimize
activate_blocked_task() and __proxy_block_task().

One exception to the lock tracking rule is the increment before
schedule_preempt_disabled() and decrement after to fix the former
mis-count.

This is necessary for mutex since a handoff can set fully blocked task
as mutex_owner and, at that moment, other tasks can start queuing
themselves onto the blocked owner which requires the wakeup path to go
via slow-path and ensure blocked waiters (if any) are activated.

Signed-off-by: K Prateek Nayak <kprateek.nayak@amd.com>
---
 include/linux/sched.h  |  6 +++++
 kernel/locking/mutex.c | 57 ++++++++++++++++++++++++++++++++----------
 2 files changed, 50 insertions(+), 13 deletions(-)

diff --git a/include/linux/sched.h b/include/linux/sched.h
index a05240b52806..cb58992814fb 100644
--- a/include/linux/sched.h
+++ b/include/linux/sched.h
@@ -1272,6 +1272,12 @@ struct task_struct {
 	struct task_struct		*sleeping_owner; /* task our blocked_node is enqueued on */
 	int				blocked_cpu;	 /* CPU where task was blocked. */
 #endif
+	/*
+	 * HACK: These bits should technically live in CONFIG_SCHED_PROXY_EXEC
+	 * block and only account when sched_proxy_exec() is true but for the
+	 * PoC keep it outside and assume proxy is always enabled ;-)
+	 */
+	unsigned int			lock_nesting;
 
 	/*
 	 * The task that is boosting this task; a back link for the current
diff --git a/kernel/locking/mutex.c b/kernel/locking/mutex.c
index 8a85912d7ee6..23a63e117c3e 100644
--- a/kernel/locking/mutex.c
+++ b/kernel/locking/mutex.c
@@ -84,7 +84,8 @@ unsigned long mutex_get_owner(struct mutex *lock)
  */
 static inline struct task_struct *__mutex_trylock_common(struct mutex *lock, bool handoff)
 {
-	unsigned long owner, curr = (unsigned long)current;
+	struct task_struct *cur_task = current;
+	unsigned long owner, curr = (unsigned long)cur_task;
 
 	owner = atomic_long_read(&lock->owner);
 	for (;;) { /* must loop, can race against a flag */
@@ -109,8 +110,10 @@ static inline struct task_struct *__mutex_trylock_common(struct mutex *lock, boo
 		}
 
 		if (atomic_long_try_cmpxchg_acquire(&lock->owner, &owner, task | flags)) {
-			if (task == curr)
+			if (task == curr) {
+				cur_task->lock_nesting++;
 				return NULL;
+			}
 			break;
 		}
 	}
@@ -153,13 +156,16 @@ EXPORT_SYMBOL(mutex_init_generic);
 static __always_inline bool __mutex_trylock_fast(struct mutex *lock)
 	__cond_acquires(true, lock)
 {
-	unsigned long curr = (unsigned long)current;
+	struct task_struct *cur_task = current;
+	unsigned long curr = (unsigned long)cur_task;
 	unsigned long zero = 0UL;
 
 	MUTEX_WARN_ON(lock->magic != lock);
 
-	if (atomic_long_try_cmpxchg_acquire(&lock->owner, &zero, curr))
+	if (atomic_long_try_cmpxchg_acquire(&lock->owner, &zero, curr)) {
+		cur_task->lock_nesting++;
 		return true;
+	}
 
 	return false;
 }
@@ -167,9 +173,15 @@ static __always_inline bool __mutex_trylock_fast(struct mutex *lock)
 static __always_inline bool __mutex_unlock_fast(struct mutex *lock)
 	__cond_releases(true, lock)
 {
-	unsigned long curr = (unsigned long)current;
+	struct task_struct *cur_task = current;
+	unsigned long curr = (unsigned long)cur_task;
+
+	if (atomic_long_try_cmpxchg_release(&lock->owner, &curr, 0UL)) {
+		cur_task->lock_nesting--;
+		return true;
+	}
 
-	return atomic_long_try_cmpxchg_release(&lock->owner, &curr, 0UL);
+	return false;
 }
 
 #else /* !CONFIG_DEBUG_LOCK_ALLOC */
@@ -723,8 +735,23 @@ __mutex_lock_common(struct mutex *lock, unsigned int state, unsigned int subclas
 
 		raw_spin_unlock_irqrestore_wake(&lock->wait_lock, flags, &wake_q);
 
+		/*
+		 * Mutex handoff can set the current as the lock owner while
+		 * it is blocked. Increment "nesting_count" on the way to
+		 * schedule() to notify the wakeup path that __mutex_owner()
+		 * can resolve to this task on the way back.
+		 */
+		current->lock_nesting++;
+
 		schedule_preempt_disabled();
 
+		/*
+		 * Fix the guard against handoff. If the task got the lock,
+		 * __mutex_trylock*() will increment the lock_nesting when the
+		 * acquisition is finalized before exiting.
+		 */
+		current->lock_nesting--;
+
 		first = lock->first_waiter == &waiter;
 
 		raw_spin_lock_irqsave(&lock->wait_lock, flags);
@@ -981,7 +1008,7 @@ EXPORT_SYMBOL_GPL(ww_mutex_lock_interruptible);
 static noinline void __sched __mutex_unlock_slowpath(struct mutex *lock, unsigned long ip)
 	__releases(lock)
 {
-	struct task_struct *donor, *next = NULL;
+	struct task_struct *cur_task, *donor, *next = NULL;
 	struct mutex_waiter *waiter;
 	unsigned long owner;
 	unsigned long flags;
@@ -997,6 +1024,10 @@ static noinline void __sched __mutex_unlock_slowpath(struct mutex *lock, unsigne
 	 * missed.
 	 */
 	guard(preempt)();
+
+	cur_task = current;
+	cur_task->lock_nesting--;
+
 	/*
 	 * Release the lock before (potentially) taking the spinlock such that
 	 * other contenders can get on with things ASAP.
@@ -1006,10 +1037,10 @@ static noinline void __sched __mutex_unlock_slowpath(struct mutex *lock, unsigne
 	 */
 	owner = atomic_long_read(&lock->owner);
 	for (;;) {
-		MUTEX_WARN_ON(__owner_task(owner) != current);
+		MUTEX_WARN_ON(__owner_task(owner) != cur_task);
 		MUTEX_WARN_ON(owner & MUTEX_FLAG_PICKUP);
 
-		if (sched_proxy_exec() && current->blocked_donor) {
+		if (sched_proxy_exec() && cur_task->blocked_donor) {
 			/* force handoff if we have a blocked_donor */
 			owner = MUTEX_FLAG_HANDOFF;
 			break;
@@ -1027,7 +1058,7 @@ static noinline void __sched __mutex_unlock_slowpath(struct mutex *lock, unsigne
 	}
 
 	raw_spin_lock_irqsave(&lock->wait_lock, flags);
-	raw_spin_lock(&current->blocked_lock);
+	raw_spin_lock(&cur_task->blocked_lock);
 	debug_mutex_unlock(lock);
 
 	if (sched_proxy_exec()) {
@@ -1036,7 +1067,7 @@ static noinline void __sched __mutex_unlock_slowpath(struct mutex *lock, unsigne
 		 * current through this lock, hand the lock to that task, as that
 		 * is the highest waiter, as selected by the scheduling function.
 		 */
-		donor = current->blocked_donor;
+		donor = cur_task->blocked_donor;
 		if (donor) {
 			struct mutex *next_lock;
 
@@ -1045,7 +1076,7 @@ static noinline void __sched __mutex_unlock_slowpath(struct mutex *lock, unsigne
 			if (next_lock == lock) {
 				next = get_task_struct(donor);
 				__clear_task_blocked_on(next, lock);
-				current->blocked_donor = NULL;
+				cur_task->blocked_donor = NULL;
 			}
 			raw_spin_unlock(&donor->blocked_lock);
 		}
@@ -1071,7 +1102,7 @@ static noinline void __sched __mutex_unlock_slowpath(struct mutex *lock, unsigne
 	if (owner & MUTEX_FLAG_HANDOFF)
 		__mutex_handoff(lock, next);
 
-	raw_spin_unlock(&current->blocked_lock);
+	raw_spin_unlock(&cur_task->blocked_lock);
 	raw_spin_unlock_irqrestore(&lock->wait_lock, flags);
 	if (next) {
 		wake_up_process(next);
-- 
2.34.1