sched: Add is_blocked task flag

Add a new is_blocked flag to the task struct. This flag is set
by try_to_block_task() and cleared by ttwu_do_wakeup() and
tracks if the task is blocked.

Traditionally this would mirror !p->on_rq, however due things
like DELAY_DEQUEUE and PROXY_EXEC, this can diverge, so its
useful to manage separately.

Additionally with this, we might be able to get rid of the
p->se.sched_delayed (ab)use in the core code (eventually).

Taken whole cloth from Peter's email:
  https://lore.kernel.org/lkml/20260501132143.GC1026330@noisy.programming.kicks-ass.net/

With a few additional p->is_blocked = 0 in a few cases where
we return current if blocked_on gets zeroed or there is
no owner. This may hint that these current special cases
might be dropped eventually.

This change also helps resolve wait-queue stalls seen with
proxy-execution. See previous patch attempts for details:
  https://lore.kernel.org/lkml/20260430215103.2978955-2-jstultz@google.com/

Reported-by: Vineeth Pillai <vineethrp@google.com>
Suggested-by: Peter Zijlstra <peterz@infradead.org>
Signed-off-by: John Stultz <jstultz@google.com>
Signed-off-by: Peter Zijlstra (Intel) <peterz@infradead.org>
Link: https://patch.msgid.link/20260512025635.2840817-7-jstultz@google.com

diff --git a/include/linux/sched.h b/include/linux/sched.h
index 5dea5b10ac99cbda750f404451c3f1d3b248c2f9..ec170663f99b308e746af1bee5ffd9c20b4c6050 100644 (file)
--- a/include/linux/sched.h
+++ b/include/linux/sched.h
@@ -852,7 +852,11 @@ struct task_struct {
        struct alloc_tag                *alloc_tag;
 #endif
 
-       int                             on_cpu;
+       u8                              on_cpu;
+       u8                              on_rq;
+       u8                              is_blocked;
+       u8                              __pad;
+
        struct __call_single_node       wake_entry;
        unsigned int                    wakee_flips;
        unsigned long                   wakee_flip_decay_ts;
@@ -867,7 +871,6 @@ struct task_struct {
         */
        int                             recent_used_cpu;
        int                             wake_cpu;
-       int                             on_rq;
 
        int                             prio;
        int                             static_prio;

diff --git a/kernel/sched/core.c b/kernel/sched/core.c
index 3f71dd9c10637b92fe06894cc66e40b39519e153..c7552869d5c4019ccdee6a663ade359b571a7d4d 100644 (file)
--- a/kernel/sched/core.c
+++ b/kernel/sched/core.c
@@ -624,6 +624,12 @@ int task_llc(const struct task_struct *p)
  *   [ The astute reader will observe that it is possible for two tasks on one
  *     CPU to have ->on_cpu = 1 at the same time. ]
  *
+ * p->is_blocked <- { 0, 1 }:
+ *
+ *   is set by try_to_block_task() and cleared by ttwu_do_wakeup() and tracks
+ *   if the task is blocked. Traditionally this would mirror p->on_rq, however
+ *   due things like DELAY_DEQUEUE and PROXY_EXEC, this can diverge.
+ *
  * task_cpu(p): is changed by set_task_cpu(), the rules are:
  *
  *  - Don't call set_task_cpu() on a blocked task:
@@ -3719,6 +3725,7 @@ ttwu_stat(struct task_struct *p, int cpu, int wake_flags)
  */
 static inline void ttwu_do_wakeup(struct task_struct *p)
 {
+       p->is_blocked = 0;
        WRITE_ONCE(p->__state, TASK_RUNNING);
        trace_sched_wakeup(p);
 }
@@ -4252,6 +4259,7 @@ int try_to_wake_up(struct task_struct *p, unsigned int state, int wake_flags)
                 *    it disabling IRQs (this allows not taking ->pi_lock).
                 */
                WARN_ON_ONCE(p->se.sched_delayed);
+               WARN_ON_ONCE(p->is_blocked);
                /* If p is current, we know we can run here, so clear blocked_on */
                clear_task_blocked_on(p, NULL);
                if (!ttwu_state_match(p, state, &success))
@@ -4563,6 +4571,7 @@ static void __sched_fork(u64 clone_flags, struct task_struct *p)
 
        /* A delayed task cannot be in clone(). */
        WARN_ON_ONCE(p->se.sched_delayed);
+       WARN_ON_ONCE(p->is_blocked);
 
 #ifdef CONFIG_FAIR_GROUP_SCHED
        p->se.cfs_rq                    = NULL;
@@ -6676,6 +6685,7 @@ static bool try_to_block_task(struct rq *rq, struct task_struct *p,
        unsigned long task_state = *task_state_p;
 
        if (signal_pending_state(task_state, p)) {
+               p->is_blocked = 0;
                WRITE_ONCE(p->__state, TASK_RUNNING);
                *task_state_p = TASK_RUNNING;
                clear_task_blocked_on(p, NULL);
@@ -6683,6 +6693,8 @@ static bool try_to_block_task(struct rq *rq, struct task_struct *p,
                return false;
        }
 
+       p->is_blocked = 1;
+
        /*
         * We check should_block after signal_pending because we
         * will want to wake the task in that case. But if
@@ -6843,6 +6855,7 @@ find_proxy_task(struct rq *rq, struct task_struct *donor, struct rq_flags *rf)
                /* if its PROXY_WAKING, do return migration or run if current */
                if (mutex == PROXY_WAKING) {
                        if (task_current(rq, p)) {
+                               p->is_blocked = 0;
                                clear_task_blocked_on(p, PROXY_WAKING);
                                return p;
                        }
@@ -6878,6 +6891,7 @@ find_proxy_task(struct rq *rq, struct task_struct *donor, struct rq_flags *rf)
                         * just run on this rq), or return-migrate the task.
                         */
                        if (task_current(rq, p)) {
+                               p->is_blocked = 0;
                                __clear_task_blocked_on(p, NULL);
                                return p;
                        }
@@ -7111,7 +7125,7 @@ pick_again:
                        clear_task_blocked_on(prev, NULL);
 
                rq_set_donor(rq, next);
-               if (unlikely(next->blocked_on)) {
+               if (unlikely(next->is_blocked && next->blocked_on)) {
                        next = find_proxy_task(rq, next, &rf);
                        if (!next) {
                                zap_balance_callbacks(rq);