sched_ext: Replace tryget_task_struct() with get_task_struct()

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 b7d4b28db7da ("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>
Signed-off-by: Tejun Heo <tj@kernel.org>

diff --git a/kernel/sched/ext.c b/kernel/sched/ext.c
index b8dd3358959d737fd015c20dac28cdc4779f14e9..64dfaf4dc5b2ffd847746d8eebc76bbcbe657de7 100644 (file)
--- 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;