Re: [PATCH] psi: Fix race when task wakes up before psi_sched_switch() adjusts flags

[K Prateek Nayak]

Hello there,

Thank you for taking a look at the patch!

On 12/26/2024 4:13 PM, Chengming Zhou wrote:
> Hi,
>
> On 2024/12/26 13:34, K Prateek Nayak wrote:
> > When running hackbench in a cgroup with bandwidth throttling enabled,
> > following PSI splat was observed:
> >
> >      psi: inconsistent task state! task=1831:hackbench cpu=8 psi_flags=14 clear=0 set=4
> >
> > When investigating the series of events leading up to the splat,
> > following sequence was observed:
> >      [008] d..2.: sched_switch: ... ==> next_comm=hackbench next_pid=1831 next_prio=120
> >          ...
> >      [008] dN.2.: dequeue_entity(task delayed): task=hackbench pid=1831 cfs_rq->throttled=0
> >      [008] dN.2.: pick_task_fair: check_cfs_rq_runtime() throttled cfs_rq on CPU8
> >      # CPU8 goes into newidle balance and releases the rq lock
> >          ...
> >      # CPU15 on same LLC Domain is trying to wakeup hackbench(pid=1831)
> >      [015] d..4.: psi_flags_change: psi: task state: task=1831:hackbench cpu=8 psi_flags=14 clear=0 set=4 final=14 # Splat (cfs_rq->throttled=1)
>
> I have a question here, why TSK_ONCPU is not set in psi_flags if
> the task hasn't arrived psi_sched_switch()?

It is set. "psi_flags" is in fact a hex value so the psi_flags is "0x14"
which is (TSK_ONCPU | TSK_RUNNING)

> >      [015] d..4.: sched_wakeup: comm=hackbench pid=1831 prio=120 target_cpu=008 # Task has woken on a throttled hierarchy
> >      [008] d..2.: sched_switch: prev_comm=hackbench prev_pid=1831 prev_prio=120 prev_state=S ==> ...
> >
> > psi_dequeue() relies on psi_sched_switch() to set the correct PSI flags
> > for the blocked entity, however, the following race is possible with
> > psi_enqueue() / psi_ttwu_dequeue() in the path from psi_dequeue() to
> > psi_sched_switch()
>
> Yeah, this race is introduced by delayed dequeue changes.
>
> In the past, a sleep task can't be migrated or enqueued before it's done in __schedule(). (finish_task(prev) clear prev->on_cpu.)

I see __block_task() doing:

    smp_store_release(&p->on_rq, 0);

wouldn't this allow the task to be migrated? P.S. I have not encountered a
case where psi_ttwu_dequeue() has occurred before a psi_sched_switch() but
looking at the code, I thought it might be possible (I might very well be
wrong)

