[RFC PATCH 00/16][PoC] sched/core: Alternate approach to sleeping-owner handling in PROXY_EXEC

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

Hello folks,

I had promised John I would share an insane idea if I got it
working(ish) so this RFC presents an alternate approach to handle
enqueued donor wakeup vs owner's chain-wakeup race via ttwu_runnable()
path and make the activate path as light as possible for tasks that
don't need to do a chain wakeup.

The series essentially breaks down John's large patch at
https://lore.kernel.org/lkml/20260807035232.1881495-9-jstultz@google.com/
into smaller chunks while introducing the new approach and fixing a few
snags encountered along the way (Patch 1 - Patch 4 are fixes that can be
discussed without getting into the rest of the RFC).

Disclaimer: Series in not bisectible in any way at the moment - related
bits are introduced together at and intermediate builds may fail.

Introduction
============

Handling sleeping owner requires proxy-chains to queue on the said owner
and perform a chain activation when the blocked owner wakes up.

Since a blocked donor can be woken up by another concurrent wake event,
the activation path becomes complicated and wakeup has to take an extra
lock (p->blocked_lock) to prevent any modifications to "p->blocked_head"
and miss activating any blocked donors.

The nesting rules for the locks are as follows:

  p->pi_lock
    __task_rq_lock()
      lock->wait_lock
        p->blocked_lock

Current chain-wakeup scheme juggles p->pi_lock, and p->blocked_lock to
prevent concurrent wake events from modifying the chain.

This RFC introduces a __task_rq_lock() based scheme that requires just
the __task_rq_lock() (always held) + p->blocked_lock (juggled for each
owner in chain) when queued tasks are detected to minimize the lock
bouncing and introduces a no-extra-locking fast-path for activation.

Trade offs are discussed at the end.

Approach
========

To implement this, we try a bunch of bad ideas:

o Waking up blocked donors linked to sleeping owner via ttwu_runnable()

To prevent chain-wakeup transitioning "p->on_rq" to QUEUED while a
concurrent wakeup is in progress, "p->on_rq" is unionized with a new
state "p->is_linked" as:

    union {
        struct {
            u8    on_rq;	/* Task on CPU */
            u8    is_linked;	/* Task on sleeping owner */
        };
        u16       needs_rq_sync;
    }

Instead of checking READ_ONCE(p->on_rq), to send task down the
ttwu_runnable() route, the new "needs_rq_sync" indicator is used which
can be atomically inspected outside rq_lock.

In ttwu_runnable() a fast-path first tries to remove the blocked donor
from the sleeping owner by just grabbing the "owner->blocked_lock" and
forces a full __task_rq_lock() if the blocked donor's state has changed
while trying to grab the blocked lock.

The chain-wakeup path takes a single __task_rq_lock() to prevent any
race from concurrent wakeups (described later).


o Proactively removing the blocked donor from owner if it wakes up

Proactively remove the blocked donor from sleeping owner in
proxy_enqueue_on_owner() if owner->on_rq is transitioned to != 0. This
ensures owner need not care about any additional donors that come in
after it transitions p->on_rq.

proxy_enqueue_on_owner() holds lock->wait_lock + owner->blocked_lock.
Since owner must grab the wait_lock again to wake up the blocked donor,
it is guaranteed that owner cannot disappear and get_task_struct() /
put_task_struct() dance is avoided.


o Migrating enqueued donors while blocking (SLEEP | MIGRATING)

To ensure a single __task_rq_lock() can handle the entire chain, a
blocked donor is migrated to a deterministic CPU  before it is
transitioned to p->on_rq = 0. This ensures the whole chain is owned by a
single CPU.


o No delayed tasks on chain

Delayed tasks pose a challenge that they are still on_rq and require
special handling to ensure the above condition of chain resolving to a
single CPU holds true. For simplicity, all delayed tasks on the chain
are fully blocked before they are queued.

This can be handled differently but I got a bit busy lately so that will
be handled in RFC v2 (if folks think this approach holds water).


o Single shot chain-wakeup (ab)using TASK_ON_RQ_MIGRATING

Tasks on chain are transitioned from being blocked to
TASK_ON_RQ_MIGRATING under __task_rq_lock(). Tasks on ttwu_runnable()
path that are trying to grab __task_rq_lock() is stalled until the
task is woken up on destination rq and the rq_lock() there is dropped.


o Tracking task that can potentially resolve to __mutex_owner()

This indicator is used to selectively set the lock holding
tasks to the slow-path.


Interleaving
============

How are different races handled (individual patches have these
highlighted too, either in comment message, or via code comments):

o Owner wakeup vs donor enqueue

  activate_task(owner)
    owner->on_rq = MIGRATING			proxy_enqueue_on_owner()
						  guard(raw_spinlock)(&owner->blocked_lock)
						  list_add(donor, owner->blocked_head)

    smp_mb() /* MIGRATING is visible */		  smp_mb() /* List addition is visible */

    if (list_empty(&owner->blocked_head))	  if (!READ_ONCE(owner->on_rq))
      return __activate_task(owner)		    block_task(donor)

    ... /* Do chain-activation */		  __proxy_dequeue_from_owner(donor)


