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

Peter Zijlstra
2025-09-12T11:55:12Z

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