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

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

Hello John,

On 10/24/2026 3:34 AM, John Stultz wrote:
> @@ -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.
> +	 */

Aren't we under rq_lock with IRQs disabled here? If we
observe a p->sleeping_owner, we should be able to use it
since we are protected by RCU.

Maybe needs a load acquire but get + put may not be necessary.

> +	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;

I feel this might be racy and needs some ordering. Consider:

  p->__state = TASK_WAKING

  sched_ttwu_pending()
    ttwu_do_activate(p)
      /* None of these are ordered */           do_activate_blocked_waiter(p)
      WRITE_ONCE(p->__state, TASK_RUNNING);       if (!p->is_blocked) /* False */
      p->is_blocked = 0;                          if (READ_ONCE(p->__state) == TASK_WAKING) /* False */
                                                  if (READ_ONCE(p->on_rq)) /* False */
      WRITE_ONCE(p->on_rq, TASK_ON_RQ_QUEUED);    ... /* NMI */


      /* p runs,
       * blocks on new mutex,                       /* Only p->pi_lock is held */
       * queues on sleeping owner
       */
                                                  ... /* Return from NMI */

                                                  scoped_guard (raw_spinlock, &p->blocked_lock)
                                                    if (!task_is_blocked(p)) /* False */

                                                  ...
                                                  activate_task(target_rq, p, en_flags);
                                                    !! Activated while being queued on sleeping owner !!


Am I missing something that prevents this?

I think there are enough defensive mechanism for this to not
make a difference from correctness standpoint but it is worth
pointing out that this is technically possible if the loads /
stores are reordered and the machine meets some terrible fate.

Since task activation is now done under p->blocked_lock, maybe
these checks can be done under the p->blocked_lock too? It'll
inevitably stretch the critical section but 4 more loads on
a relatively hot cacheline shouldn't hurt much?
> +
> +		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;
> +		}
> +		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);
> +		}
> +		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);

This should be wakeup_preempt(rq, p, en_flags) instead?

> +		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

Insane jugglery but super impressed that you got everything right :-)

> +	}
> +}
> +
> +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);

Is this hunk necessary?

>  	}
>  
>  	/*
> @@ -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);
>  	if (success)
>  		ttwu_stat(p, task_cpu(p), wake_flags);
>  
> @@ -6822,21 +7079,6 @@ static bool try_to_block_task(struct rq *rq, struct task_struct *p,
>  }
>  
>  #ifdef CONFIG_SCHED_PROXY_EXEC
> -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 inline struct task_struct *proxy_resched_idle(struct rq *rq)
>  {
>  	put_prev_set_next_task(rq, rq->donor, rq->idle);
> @@ -6984,6 +7226,29 @@ static void proxy_migrate_task(struct rq *rq, struct rq_flags *rf,
>  	proxy_reacquire_rq_lock(rq, rf);
>  }
>  
> +static void proxy_enqueue_on_owner(struct rq *rq, struct task_struct *owner,
> +				   struct task_struct *p)
> +{
> +	lockdep_assert_rq_held(rq);
> +	lockdep_assert_held(&owner->blocked_lock);
> +	/*
> +	 * ttwu_activate() will pick them up and place them on whatever rq
> +	 * @owner will run next.
> +	 */
> +	WARN_ON(p == owner);
> +	WARN_ON(!p->on_rq);
> +	WARN_ON(p->sleeping_owner);
> +	get_task_struct(owner);
> +	WRITE_ONCE(p->sleeping_owner, owner);
> +	/*
> +	 * ttwu_do_activate must not have a chance to activate p
> +	 * elsewhere before it's fully extricated from its old rq.
> +	 */
> +	list_add(&p->blocked_node, &owner->blocked_head);
> +	proxy_resched_idle(rq);
> +	block_task(rq, p, READ_ONCE(p->__state));
> +}
> +
>  #define MAX_PROXY_CHAIN_DEPTH 1024
>  
>  /*
> @@ -7087,11 +7352,37 @@ find_proxy_task(struct rq *rq, struct task_struct *donor, struct rq_flags *rf)
>  		}
>  
>  		if (!READ_ONCE(owner->on_rq) || owner->se.sched_delayed) {
> -			/* XXX Don't handle blocked owners/delayed dequeue yet */
> +			/*
> +			 * rq->curr must not be added to the blocked_head list or else
> +			 * ttwu_do_activate could enqueue it elsewhere before it switches
> +			 * out here. The approach to avoid this is the same as in the
> +			 * migrate_task case.
> +			 */
>  			if (curr_in_chain)
>  				goto resched_idle;
> -			__clear_task_blocked_on(p, NULL);
> -			goto deactivate;
> +			/*
> +			 * If !@owner->on_rq, holding @rq->lock will not pin the task,
> +			 * so we cannot drop @mutex->wait_lock until we're sure its a blocked
> +			 * task on this rq.
> +			 *
> +			 * We use @owner->blocked_lock to serialize against ttwu_activate().
> +			 * Either we see its new owner->on_rq or it will see our list_add().
> +			 */

nit. Last two lines can be part on next comment to keep it
next to the p->on_rq check.

> +			WARN_ON(owner == p);
> +			raw_spin_unlock(&p->blocked_lock);
> +			raw_spin_lock(&owner->blocked_lock);
> +			/*
> +			 * Before actually adding to the sleeping owner, double check
> +			 * we didn't race with an owner wakeup before grabbing the
> +			 * owner's blocked_lock.
> +			 */
> +			if (!READ_ONCE(owner->on_rq) || owner->se.sched_delayed)
> +				proxy_enqueue_on_owner(rq, owner, p);
> +
> +			raw_spin_unlock(&owner->blocked_lock);
> +			raw_spin_lock(&p->blocked_lock);
> +
> +			return NULL; /* retry task selection */
>  		}
>  
>  		owner_cpu = task_cpu(owner);

-- 
Thanks and Regards,
Prateek