o Chain-wakeup vs donor wakeup

  proxy_activate_donor_task()			try_to_wake_up(donor)

    rq_lock(rq_of(onwner->blocked_cpu))		  proxy_try_dequeue_from_owner(p)
      guard(raw_spinlock)(&owner->blocked_lock)	    guard(raw_spinlock)(&owner->blocked_lock)

      for_each(donor, owner->blocked_head)
        donor->on_rq = MIGRATING;
        smp_mb();				    if (donor->on_rq)
						      goto slowpath;
        ... /* More stuff */
    rq_unlock(rq_of(onwner->blocked_cpu));	    __task_rq_lock(donor)
						      if (donor->on_rq == MIGRATING)
						        cpu_relax();

    __proxy_dequeue_from_owner(donor)
    __activate_task(donor)
      p->on_rq = QUEUED;			    /* Acquired */
						    if (p->is_linked)
						      proxy_dequeue_from_owner(donor)

						    /* Do wakeup. */

Chain wakeup is covered extensively in Patch 14.


Series breakdown
================

Patch 01-04: Fixes to existing bits (Patch 4 needs more thinking)
Patch 05-07: Donor queuing with some convenience added to skip delayed
             handling later on.
Patch 08-13: Putting tasks on ttwu_running() path and swapping
             task_cpu() while blocking
Patch 14   : Chain wake-up
Patch 15-16: Optimizations to limit slowpath.


Trade-off
=========

Advantages:

o Only need to juggle blocked_lock(s) under a single rq_lock().
o Re-use ttwu_runnable bit to handle removal of proxy donor.
o No need to do get_task_stuct() / put_task_struct() juggling.

Disadvantages:

o Possible increased rq_lock contention on the chain-wakeup path but
  those events are generally rare.
o Lack of delayed task handling since the src_rq need to be locked
  separately to migrate it.

For the delayed handling, it is possible to use proxy_migrate_task() to
move the task after blocking. Series does not implement this yet to
limit the number of bad ideas.


How to apply
============

This series is based on John's tree at 

  https://github.com/johnstultz-work/linux-dev.git proxy-exec-v31-7.2-rc4

at commit 18bea15ce9fa ("sched: Migrate whole chain in
proxy_migrate_task()"). For convenience, a tree is present at:

  https://github.com/kudureranganath/linux.git  sched/proxy/sleeping_owner_rfc_v1

Sorry in advance to folks whose VPNs hate dealing with Github but I'm
told pulling it locally and inspecting the patches on tree works better.


Testing
=======

Tested with perf bench sched messaging, test-ww_mutex, and a kernel
module that implements a duty-cycle with sleep_ms() + busy-loop in
kernel.

I made sure to build and test !SCHED_PROXY_EXEC version this time around
to save John some trouble :-)


Acknowledgement
===============

Most of the good work is by John and all the bad ideas are by me.
Sorry in advance.

---
K Prateek Nayak (16):
  sched/core: Break activation of blocked task into a separate helper
  sched/core: Use enqueue/dequeue flags instead of
    task_on_rq_migrating()
  sched/fair: Use enqueue flags for DO_ATTACH in update_load_avg()
  sched/core: Activate blocked donor when no owner is found
  sched/core: Do not queue blocked donor on a delayed owner
  sched/core: Queue blocked donor onto sleeping owner for chain-wakeup
  sched/core: Avoid delaying blocked donors queued on sleeping owner
  sched/deadline: Prepare for blocking and proxy activation with
    MIGRATING flag
  sched/core: Track CPU where the task was blocked on
  sched/core: Introduce p->is_linked to track if task is queued on
    sleeping owner
  sched/core: Prepare to inspect ->is_linked alongside ->on_rq during
    wakeup
  sched:core: Add MIGRATING flags when blocking and activating linked
    donors
  sched/core: Use p->is_linked state to unlink from sleeping owner early
  sched/core: Introduce chain-wakeup to activate blocked donors
  locking/mutex: Track locks owned by a task in a per-task counter
  sched/core: Set activation of non lock-holders to fast-path

 include/linux/sched.h   |  31 ++-
 init/init_task.c        |   5 +
 kernel/fork.c           |   5 +
 kernel/locking/mutex.c  |  57 +++-
 kernel/sched/core.c     | 580 ++++++++++++++++++++++++++++++++++++++--
 kernel/sched/deadline.c |  15 +-
 kernel/sched/fair.c     |  16 +-
 kernel/sched/sched.h    |   6 +
 kernel/sched/stats.h    |   2 +-
 9 files changed, 666 insertions(+), 51 deletions(-)


base-commit: 18bea15ce9faaa0c1bd827012e6c60dbf35578c7
-- 
2.34.1