[PATCH v33 1/7] sched/psi: Preserve sleep-persistent PSI flags on off-CPU sleeping dequeues

[John Stultz <jstultz@google.com>]

Normally, when a task goes to sleep (DEQUEUE_SLEEP),
psi_dequeue() does nothing because psi_sched_switch() ->
psi_task_switch() immediately follows and handles the
runnable-to-sleeping PSI transition (clearing TSK_ONCPU,
TSK_RUNNING, and TSK_MEMSTALL_RUNNING, setting TSK_IOWAIT, and
leaving sleep-persistent TSK_MEMSTALL set).

With Proxy Execution, mutex-blocked tasks stay on the runqueue
across context switches, so psi_sched_switch() runs with
sleep=false and only clears TSK_ONCPU. If the blocked donor task
is later deactivated via block_task(), psi_dequeue() is called
with DEQUEUE_SLEEP while TSK_ONCPU is already cleared.

Commit c2ae8b0df2d1 ("sched/core: Fix psi_dequeue() for Proxy
Execution") checked for (flags & DEQUEUE_SLEEP) && (p->psi_flags
& TSK_ONCPU), falling through to
psi_task_change(p, p->psi_flags, 0) when TSK_ONCPU was not set.
However, clearing all of p->psi_flags treats the sleeping
dequeue as a migration and wipes sleep-persistent flags like
TSK_MEMSTALL (and fails to set TSK_IOWAIT).

When the task later wakes up on the same cpu without
ENQUEUE_MIGRATED, psi_enqueue() assumes TSK_MEMSTALL is still
set from when the task went to sleep and only sets TSK_RUNNING |
TSK_MEMSTALL_RUNNING. This leaves things in an inconsistent
state (TSK_MEMSTALL_RUNNING set without TSK_MEMSTALL), causing
psi.mem.full to exceed psi.mem.some, delta underflows in
get_recent_times(), and spurious psi.mem.full events.

So fix psi_dequeue() such that when DEQUEUE_SLEEP is set and
TSK_ONCPU is not set, we do the same runnable-to-sleeping flag
transition as psi_task_switch() (clearing TSK_RUNNING and
TSK_MEMSTALL_RUNNING, setting TSK_IOWAIT, and preserving
TSK_MEMSTALL).

Fixes: c2ae8b0df2d1 ("sched/core: Fix psi_dequeue() for Proxy Execution")
Reported-by: David Stevens <stevensd@google.com>
Signed-off-by: John Stultz <jstultz@google.com>
---
Cc: Peter Zijlstra <peterz@infradead.org>
Cc: Ingo Molnar <mingo@redhat.com>
Cc: K Prateek Nayak <kprateek.nayak@amd.com>
Cc: Thomas Gleixner <tglx@linutronix.de>
Cc: Juri Lelli <juri.lelli@redhat.com>
Cc: Vincent Guittot <vincent.guittot@linaro.org>
Cc: Valentin Schneider <vschneid@redhat.com>
Cc: Suleiman Souhlal <suleiman@google.com>
Cc: Andrea Righi <arighi@nvidia.com>
Cc: Metin Kaya <Metin.Kaya@arm.com>
Cc: Dietmar Eggemann <dietmar.eggemann@arm.com>
Cc: Steven Rostedt <rostedt@goodmis.org>
Cc: Zimuzo Ezeozue <zezeozue@google.com>
Cc: Waiman Long <longman@redhat.com>
Cc: Boqun Feng <boqun.feng@gmail.com>
Cc: "Paul E. McKenney" <paulmck@kernel.org>
Cc: Joel Fernandes <joelagnelf@nvidia.com>
Cc: Qais Yousef <qyousef@layalina.io>
Cc: Xuewen Yan <xuewen.yan94@gmail.com>
Cc: kuyo chang <kuyo.chang@mediatek.com>
Cc: hupu <hupu.gm@gmail.com>
Cc: Zhan Xusheng <zhanxusheng@xiaomi.com>
Cc: zhidao su (Xiaomi) <soolaugust@gmail.com>
Cc: Hui Su <sh_def@163.com>
Cc: Shubhang Kaushik (Ampere) <sh@gentwo.org>
Cc: kernel-team@android.com
---
 kernel/sched/stats.h | 24 ++++++++++++++++++++----
 1 file changed, 20 insertions(+), 4 deletions(-)

diff --git a/kernel/sched/stats.h b/kernel/sched/stats.h
index ebe0a7765f98e..20820fe1021a1 100644
--- a/kernel/sched/stats.h
+++ b/kernel/sched/stats.h
@@ -175,12 +175,28 @@ static inline void psi_dequeue(struct task_struct *p, int flags)
 	 * Do nothing here.
 	 *
 	 * In the SCHED_PROXY_EXECUTION case we may do sleeping
-	 * dequeues that are not followed by a task switch, so check
-	 * TSK_ONCPU is set to ensure the task switch is imminent.
-	 * Otherwise clear the flags as usual.
+	 * dequeues that are not followed by a task switch. If TSK_ONCPU
+	 * is not set (meaning we already switched away while the task
+	 * stayed on the rq as a blocked donor), handle the sleeping
+	 * state transition here so sleep-persistent states (TSK_MEMSTALL,
+	 * TSK_IOWAIT) are preserved.
 	 */
-	if ((flags & DEQUEUE_SLEEP) && (p->psi_flags & TSK_ONCPU))
+	if (flags & DEQUEUE_SLEEP) {
+		int clear = TSK_RUNNING, set = 0;
+
+		if (p->psi_flags & TSK_ONCPU)
+			return;
+		if (p->se.sched_delayed)
+			return;
+
+		if (p->in_memstall)
+			clear |= TSK_MEMSTALL_RUNNING;
+		if (p->in_iowait)
+			set |= TSK_IOWAIT;
+
+		psi_task_change(p, clear, set);
 		return;
+	}
 
 	/*
 	 * When migrating a task to another CPU, clear all psi
-- 
2.56.0.257.gf7497c7d27-goog