Booting a machine with CONFIG_SCHED_PROXY_EXEC=y leads to a PSI warning
similar to following early into the boot:

    psi: inconsistent task state! task=... cpu=... psi_flags=4 clear=0 set=4

Investigating the set of events that led to warning indicated that
psi_sched_switch() never dequeued the signals for a blocked task since
it was retained on the runqueue as a potential donor but the
psi_enqueue() that follows the wakeup tries to re-enqueue these signals
for the blocked tasks thus leading to the inconsistent state warning.

When fixing PSI state corruption for delayed dequeue in commit
c6508124193d ("sched/psi: Fix mistaken CPU pressure indication after
corrupted task state bug"), Johannes mentioned that delayed tasks are
in-fact blocked and the PSI signals should treat them as such in [1].

With proxy execution, the same argument holds true for blocked tasks
that are kept around to donate the vruntime context - the tasks are
essentially blocked and their PSI signals should be treated as such
until they are woken up again.

Treat is_proxy_task() as blocked for psi_sched_switch() and call
psi_enqueue(ENQUEUE_WAKEUP) when the "sched_proxy" signal is cleared.
For all the transitions in-between, treat the task similar to a delayed
task and just move the block signals on migration of blocked donor.

Fixes: be41bde4c3a8 ("sched: Add an initial sketch of the find_proxy_task() function")
Link: https://lore.kernel.org/all/20241010130316.GA181795@cmpxchg.org/ [1]
Signed-off-by: K Prateek Nayak <kprateek.nayak@amd.com>
---
 kernel/sched/core.c  | 9 +++++++--
 kernel/sched/stats.h | 2 +-
 2 files changed, 8 insertions(+), 3 deletions(-)

diff --git a/kernel/sched/core.c b/kernel/sched/core.c
index d6265f38e93a..765365b81b12 100644
--- a/kernel/sched/core.c
+++ b/kernel/sched/core.c
@@ -3696,8 +3696,12 @@ static int ttwu_runnable(struct task_struct *p, int wake_flags)
 	rq = __task_rq_lock(p, &rf);
 	if (task_on_rq_queued(p)) {
 		update_rq_clock(rq);
-		if (is_proxy_task(p))
+		if (is_proxy_task(p)) {
 			clear_task_proxy(p);
+			/* Task was never fully blocked to be delayed. */
+			WARN_ON_ONCE(p->se.sched_delayed);
+			psi_enqueue(p, ENQUEUE_WAKEUP);
+		}
 		if (p->se.sched_delayed)
 			enqueue_task(rq, p, ENQUEUE_NOCLOCK | ENQUEUE_DELAYED);
 		if (!task_on_cpu(rq, p)) {
@@ -6903,7 +6907,8 @@ static void __sched notrace __schedule(int sched_mode)
 
 		psi_account_irqtime(rq, prev, next);
 		psi_sched_switch(prev, next, !task_on_rq_queued(prev) ||
-					     prev->se.sched_delayed);
+					     prev->se.sched_delayed ||
+					     is_proxy_task(prev));
 
 		trace_sched_switch(preempt, prev, next, prev_state);
 
diff --git a/kernel/sched/stats.h b/kernel/sched/stats.h
index 3323d773fec3..b2d7461c1ea9 100644
--- a/kernel/sched/stats.h
+++ b/kernel/sched/stats.h
@@ -142,7 +142,7 @@ static inline void psi_enqueue(struct task_struct *p, int flags)
 	if (task_on_cpu(task_rq(p), p))
 		return;
 
-	if (p->se.sched_delayed) {
+	if (p->se.sched_delayed || is_proxy_task(p)) {
 		/* CPU migration of "sleeping" task */
 		WARN_ON_ONCE(!(flags & ENQUEUE_MIGRATED));
 		if (p->in_memstall)
-- 
2.34.1