[PATCH sched_ext/for-7.4] sched_ext: Drop the locks periodically while walking a user DSQ for re-enqueue

[Tejun Heo <tj@kernel.org>]

reenq_user() drops dsq->lock around each re-enqueue and its rq lock every
SCX_TASK_ITER_BATCH re-enqueues. A walk filtered by SCX_REENQ_MARKED skips
unmarked tasks without dropping either, so a long run of unmarked tasks is
walked in one hold of both locks.

Count the visited tasks instead of the re-enqueued ones and drop both locks
every SCX_TASK_ITER_BATCH visits. An unfiltered walk drops the locks as
often as before.

Suggested-by: Andrea Righi <arighi@nvidia.com>
Link: https://lore.kernel.org/all/atP-LIKpxPhR-Nuc@gpd4/
Signed-off-by: Tejun Heo <tj@kernel.org>
---
 kernel/sched/ext/ext.c | 24 ++++++++++++++++--------
 1 file changed, 16 insertions(+), 8 deletions(-)

diff --git a/kernel/sched/ext/ext.c b/kernel/sched/ext/ext.c
index 14fe63428036..07857039aee6 100644
--- a/kernel/sched/ext/ext.c
+++ b/kernel/sched/ext/ext.c
@@ -5111,7 +5111,7 @@ static void reenq_user(struct rq *rq, struct scx_dispatch_q *dsq, u64 reenq_flag
 	struct scx_sched *sch = dsq->sched;
 	struct scx_dsq_list_node cursor = INIT_DSQ_LIST_CURSOR(cursor, dsq, 0);
 	struct task_struct *p;
-	s32 nr_enqueued = 0;
+	s32 nr_visited = 0;
 
 	lockdep_assert_rq_held(rq);
 
@@ -5121,6 +5121,21 @@ static void reenq_user(struct rq *rq, struct scx_dispatch_q *dsq, u64 reenq_flag
 		struct rq *task_rq;
 		u32 reason;
 
+		/*
+		 * Count skipped tasks too so that a long walk doesn't hold the
+		 * locks throughout. The cursor keeps the position.
+		 */
+		if (!(++nr_visited % SCX_TASK_ITER_BATCH)) {
+			raw_spin_unlock(&dsq->lock);
+			if (locked_rq) {
+				scx_rq_lock_drop(locked_rq);
+				raw_spin_rq_unlock(locked_rq);
+				locked_rq = NULL;
+			}
+			cpu_relax();
+			raw_spin_lock(&dsq->lock);
+		}
+
 		p = nldsq_cursor_next_task(&cursor, dsq);
 		if (!p)
 			break;
@@ -5157,13 +5172,6 @@ static void reenq_user(struct rq *rq, struct scx_dispatch_q *dsq, u64 reenq_flag
 
 		scx_do_enqueue_task(task_rq, p, SCX_ENQ_REENQ, -1);
 
-		if (!(++nr_enqueued % SCX_TASK_ITER_BATCH)) {
-			scx_rq_lock_drop(locked_rq);
-			raw_spin_rq_unlock(locked_rq);
-			locked_rq = NULL;
-			cpu_relax();
-		}
-
 		raw_spin_lock(&dsq->lock);
 	}
 
-- 
2.56.0