Re: [PATCH 12/14] sched: Add shared runqueue locking to __task_rq_lock()

[Peter Zijlstra <peterz@infradead.org>]

On Thu, Sep 11, 2025 at 02:19:57PM -1000, Tejun Heo wrote:
> Hello,
> 
> On Wed, Sep 10, 2025 at 05:44:21PM +0200, Peter Zijlstra wrote:
> > @@ -703,17 +703,24 @@ void double_rq_lock(struct rq *rq1, stru
> >  struct rq *__task_rq_lock(struct task_struct *p, struct rq_flags *rf)
> >  	__acquires(rq->lock)
> >  {
> > +	raw_spinlock_t *slock;
> >  	struct rq *rq;
> >  
> >  	lockdep_assert_held(&p->pi_lock);
> >  
> >  	for (;;) {
> >  		rq = task_rq(p);
> > +		slock = p->srq_lock;
> >  		raw_spin_rq_lock(rq);
> > -		if (likely(rq == task_rq(p) && !task_on_rq_migrating(p))) {
> > +		if (slock)
> > +			raw_spin_lock(slock);
> > +		if (likely(rq == task_rq(p) && !task_on_rq_migrating(p) &&
> > +			   (!slock || p->srq_lock == slock))) {
> >  			rq_pin_lock(rq, rf);
> >  			return rq;
> >  		}

Yeah, I think that needs to change a little. Perhaps something like:

	slock2 = p->srq_lock;
	if (... && (!slock2 || slock2 == slock))

> With the !slock condition, the following scenario is possible:
> 
>   __task_rq_lock()
>      slock = p->srq_lock; /* NULL */
>                                                 dispatch_enqueue()
>                                                   p->srq_lock = &dsq->lock;
>                                                 enqueue finishes
>      raw_spin_rq_lock(rq);
>      rq is the same, $slock is NULL, return
>   do something assuming p is locked down        p gets dispatched to another rq
> 
> I'm unclear on when p->srq_lock would be safe to set and clear, so the goal
> is that whoever does [__]task_rq_lock() ends up waiting on the dsq lock that
> the task is queued on, and if we can exclude other sched operations that
> way, we don't have to hold source rq lock when moving the task to another rq
> for execution, right?

Indeed. If !p->srq_lock then task_rq(p)->lock must be sufficient.

So for enqueue, which sets p->srq_lock, this must be done while holding
task_rq(p)->lock.

So the above example should be serialized on task_rq(p)->lock, since
__task_rq_lock() holds it, enqueue cannot happen. Conversely, if enqueue
holds task_rq(p)->lock, then __task_rq_lock() will have to wait for
that, and then observe the newly set p->srq_lock and cycle to take that.

> In the last patch, it's set on dispatch_enqueue() and cleared when the task
> leaves the DSQ. Let's consider a simple scenario where a task gets enqueued,
> gets put on a non-local DSQ and then dispatched to a local DSQ, Assuming
> everything works out and we don't have to lock the source rq for migration,
> we'd be depending on task_rq_lock() reliably hitting p->srq_lock to avoid
> races, but I'm not sure how this would work. Let's say p is currently
> associated with CPU1 on a non-local DSQ w/ p->srq_lock set to its source
> DSQ.
> 
>   pick_task_ext() on CPU0               task property change on CPU1
>     locks the DSQ
>     picks p      
>     task_unlink_from_dsq()              task_rq_lock();
>       p->srq_lock = NULL;                 lock rq on CPU1
>     p is moved to local DSQ               sees p->src_lock == NULL
>                                           return
>   p starts running
>   anything can happen
>                                         proceed with property change

Hmm, the thinking was that if !p->srq_lock then task_rq(p)->lock should
be sufficient.

We must do set_task_cpu(0) before task_unlink_from_dsq() (and I got this
order wrong in yesterday's email).

  pick_task_ext() on CPU0		
    lock DSQ
    pick p
    set_task_cpu(0)			task_rq_lock()
    task_unlink_from_dsq()		  if !p->srq_lock, then task_rq(p) == 0
      p->srq_lock = NULL;
    p is moved to local DSQ

Perhaps the p->srq_lock store should be store-release, so that the cpu
store is before.

Then if we observe p->srq_lock, we'll serialize against DSQ and all is
well, if we observe !p->srq_lock then we must also observe task_rq(p) ==
0 and then we'll serialize on rq->lock.


Now let me see if there isn't an ABA issue here, consider:

pre: task_cpu(p) != 2, p->srq_lock = NULL

  CPU0				CPU1				CPU2

  __task_rq_lock()		enqueue_task_scx() 		pick_task_scx()

				rq = task_rq(p);
				LOCK rq->lock
  rq = task_rq(p)
  LOCK rq->lock
    .. waits
				LOCK dsq->lock
				enqueue on dsq
				p->srq_lock = &dsq->lock	
				UNLOCK dsq->lock
								LOCK dsq->lock
								pick p
				UNLOCK rq->lock
								set_task_cpu(2)
								task_unlink_from_dsq()
								  p->srq_lock = NULL;
								UNLOCK dsq->lock
    .. resumes
 
At this point our CPU0's __task_rq_lock():

  - if it observes p->srq_lock, it will cycle taking that, only to then
    find out p->srq_lock is no longer set, but then it must also see
    task_rq() has changed, so the next cycle will block on CPU2's
    rq->lock.

  - if it observes !p->srq_lock, then it cannot be the initial NULL,
    since the initial task_rq(p)->lock ordering prohibits this. So it
    must be the second NULL, which then also mandates we see the CPU
    change and we'll cycle to take CPU2's rq->lock.

That is, I _think_ we're okay :-)