Re: [PATCH v33 7/7] sched/proxy: Add deactivated (sleeping) owner handling to find_proxy_task()

[Andrea Righi <arighi@nvidia.com>]

Hi John and Peter,

On Fri, Oct 23, 2026 at 10:04:18PM +0000, John Stultz wrote:
> From: Peter Zijlstra <peterz@infradead.org>
> 
> If the blocked_on chain resolves to a sleeping owner, deactivate
> the donor task, and enqueue it on the sleeping owner task.
> Then re-activate it later when the owner is woken up.
> 
> NOTE: This has been particularly challenging to get working
> properly, and some of the locking is particularly awkward. I'd
> very much appreciate review and feedback for ways to simplify
> this.
> 
> Signed-off-by: Peter Zijlstra (Intel) <peterz@infradead.org>
> Signed-off-by: Juri Lelli <juri.lelli@redhat.com>
> Signed-off-by: Valentin Schneider <valentin.schneider@arm.com>
> Signed-off-by: Connor O'Brien <connoro@google.com>
> [jstultz: This was broken out from the larger proxy() patch]
> Signed-off-by: John Stultz <jstultz@google.com>
> ---

...

> diff --git a/include/linux/sched.h b/include/linux/sched.h
> index 12a7e094a3e6b..3fae2f1438583 100644
> --- a/include/linux/sched.h
> +++ b/include/linux/sched.h
> @@ -1266,6 +1266,13 @@ struct task_struct {
>  
>  	struct mutex			*blocked_on;	/* lock we're blocked on */
>  	raw_spinlock_t			blocked_lock;
> +#ifdef CONFIG_SCHED_PROXY_EXEC
> +	struct list_head		blocked_head;  /* tasks blocked on this task */
> +	struct list_head		blocked_node;  /* our entry on someone elses blocked_head */
> +	/* Node for list of tasks to process blocked_head list for blocked entitiy activations */

nit: s/entitiy/entity/

> +	struct list_head		blocked_activation_node;
> +	struct task_struct		*sleeping_owner; /* task our blocked_node is enqueued on */
> +#endif
>  
>  	/*
>  	 * The task that is boosting this task; a back link for the current
> diff --git a/include/linux/sched/task.h b/include/linux/sched/task.h
> index 90ed5bc3c7afe..26c19ee980456 100644
> --- a/include/linux/sched/task.h
> +++ b/include/linux/sched/task.h
> @@ -160,6 +160,7 @@ static inline void put_task_struct(struct task_struct *t)
>  }
>  
>  DEFINE_FREE(put_task, struct task_struct *, if (_T) put_task_struct(_T))
> +DEFINE_GUARD(get_task, struct task_struct *, get_task_struct(_T), put_task_struct(_T))
>  
>  static inline void put_task_struct_many(struct task_struct *t, int nr)
>  {
> diff --git a/init/init_task.c b/init/init_task.c
> index ce7c2b07d855d..aed3f7d8a5cc3 100644
> --- a/init/init_task.c
> +++ b/init/init_task.c
> @@ -213,6 +213,12 @@ struct task_struct init_task __aligned(L1_CACHE_BYTES) = {
>  						 &init_task.alloc_lock),
>  #endif
>  	.blocked_donor = NULL,
> +#ifdef CONFIG_SCHED_PROXY_EXEC
> +	.blocked_head = LIST_HEAD_INIT(init_task.blocked_head),
> +	.blocked_node = LIST_HEAD_INIT(init_task.blocked_node),
> +	.blocked_activation_node = LIST_HEAD_INIT(init_task.blocked_activation_node),
> +	.sleeping_owner = NULL,
> +#endif
>  #ifdef CONFIG_RT_MUTEXES
>  	.pi_waiters	= RB_ROOT_CACHED,
>  	.pi_top_task	= NULL,
> diff --git a/kernel/fork.c b/kernel/fork.c
> index 394906afafd30..f87f806c2958b 100644
> --- a/kernel/fork.c
> +++ b/kernel/fork.c
> @@ -2277,6 +2277,12 @@ __latent_entropy struct task_struct *copy_process(
>  
>  	p->blocked_on = NULL; /* not blocked yet */
>  	p->blocked_donor = NULL; /* nobody is boosting p yet */
> +#ifdef CONFIG_SCHED_PROXY_EXEC
> +	INIT_LIST_HEAD(&p->blocked_head);
> +	INIT_LIST_HEAD(&p->blocked_node);
> +	INIT_LIST_HEAD(&p->blocked_activation_node);
> +	p->sleeping_owner = NULL;
> +#endif
>  
>  #ifdef CONFIG_BCACHE
>  	p->sequential_io	= 0;
> diff --git a/kernel/sched/core.c b/kernel/sched/core.c
> index 0cb41a884ea70..be3ce8b554fed 100644
> --- a/kernel/sched/core.c
> +++ b/kernel/sched/core.c
> @@ -2229,7 +2229,7 @@ inline bool dequeue_task(struct rq *rq, struct task_struct *p, int flags)
>  	return p->sched_class->dequeue_task(rq, p, flags);
>  }
>  
> -void activate_task(struct rq *rq, struct task_struct *p, int flags)
> +static inline void __activate_task(struct rq *rq, struct task_struct *p, int flags)
>  {
>  	if (task_on_rq_migrating(p))
>  		flags |= ENQUEUE_MIGRATED;
> @@ -2240,6 +2240,76 @@ void activate_task(struct rq *rq, struct task_struct *p, int flags)
>  	ASSERT_EXCLUSIVE_WRITER(p->on_rq);
>  }
>  
> +#ifdef CONFIG_SCHED_PROXY_EXEC
> +static inline
> +void __proxy_remove_from_sleeping_owner(struct task_struct *owner, struct task_struct *p)
> +{
> +	lockdep_assert_held(&owner->blocked_lock);
> +
> +	BUG_ON(p->sleeping_owner != owner);
> +
> +	list_del_init(&p->blocked_node);
> +	WRITE_ONCE(p->sleeping_owner, NULL);
> +	put_task_struct(owner); // matches get in proxy_enqueue_on_owner
> +}
> +
> +static inline void proxy_remove_from_sleeping_owner(struct task_struct *p)
> +{
> +	struct task_struct *owner;
> +
> +retry:
> +	owner = data_race(READ_ONCE(p->sleeping_owner));
> +	if (!owner)
> +		return;
> +
> +	/*
> +	 * __proxy_remove_from_sleeping_owner() does a
> +	 * put on owner to match the get done in
> +	 * proxy_enqueue_on_owner(). If that put is the
> +	 * last one and it frees owner, we'd be freeing
> +	 * a lock we held. So get/put owner around its
> +	 * usage her to ensure that doesn't happen.

nit: s/usage her/usage here/

> +	 */
> +	scoped_guard(get_task, owner) {
> +		guard(raw_spinlock)(&owner->blocked_lock);
> +
> +		if (owner != READ_ONCE(p->sleeping_owner))
> +			goto retry;
> +		__proxy_remove_from_sleeping_owner(owner, p);
> +	}
> +}
> +
> +void activate_task(struct rq *rq, struct task_struct *p, int en_flags)
> +{
> +	lockdep_assert_rq_held(rq);
> +
> +	if (!sched_proxy_exec()) {
> +		__activate_task(rq, p, en_flags);
> +		return;
> +	}
> +
> +	proxy_remove_from_sleeping_owner(p);
> +	/*
> +	 * By calling __activate_task() with blocked_lock held, we
> +	 * order against the find_proxy_task() blocked_task case
> +	 * such that no more blocked tasks will be enqueued on p
> +	 * once we release p->blocked_lock.
> +	 */
> +	guard(raw_spinlock)(&p->blocked_lock);
> +	WARN_ON(task_cpu(p) != cpu_of(rq));
> +	__activate_task(rq, p, en_flags);
> +}
> +#else
> +static inline void proxy_remove_from_sleeping_owner(struct task_struct *p)
> +{
> +}
> +
> +void activate_task(struct rq *rq, struct task_struct *p, int en_flags)
> +{
> +	__activate_task(rq, p, en_flags);
> +}
> +#endif
> +
>  void deactivate_task(struct rq *rq, struct task_struct *p, int flags)
>  {
>  	WARN_ON_ONCE(flags & DEQUEUE_SLEEP);
> @@ -3815,6 +3885,184 @@ static inline void proxy_reset_donor(struct rq *rq)
>  	resched_curr(rq);
>  }
>  
> +static inline void proxy_set_task_cpu(struct task_struct *p, int cpu)
> +{
> +	unsigned int wake_cpu;
> +
> +	/*
> +	 * Since we are enqueuing a blocked task on a cpu it may
> +	 * not be able to run on, preserve wake_cpu when we
> +	 * __set_task_cpu so we can return the task to where it
> +	 * was previously runnable.
> +	 */
> +	wake_cpu = p->wake_cpu;
> +	__set_task_cpu(p, cpu);
> +	p->wake_cpu = wake_cpu;
> +}
> +
> +static void do_activate_blocked_waiter(struct rq *target_rq, struct task_struct *p, int en_flags)
> +{
> +	struct rq_flags rf;
> +	int target_cpu = cpu_of(target_rq);
> +
> +	scoped_guard (raw_spinlock_irqsave, &p->pi_lock) {
> +		/* If task is no longer blocked, don't activate */
> +		if (!p->is_blocked)
> +			return;
> +
> +		/* Avoid racing with ttwu */
> +		if (READ_ONCE(p->__state) == TASK_WAKING)
> +			return;
> +
> +		if (READ_ONCE(p->on_rq)) {
> +			/*
> +			 * We raced with a non mutex handoff activation of p.
> +			 * That activation will also take care of activating
> +			 * all of the tasks after p in the blocked_head list,
> +			 * so we're done here.
> +			 */
> +			return;
> +		}

Is a racing activation the only possible case here? When
proxy_enqueue_on_owner() calls block_task() on a fair task that isn't
eligible, DELAY_DEQUEUE keeps it on_rq with se.sched_delayed set. In
that case p simply won't donate until it's woken by the mutex, so
nothing nasty shoudl happen, but maybe it's worth updating the comment
to describe this case as well?

> +		if (task_on_cpu(task_rq(p), p)) {
> +			/*
> +			 * Its possible this activation is very late, and
> +			 * we already were woken up and are running on a
> +			 * different cpu. If that task blocked, it could be
> +			 * dequeued (so on_rq == 0), but still on_cpu.
> +			 * Bail in this case, as we definitely don't want to
> +			 * activate a task when its on_cpu elsewhere.
> +			 */
> +			return;
> +		}
> +
> +		scoped_guard (raw_spinlock, &p->blocked_lock) {
> +			/*
> +			 * Again, if activation is delayed, it's possible the
> +			 * task has been woken and is no longer blocked_on,
> +			 * but went back to sleep, so double check blocked_on
> +			 * before activating onto the owner's rq
> +			 */
> +			if (!task_is_blocked(p))
> +				return;
> +		}
> +
> +		/*
> +		 * Have to make sure we handle nr_iowait adjustment before
> +		 * we call proxy_set_task_cpu() to ensure we are adjusting
> +		 * the same runqueue we left (where block_task()
> +		 * incremented nr_iowait).
> +		 */
> +		if (p->in_iowait) {
> +			delayacct_blkio_end(p);
> +			atomic_dec(&task_rq(p)->nr_iowait);
> +		}
> +
> +		if (task_cpu(p) != target_cpu) {
> +			en_flags |= ENQUEUE_MIGRATED;
> +			psi_ttwu_dequeue(p);
> +			proxy_set_task_cpu(p, target_cpu);
> +		}

Can this break SCHED_DEADLINE bw accounting?

p is sleeping here (not TASK_WAKING, not TASK_ON_RQ_MIGRATING), so
__set_task_cpu() skips migrate_task_rq_dl(): the old rq never does
sub_rq_bw(), while ENQUEUE_MIGRATED makes the new rq do add_rq_bw(), so
this_bw is leaked on the old rq. And if p is still dl_non_contending,
task_contending() on the new rq just cancels the inactive timer, leaving
running_bw on the old rq forever and making it underflow on the new rq
the next time p blocks.

set_task_cpu() doesn't help, because migrate_task_rq_dl() bails out
unless TASK_WAKING, so I'm not sure how to fix this case...

> +		rq_lock_irqsave(target_rq, &rf);
> +		/*
> +		 * proxy_enqueue_on_owner() called block_task() which
> +		 * increments nr_uninterruptible, so we need to reverse
> +		 * that when we activate the blocked waiter
> +		 */
> +		if (p->sched_contributes_to_load)
> +			target_rq->nr_uninterruptible--;
> +		update_rq_clock(target_rq);
> +		activate_task(target_rq, p, en_flags);
> +		resched_curr(target_rq);
> +		rq_unlock_irqrestore(target_rq, &rf);
> +	}
> +}
> +
> +static void activate_blocked_waiters(struct rq *target_rq,
> +				     struct task_struct *owner,
> +				     int wake_flags)
> +{
> +	struct list_head bal_head;
> +	unsigned long flags;
> +	int en_flags = ENQUEUE_WAKEUP | ENQUEUE_NOCLOCK;
> +
> +	if (!sched_proxy_exec())
> +		return;
> +
> +	/*
> +	 * A whole bunch of waiting donor tasks back this blocked
> +	 * lock owner task, wake them all up to give this task its
> +	 * 'fair' share.
> +	 *
> +	 * This is a little unique here and the locking is messy.
> +	 * At this point we only hold the blocked_lock, so the
> +	 * owner task may be able to run and do all sorts of
> +	 * things while we are processing the blocked_head list,
> +	 * including going back to sleep, which can cause tasks
> +	 * to be added to the owners->blocked_head while we are
> +	 * processing it!
> +	 * Thus, we pull the entire list off the owner->blocked_head
> +	 * here so that we will only process a finite amount of
> +	 * tasks. Tasks added after this will be processed by the
> +	 * future wake events.
> +	 * Even though we have pulled the list off the blocked_head
> +	 * the removed list is *still* "owned" and serialized by
> +	 * the owner->blocked_lock! As we have to serialize against
> +	 * mid-chain wakeups, who may try to remove themselves from
> +	 * the list.
> +	 */
> +	raw_spin_lock_irqsave(&owner->blocked_lock, flags);
> +	if (!list_empty(&owner->blocked_activation_node)) {
> +		raw_spin_unlock_irqrestore(&owner->blocked_lock, flags);
> +		return;
> +	}
> +
> +	if (list_empty(&owner->blocked_head)) {
> +		raw_spin_unlock_irqrestore(&owner->blocked_lock, flags);
> +		return;
> +	}
> +
> +	get_task_struct(owner);
> +	INIT_LIST_HEAD(&bal_head);
> +	list_add_tail(&owner->blocked_activation_node, &bal_head);
> +	raw_spin_unlock_irqrestore(&owner->blocked_lock, flags);
> +
> +	while (!list_empty(&bal_head)) {
> +		struct list_head tmp_head;
> +
> +		INIT_LIST_HEAD(&tmp_head);
> +		owner = list_first_entry(&bal_head, struct task_struct, blocked_activation_node);
> +
> +		raw_spin_lock_irqsave(&owner->blocked_lock, flags);
> +		list_replace_init(&owner->blocked_head, &tmp_head);
> +		list_del_init(&owner->blocked_activation_node);
> +		while (!list_empty(&tmp_head)) {
> +			struct task_struct *p;
> +
> +			p = list_first_entry(&tmp_head,
> +					     struct task_struct,
> +					     blocked_node);
> +			WARN_ON(p == owner);
> +			__proxy_remove_from_sleeping_owner(owner, p);
> +			raw_spin_unlock_irqrestore(&owner->blocked_lock, flags);
> +
> +			do_activate_blocked_waiter(target_rq, p, en_flags);
> +
> +			raw_spin_lock_irqsave(&p->blocked_lock, flags);
> +			if (list_empty(&p->blocked_activation_node)) {
> +				get_task_struct(p);
> +				list_add_tail(&p->blocked_activation_node, &bal_head);
> +			}
> +			raw_spin_unlock_irqrestore(&p->blocked_lock, flags);
> +
> +			raw_spin_lock_irqsave(&owner->blocked_lock, flags);
> +		}
> +		raw_spin_unlock_irqrestore(&owner->blocked_lock, flags);
> +		put_task_struct(owner); // put matches get prior to adding to local bal_head
> +	}
> +}
> +
> +static inline struct task_struct *proxy_resched_idle(struct rq *rq);
> +
>  /*
>   * Checks to see if task p has been proxy-migrated to another rq
>   * and needs to be returned. If so, we deactivate the task here
> @@ -3872,6 +4120,10 @@ static inline bool proxy_needs_return(struct rq *rq, struct task_struct *p)
>  {
>  	return false;
>  }
> +
> +static inline void activate_blocked_waiters(struct rq *target_rq,
> +					    struct task_struct *owner,
> +					    int wake_flags) {}
>  #endif /* CONFIG_SCHED_PROXY_EXEC */
>  
>  static void
> @@ -3945,8 +4197,10 @@ static int ttwu_runnable(struct task_struct *p, int wake_flags)
>  
>  	update_rq_clock(rq);
>  	if (p->is_blocked) {
> -		if (p->se.sched_delayed)
> +		if (p->se.sched_delayed) {
> +			proxy_remove_from_sleeping_owner(p);
>  			enqueue_task(rq, p, ENQUEUE_NOCLOCK | ENQUEUE_DELAYED);
> +		}
>  		if (proxy_needs_return(rq, p))
>  			return 0;
>  	}
> @@ -3975,13 +4229,15 @@ void sched_ttwu_pending(void *arg)
>  	update_rq_clock(rq);
>  
>  	llist_for_each_entry_safe(p, t, llist, wake_entry.llist) {
> +		int wake_flags;
>  		if (WARN_ON_ONCE(p->on_cpu))
>  			smp_cond_load_acquire(&p->on_cpu, !VAL);
>  
>  		if (WARN_ON_ONCE(task_cpu(p) != cpu_of(rq)))
>  			set_task_cpu(p, cpu_of(rq));
>  
> -		ttwu_do_activate(rq, p, p->sched_remote_wakeup ? WF_MIGRATED : 0, &rf);
> +		wake_flags = p->sched_remote_wakeup ? WF_MIGRATED : 0;
> +		ttwu_do_activate(rq, p, wake_flags, &rf);
>  	}
>  
>  	/*
> @@ -4488,6 +4744,7 @@ int try_to_wake_up(struct task_struct *p, unsigned int state, int wake_flags)
>  		ttwu_queue(p, cpu, wake_flags);
>  	}
>  out:
> +	activate_blocked_waiters(cpu_rq(task_cpu(p)), p, wake_flags);

We can get here also when ttwu_state_match() fails, e.g., a signal sent
to an owner sleeping in TASK_UNINTERRUPTIBLE. In that case the owner
never goes through select_task_rq() and task_cpu(owner) is the CPU where
it last slept, which may have gone offline in the meantime (hotplug
only migrates queued tasks I think).

Should this be done only if (success)?

Thanks,
-Andrea