> Now, ttwu_runnable() can call enqueue_task() on the delayed dequeue task
> to bring it schedulable.
>
> But migration is still impossible, since it's still running on this cpu,
> so no psi_ttwu_dequeue(), only psi_enqueue() can happen, right?
>
> (Actually, there we can enqueue_task() for any sleep task, including
> those are not delayed dequeue, if select_task_rq() returns same cpu
> as task_cpu(p) to optimize wakeup latency, maybe need to submit a patch
> later.)
>
> >      __schedule()
> >     rq_lock(rq)
> >         try_to_block_task(p)
> >         psi_dequeue()
> >         [ psi_task_switch() is responsible
> >           for adjusting the PSI flags ]
> >         put_prev_entity(&p->se)            try_to_wake_up(p)
> >         # no runnable task on rq->cfs            ...
> >         sched_balance_newidle()
> >         raw_spin_rq_unlock(rq)                __task_rq_lock(p)
> >         ...                        psi_enqueue()/psi_ttwu_dequeue() [Woops!]
> >                                 __task_rq_unlock(p)
> >         raw_spin_rq_lock(rq)
> >         ...
> >         [ p was re-enqueued or has migrated away ]
>
> Here ttwu_runnable() call enqueue_task() for delayed dequeue task.
>
> migration can't happen since p->on_cpu is still true.
>
> >         ...
> >         psi_task_switch() [Too late!]
> >     raw_spin_rq_unlock(rq)
> >
> > The wakeup context will see the flags for a running task when the flags
> > should have reflected the task being blocked. Similarly, a migration
> > context in the wakeup path can clear the flags that psi_sched_switch()
> > assumes will be set (TSK_ONCPU / TSK_RUNNING)
>
> In this ttwu_runnable() -> enqueue_task() case, I think psi_enqueue()
> should do nothing at all.
>
> Why? Because psi_dequeue() is deferred to psi_sched_switch(), so from
> PSI POV, this task hasn't gone sleep at all, so psi_enqueue() should NOT
> change any state too. (It's not a wakeup or migration from PSI POV.)

There I imagined that newidle_balance() can still pull a task that can
be selected before prev and with the current implementation where
calling try_to_block_task() would still mark it as blocked and the flags
would again be inconsistent.

> And the current code of "psi_sched_switch(prev, next, block);" looks
> buggy to me too! The "block" value is from try_to_block_task(), then
> pick_next_task() may drop and gain rq lock, so we can't use the stale
> value for psi_sched_switch().
>
> Before we used "task_on_rq_queued(prev)", now we have to also consider
> delayed dequeue case, so it should be:
>
> "!task_on_rq_queued(prev) || prev->se.sched_delayed"

Peter had suggested the current approach as opposed to that on:
https://lore.kernel.org/lkml/20241004123506.GR18071@xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx/
we can perhaps revisit that in light of this.

Again, lot of the observations in the cover letter are from auditing the
code itself and I might have missed something; any and all comments are
greatly appreciated.

--
Thanks and Regards,
Prateek

> Thanks!
>
> > Since the TSK_ONCPU flag has to be modified with the rq lock of
> > task_cpu() held, use a combination of task_cpu() and TSK_ONCPU checks to
> > prevent the race. Specifically:
> >
> > o psi_enqueue() will clear the TSK_ONCPU flag when it encounters one.
> >    psi_enqueue() will only be called with TSK_ONCPU set when the task is
> >    being requeued on the same CPU. If the task was migrated,
> >    psi_ttwu_dequeue() would have already cleared the PSI flags.
> >
> >    psi_enqueue() cannot guarantee that this same task will be picked
> >    again when the scheduling CPU returns from newidle balance which is
> >    why it clears the TSK_ONCPU to mimic a net result of sleep + wakeup
> >    without migration.
> >
> > o When psi_sched_switch() observes that prev's task_cpu() has changes or
> >    the TSK_ONCPU flag is not set, a wakeup has raced with the
> >    psi_sched_switch() trying to adjust the dequeue flag. If the next is
> >    same as the prev, psi_sched_switch() has to now set the TSK_ONCPU flag
> >    again. Otherwise, psi_enqueue() or psi_ttwu_dequeue() would have
> >    already adjusted the PSI flags and no further changes are required
> >    to prev's PSI flags.
> >
> > With the introduction of DELAY_DEQUEUE, the requeue path is considerably
> > shortened and with the addition of bandwidth throttling in the
> > __schedule() path, the race window is large enough to observed this
> > issue.
> >
> > Fixes: 4117cebf1a9f ("psi: Optimize task switch inside shared cgroups")
> > Signed-off-by: K Prateek Nayak <kprateek.nayak@xxxxxxx>
> > ---
> > This patch is based on tip:sched/core at commit af98d8a36a96
> > ("sched/fair: Fix CPU bandwidth limit bypass during CPU hotplug")
> >
> > Reproducer for the PSI splat:
> >
> >    mkdir /sys/fs/cgroup/test
> >    echo $$ > /sys/fs/cgroup/test/cgroup.procs
> >    # Ridiculous limit on SMP to throttle multiple rqs at once
> >    echo "50000 100000" > /sys/fs/cgroup/test/cpu.max
> >    perf bench sched messaging -t -p -l 100000 -g 16
> >
> > This worked reliably on my 3rd Generation EPYC System (2 x 64C/128T) but
> > also on a 32 vCPU VM.
> > ---
> > [..snip..]