Re: [RFC][PATCH v14 5/7] sched: Add an initial sketch of the find_proxy_task() function

Peter Zijlstra
2024-12-14T00:06:08Z

On Mon, Nov 25, 2024 at 11:51:59AM -0800, John Stultz wrote:

> diff --git a/kernel/sched/core.c b/kernel/sched/core.c
> index f8714050b6d0d..b492506d33415 100644
> --- a/kernel/sched/core.c
> +++ b/kernel/sched/core.c
> @@ -5052,6 +5052,34 @@ static void do_balance_callbacks(struct rq *rq, struct balance_callback *head)
> }
> }
>
> +/*
> + * Only called from __schedule context
> + *
> + * There are some cases where we are going to re-do the action
> + * that added the balance callbacks. We may not be in a state
> + * where we can run them, so just zap them so they can be
> + * properly re-added on the next time around. This is similar
> + * handling to running the callbacks, except we just don't call
> + * them.
> + */

Which specific callbacks are this? sched_core_balance()?

In general, shooting down all callbacks like this makes me feel somewhat
uncomfortable.

> +#ifdef CONFIG_SCHED_PROXY_EXEC
> +
> +static inline struct task_struct *
> +proxy_resched_idle(struct rq *rq)
> +{
> + put_prev_task(rq, rq->donor);
> + rq_set_donor(rq, rq->idle);
> + set_next_task(rq, rq->idle);
> + set_tsk_need_resched(rq->idle);
> + return rq->idle;
> +}
> +
> +static bool proxy_deactivate(struct rq *rq, struct task_struct *donor)
> +{
> + unsigned long state = READ_ONCE(donor->__state);
> +
> + /* Don't deactivate if the state has been changed to TASK_RUNNING */
> + if (state == TASK_RUNNING)
> + return false;
> + /*
> + * Because we got donor from pick_next_task, it is *crucial*
> + * that we call proxy_resched_idle before we deactivate it.
> + * As once we deactivate donor, donor->on_rq is set to zero,
> + * which allows ttwu to immediately try to wake the task on
> + * another rq. So we cannot use *any* references to donor
> + * after that point. So things like cfs_rq->curr or rq->donor
> + * need to be changed from next *before* we deactivate.
> + */
> + proxy_resched_idle(rq);
> + return try_to_block_task(rq, donor, state, true);
> +}
> +
> +/*
> + * Initial simple proxy that just returns the task if it's waking
> + * or deactivates the blocked task so we can pick something that
> + * isn't blocked.
> + */
> +static struct task_struct *
> +find_proxy_task(struct rq *rq, struct task_struct *donor, struct rq_flags *rf)
> +{
> + struct task_struct *p = donor;
> + struct mutex *mutex;
> +
> + mutex = p->blocked_on;
> + /* Something changed in the chain, so pick again */
> + if (!mutex)
> + return NULL;
> + /*
> + * By taking mutex->wait_lock we hold off concurrent mutex_unlock()
> + * and ensure @owner sticks around.
> + */
> + raw_spin_lock(&mutex->wait_lock);
> + raw_spin_lock(&p->blocked_lock);

I'm still wondering what this blocked_lock does, that previous patch had
it mirror wait_mutex too, so far I don't see the point.

> +
> + /* Check again that p is blocked with blocked_lock held */
> + if (!task_is_blocked(p) || mutex != get_task_blocked_on(p)) {
> + /*
> + * Something changed in the blocked_on chain and
> + * we don't know if only at this level. So, let's
> + * just bail out completely and let __schedule
> + * figure things out (pick_again loop).
> + */
> + goto out;
> + }
> + if (!proxy_deactivate(rq, donor))
> + /* XXX: This hack won't work when we get to migrations */
> + donor->blocked_on_state = BO_RUNNABLE;
> +
> +out:
> + raw_spin_unlock(&p->blocked_lock);
> + raw_spin_unlock(&mutex->wait_lock);
> + return NULL;
> +}
> +#else /* SCHED_PROXY_EXEC */
> +static struct task_struct *
> +find_proxy_task(struct rq *rq, struct task_struct *donor, struct rq_flags *rf)
> +{
> + WARN_ONCE(1, "This should never be called in the !SCHED_PROXY_EXEC case\n");
> + return donor;
> +}
> +#endif /* SCHED_PROXY_EXEC */
> +
> /*
> * __schedule() is the main scheduler function.
> *
> @@ -6732,12 +6845,22 @@ static void __sched notrace __schedule(int sched_mode)
> goto picked;
> }
> } else if (!preempt && prev_state) {
> - block = try_to_block_task(rq, prev, prev_state);
> + block = try_to_block_task(rq, prev, prev_state,
> + !task_is_blocked(prev));
> switch_count = &prev->nvcsw;
> }
>
> - next = pick_next_task(rq, prev, &rf);
> +pick_again:
> + next = pick_next_task(rq, rq->donor, &rf);
> rq_set_donor(rq, next);
> + if (unlikely(task_is_blocked(next))) {
> + next = find_proxy_task(rq, next, &rf);
> + if (!next) {
> + /* zap the balance_callbacks before picking again */
> + zap_balance_callbacks(rq);
> + goto pick_again;
> + }
> + }
> picked:
> clear_tsk_need_resched(prev);
> clear_preempt_need_resched();