Re: [PATCH v2] sched: set TIF_NEED_RESCHED before calling __trace_set_need_resched()

[Peter Zijlstra <peterz@infradead.org>]

On Wed, Jul 01, 2026 at 12:21:14PM +0530, K Prateek Nayak wrote:

> diff --git a/kernel/sched/core.c b/kernel/sched/core.c
> index 7cbd541f656f..bd2f7fb87dc9 100644
> --- a/kernel/sched/core.c
> +++ b/kernel/sched/core.c
> @@ -1049,9 +1049,16 @@ static inline void hrtick_schedule_exit(struct rq *rq) { }
>   * this avoids any races wrt polling state changes and thereby avoids
>   * spurious IPIs.
>   */
> -static inline bool set_nr_and_not_polling(struct thread_info *ti, int tif)
> +static inline bool set_nr_and_not_polling(struct rq *rq, int tif)
>  {
> -	return !(fetch_or(&ti->flags, 1 << tif) & _TIF_POLLING_NRFLAG);
> +	struct task_struct *curr = rq->curr;
> +	struct thread_info *ti = task_thread_info(curr);
> +	unsigned long old_flags = fetch_or(&ti->flags, 1 << tif);
> +
> +	if (trace_sched_set_need_resched_tp_enabled() && !(old_flags & (1 << tif)))
> +		trace_call__sched_set_need_resched_tp(curr, cpu_of(rq), tif);
> +
> +	return !(old_flags & _TIF_POLLING_NRFLAG);
>  }
>  
>  /*
> @@ -1072,13 +1079,20 @@ static bool set_nr_if_polling(struct task_struct *p)
>  			return true;
>  	} while (!try_cmpxchg(&ti->flags, &val, val | _TIF_NEED_RESCHED));
>  
> +	trace_sched_set_need_resched_tp(p, task_cpu(p), TIF_NEED_RESCHED);
>  	return true;
>  }
>  
>  #else
> -static inline bool set_nr_and_not_polling(struct thread_info *ti, int tif)
> +static inline bool set_nr_and_not_polling(struct rq *rq, int tif)
>  {
> -	set_ti_thread_flag(ti, tif);
> +	struct task_struct *curr = rq->curr;
> +	struct thread_info *ti = task_thread_info(curr);
> +	int set = test_and_set_ti_thread_flag(ti, tif);
> +
> +	if (trace_sched_set_need_resched_tp_enabled() && !set)
> +		trace_call__sched_set_need_resched_tp(curr, cpu_of(rq), tif);
> +
>  	return true;
>  }
>  
> @@ -1186,7 +1200,6 @@ static void __resched_curr(struct rq *rq, int tif)
>  {
>  	struct task_struct *curr = rq->curr;
>  	struct thread_info *cti = task_thread_info(curr);
> -	bool need_ipi;
>  	int cpu;
>  
>  	lockdep_assert_rq_held(rq);
> @@ -1204,16 +1217,16 @@ static void __resched_curr(struct rq *rq, int tif)
>  	cpu = cpu_of(rq);
>  
>  	if (cpu == smp_processor_id()) {
> -		set_ti_thread_flag(cti, tif);
> +		int set = test_and_set_ti_thread_flag(cti, tif);
> +
> +		if (trace_sched_set_need_resched_tp_enabled() && !set)
> +			trace_call__sched_set_need_resched_tp(curr, cpu, tif);
>  		if (tif == TIF_NEED_RESCHED)
>  			set_preempt_need_resched();
> -		trace_sched_set_need_resched_tp(curr, cpu, tif);
>  		return;
>  	}

I can't help but notice that the local and !POLLING cases show
remarkable similarity. Just not sure extracting that isn't going to make
a mess.

Anyway, yes this looks about right.