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:35, Chengming Zhou wrote:
> 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!

Just made a quick fix and tested passed using your script.

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)) {
                 /* Wakeup of new or sleeping task */
                 if (p->in_iowait)
                         clear |= TSK_IOWAIT;