From 493a9170800259fa92e14ca14b108cd13abef29c Mon Sep 17 00:00:00 2001
From: John Stultz <jstultz@google.com>
Date: Mon, 18 Dec 2023 14:23:11 -0800
Subject: [PATCH] ANDROID: sched: Refactor dl/rt find_lowest/latest_rq logic

This pulls re-validation logic done in find_lowest_rq
and find_latest_rq after re-acquiring the rq locks out into its
own function.

This allows us to later use a more complicated validation
check for chain-migration when using proxy-execution.

Cc: Joel Fernandes <joelaf@google.com>
Cc: Qais Yousef <qyousef@layalina.io>
Cc: Ingo Molnar <mingo@redhat.com>
Cc: Peter Zijlstra <peterz@infradead.org>
Cc: Juri Lelli <juri.lelli@redhat.com>
Cc: Vincent Guittot <vincent.guittot@linaro.org>
Cc: Dietmar Eggemann <dietmar.eggemann@arm.com>
Cc: Valentin Schneider <vschneid@redhat.com>
Cc: Steven Rostedt <rostedt@goodmis.org>
Cc: Ben Segall <bsegall@google.com>
Cc: Zimuzo Ezeozue <zezeozue@google.com>
Cc: Mel Gorman <mgorman@suse.de>
Cc: Will Deacon <will@kernel.org>
Cc: Waiman Long <longman@redhat.com>
Cc: Boqun Feng <boqun.feng@gmail.com>
Cc: "Paul E. McKenney" <paulmck@kernel.org>
Cc: Metin Kaya <Metin.Kaya@arm.com>
Cc: Xuewen Yan <xuewen.yan94@gmail.com>
Cc: K Prateek Nayak <kprateek.nayak@amd.com>
Cc: Thomas Gleixner <tglx@linutronix.de>
Cc: Daniel Lezcano <daniel.lezcano@linaro.org>
Cc: kernel-team@android.com
Change-Id: I6a1b60ab17659adec750d322133da5e11043b90c
Signed-off-by: John Stultz <jstultz@google.com>
Bug: 306081722
---
v8:
* Typo fixups suggested by Metin Kaya
---
 kernel/sched/deadline.c | 31 ++++++++++++++++++++-----
 kernel/sched/rt.c       | 50 ++++++++++++++++++++++++++++-------------
 2 files changed, 59 insertions(+), 22 deletions(-)

diff --git a/kernel/sched/deadline.c b/kernel/sched/deadline.c
index 0ad99fd124e7..ae1fb7bf8948 100644
--- a/kernel/sched/deadline.c
+++ b/kernel/sched/deadline.c
@@ -2614,6 +2614,30 @@ static int find_later_rq(struct task_struct *sched_ctx, struct task_struct *exec
 	return -1;
 }
 
+static inline bool dl_revalidate_rq_state(struct task_struct *task, struct rq *rq,
+					  struct rq *later)
+{
+	if (task_rq(task) != rq)
+		return false;
+
+	if (!cpumask_test_cpu(later->cpu, &task->cpus_mask))
+		return false;
+
+	if (task_on_cpu(rq, task))
+		return false;
+
+	if (!dl_task(task))
+		return false;
+
+	if (is_migration_disabled(task))
+		return false;
+
+	if (!task_on_rq_queued(task))
+		return false;
+
+	return true;
+}
+
 /* Locks the rq it finds */
 static struct rq *find_lock_later_rq(struct task_struct *task, struct rq *rq)
 {
@@ -2645,12 +2669,7 @@ static struct rq *find_lock_later_rq(struct task_struct *task, struct rq *rq)
 
 		/* Retry if something changed. */
 		if (double_lock_balance(rq, later_rq)) {
-			if (unlikely(task_rq(task) != rq ||
-				     !cpumask_test_cpu(later_rq->cpu, &task->cpus_mask) ||
-				     task_on_cpu(rq, task) ||
-				     !dl_task(task) ||
-				     is_migration_disabled(task) ||
-				     !task_on_rq_queued(task))) {
+			if (unlikely(!dl_revalidate_rq_state(task, rq, later_rq))) {
 				double_unlock_balance(rq, later_rq);
 				later_rq = NULL;
 				break;
diff --git a/kernel/sched/rt.c b/kernel/sched/rt.c
index 3d590e518e41..737e9401484f 100644
--- a/kernel/sched/rt.c
+++ b/kernel/sched/rt.c
@@ -2004,6 +2004,39 @@ static int find_lowest_rq(struct task_struct *sched_ctx, struct task_struct *exe
 	return -1;
 }
 
+static inline bool rt_revalidate_rq_state(struct task_struct *task, struct rq *rq,
+					  struct rq *lowest)
+{
+	/*
+	 * We had to unlock the run queue. In
+	 * the mean time, task could have
+	 * migrated already or had its affinity changed.
+	 * Also make sure that it wasn't scheduled on its rq.
+	 * It is possible the task was scheduled, set
+	 * "migrate_disabled" and then got preempted, so we must
+	 * check the task migration disable flag here too.
+	 */
+	if (task_rq(task) != rq)
+		return false;
+
+	if (!cpumask_test_cpu(lowest->cpu, &task->cpus_mask))
+		return false;
+
+	if (task_on_cpu(rq, task))
+		return false;
+
+	if (!rt_task(task))
+		return false;
+
+	if (is_migration_disabled(task))
+		return false;
+
+	if (!task_on_rq_queued(task))
+		return false;
+
+	return true;
+}
+
 /* Will lock the rq it finds */
 static struct rq *find_lock_lowest_rq(struct task_struct *task, struct rq *rq)
 {
@@ -2033,22 +2066,7 @@ static struct rq *find_lock_lowest_rq(struct task_struct *task, struct rq *rq)
 
 		/* if the prio of this runqueue changed, try again */
 		if (double_lock_balance(rq, lowest_rq)) {
-			/*
-			 * We had to unlock the run queue. In
-			 * the mean time, task could have
-			 * migrated already or had its affinity changed.
-			 * Also make sure that it wasn't scheduled on its rq.
-			 * It is possible the task was scheduled, set
-			 * "migrate_disabled" and then got preempted, so we must
-			 * check the task migration disable flag here too.
-			 */
-			if (unlikely(task_rq(task) != rq ||
-				     !cpumask_test_cpu(lowest_rq->cpu, &task->cpus_mask) ||
-				     task_on_cpu(rq, task) ||
-				     !rt_task(task) ||
-				     is_migration_disabled(task) ||
-				     !task_on_rq_queued(task))) {
-
+			if (unlikely(!rt_revalidate_rq_state(task, rq, lowest_rq))) {
 				double_unlock_balance(rq, lowest_rq);
 				lowest_rq = NULL;
 				break;