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

[Chengming Zhou <chengming.zhou@linux.dev>]

On 2024/12/26 19:04, K Prateek Nayak wrote:
> 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)

Ah, right :)

> 
>>
>>>      [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);

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).

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!