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

[K Prateek Nayak <kprateek.nayak@amd.com>]

Hello there,

On 12/26/2024 9:12 PM, Chengming Zhou wrote:
> [..snip..]
>>>>>
>>>>> 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);
>>
>> Right, p->on_rq is cleared if not delayed dequeue.
>>
>>>
>>> wouldn't this allow the task to be migrated? P.S. I have not encountered a
>>
>> But p->on_cpu hasn't been cleared until finish_task(prev).

I forgot those are two different flags! Thank you for clarifying :)

>>
>> We can see in `can_migrate_task()`, we can't migrate `task_on_cpu()`
>> task. Anyway, its code still running on the cpu, we can't migrate it
>> to another cpu and run its code concurrently.
>>
>>> 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.
>>
>> Yes, it's blocked by try_to_block_task(), just ignored (deferred) by
>> PSI. During the time before `psi_sched_switch()`, it maybe enqueued
>> by ttwu(), which should be ignored by PSI too.
>>
>> At last `psi_sched_switch()` called with "sleep == false", just don't
>> notice this transient dequeue & enqueue operations. And the flags are
>> still consistent.
>>
>>>
>>>>
>>>> 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@noisy.programming.kicks-ass.net/
>>> we can perhaps revisit that in light of this.
>>
>> Ok, this "block" value is stale when `psi_sched_switch()` called.
>>
>> So we would see these inconsistent psi flags changes.
>>
>>>
>>> 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!
> 
> Just made a quick fix and tested passed using your script.

Thank you! The diff seems to be malformed as a result of whitespaces but
I was able to test if by recreating the diff. Feel free to add:

Reported-by: K Prateek Nayak <kprateek.nayak@amd.com>
Closes: https://lore.kernel.org/lkml/20241226053441.1110-1-kprateek.nayak@amd.com/
Tested-by: K Prateek Nayak <kprateek.nayak@amd.com>

If you can give your sign off, I could add a commit message and send it on
your behalf too.

> 
> diff --git a/kernel/sched/core.c b/kernel/sched/core.c
> index 3e5a6bf587f9..065ac76c47f9 100644
> --- a/kernel/sched/core.c
> +++ b/kernel/sched/core.c
> @@ -6641,7 +6641,6 @@ static void __sched notrace __schedule(int sched_mode)
>           * as a preemption by schedule_debug() and RCU.
>           */
>          bool preempt = sched_mode > SM_NONE;
> -       bool block = false;
>          unsigned long *switch_count;
>          unsigned long prev_state;
>          struct rq_flags rf;
> @@ -6702,7 +6701,7 @@ static void __sched notrace __schedule(int sched_mode)
>                          goto picked;
>                  }
>          } else if (!preempt && prev_state) {
> -               block = try_to_block_task(rq, prev, prev_state);
> +               try_to_block_task(rq, prev, prev_state);
>                  switch_count = &prev->nvcsw;
>          }
> 
> @@ -6748,7 +6747,8 @@ static void __sched notrace __schedule(int sched_mode)
> 
>                  migrate_disable_switch(rq, prev);
>                  psi_account_irqtime(rq, prev, next);
> -               psi_sched_switch(prev, next, block);
> +               psi_sched_switch(prev, next, !task_on_rq_queued(prev) ||
> +                                               prev->se.sched_delayed);
> 
>                  trace_sched_switch(preempt, prev, next, prev_state);
> 
> diff --git a/kernel/sched/stats.h b/kernel/sched/stats.h
> index 8ee0add5a48a..65efe45fcc77 100644
> --- a/kernel/sched/stats.h
> +++ b/kernel/sched/stats.h
> @@ -150,7 +150,7 @@ static inline void psi_enqueue(struct task_struct *p, int flags)
>                  set = TSK_RUNNING;
>                  if (p->in_memstall)
>                          set |= TSK_MEMSTALL | TSK_MEMSTALL_RUNNING;
> -       } else {
> +       } else if (!task_on_cpu(task_rq(p), p)) {

One small nit. here

If the task is on CPU at this point, both set and clear are 0 but
psi_task_change() is still called and I don't see it bailing out if it
doesn't have to adjust any flags.

Can we instead just do an early return if task_on_cpu(task_rq(p), p)
returns true? I've tested that version too and I haven't seen any
splats.

>                  /* Wakeup of new or sleeping task */
>                  if (p->in_iowait)
>                          clear |= TSK_IOWAIT;

-- 
Thanks and Regards,
Prateek