[PATCH sched_ext/for-7.2] sched_ext: Replace tryget_task_struct() with get_task_struct()

[Andrea Righi <arighi@nvidia.com>]

The tryget_task_struct() calls in scx_sub_disable(),
scx_root_enable_workfn() and scx_sub_enable_workfn() can never fail at
the points they're invoked:

 - scx_root_enable_workfn() iterates over scx_tasks under scx_tasks_lock
   and rq lock. sched_ext_dead() removes tasks from scx_tasks under the
   same scx_tasks_lock before put_task_struct_rcu_user() runs in
   finish_task_switch(). So any task observed in scx_tasks must have
   usage > 0; put_task_struct_rcu_user() hasn't been called and the
   delayed_put_task_struct() callback that decrements usage cannot have
   been queued.

 - scx_sub_disable() and scx_sub_enable_workfn() iterate via
   css_task_iter, which takes a reference on each task in
   css_task_iter_next() and holds it until the next iter_next() call, so
   usage > 0 is guaranteed by the iter itself.

The actual filter for dead tasks is the SCX_TASK_DEAD check inside
scx_task_iter_next_locked(), not tryget; tryget only fails on zero
usage, a state that can't be reached for tasks visible to these iters.

Commit 54d1429e7180 ("sched_ext: Use SCX_TASK_READY test instead of
tryget_task_struct() during class switch") removed an analogous tryget
in the class-switch loop. Convert the remaining tryget calls to plain
get_task_struct() and update the comment in scx_root_enable_workfn()
that suggested tasks could be observed with zero @usage waiting for an
RCU grace period.

Link: https://lore.kernel.org/all/agCLBxHEUqWIepx8@google.com
Suggested-by: Alice Ryhl <aliceryhl@google.com>
Signed-off-by: Andrea Righi <arighi@nvidia.com>
---
 kernel/sched/ext.c | 27 ++++++++++++---------------
 1 file changed, 12 insertions(+), 15 deletions(-)

diff --git a/kernel/sched/ext.c b/kernel/sched/ext.c
index b8dd3358959d7..64dfaf4dc5b2f 100644
--- a/kernel/sched/ext.c
+++ b/kernel/sched/ext.c
@@ -5933,14 +5933,11 @@ static void scx_sub_disable(struct scx_sched *sch)
 		WARN_ON_ONCE(!scx_task_on_sched(sch, p));
 
 		/*
-		 * If $p is about to be freed, nothing prevents $sch from
-		 * unloading before $p reaches sched_ext_free(). Disable and
-		 * exit $p right away.
+		 * @p is pinned by the iter: css_task_iter_next() takes a
+		 * reference and holds it until the next iter_next() call, so
+		 * @p->usage is guaranteed > 0.
 		 */
-		if (!tryget_task_struct(p)) {
-			scx_disable_and_exit_task(sch, p);
-			continue;
-		}
+		get_task_struct(p);
 
 		scx_task_iter_unlock(&sti);
 
@@ -7181,12 +7178,13 @@ static void scx_root_enable_workfn(struct kthread_work *work)
 	scx_task_iter_start(&sti, NULL);
 	while ((p = scx_task_iter_next_locked(&sti))) {
 		/*
-		 * @p may already be dead, have lost all its usages counts and
-		 * be waiting for RCU grace period before being freed. @p can't
-		 * be initialized for SCX in such cases and should be ignored.
+		 * @p is in scx_tasks under scx_tasks_lock, and SCX_TASK_DEAD
+		 * tasks are filtered by scx_task_iter_next_locked().
+		 * sched_ext_dead() removes @p from scx_tasks under the same
+		 * lock before put_task_struct_rcu_user() runs, so @p->usage
+		 * is guaranteed > 0 here.
 		 */
-		if (!tryget_task_struct(p))
-			continue;
+		get_task_struct(p);
 
 		/*
 		 * Set %INIT_BEGIN under the iter's rq lock so that a concurrent
@@ -7487,9 +7485,8 @@ static void scx_sub_enable_workfn(struct kthread_work *work)
 		if (p->scx.flags & SCX_TASK_SUB_INIT)
 			continue;
 
-		/* see scx_root_enable() */
-		if (!tryget_task_struct(p))
-			continue;
+		/* @p is pinned by the iter; see scx_sub_disable() */
+		get_task_struct(p);
 
 		if (!assert_task_ready_or_enabled(p)) {
 			ret = -EINVAL;
-- 
2.54.0