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/
Unified · Split
| @@ -852,7 +852,11 @@ struct task_struct { · 851 unmodified lines | ||
| 852 | 852 | struct alloc_tag *alloc_tag; |
| 853 | 853 | #endif |
| 854 | 854 | |
| 855 | − int on_cpu; | |
| 855 | + u8 on_cpu; | |
| 856 | + u8 on_rq; | |
| 857 | + u8 is_blocked; | |
| 858 | + u8 __pad; | |
| 859 | + | |
| 856 | 860 | struct __call_single_node wake_entry; |
| 857 | 861 | unsigned int wakee_flips; |
| 858 | 862 | unsigned long wakee_flip_decay_ts; |
| @@ -867,7 +871,6 @@ struct task_struct { · 8 unmodified lines | ||
| 867 | 871 | */ |
| 868 | 872 | int recent_used_cpu; |
| 869 | 873 | int wake_cpu; |
| 870 | − int on_rq; | |
| 871 | 874 | |
| 872 | 875 | int prio; |
| 873 | 876 | int static_prio; |
| @@ -624,6 +624,12 @@ int task_llc(const struct task_struct *p) · 623 unmodified lines | ||
| 624 | 624 | * [ The astute reader will observe that it is possible for two tasks on one |
| 625 | 625 | * CPU to have ->on_cpu = 1 at the same time. ] |
| 626 | 626 | * |
| 627 | + * p->is_blocked <- { 0, 1 }: | |
| 628 | + * | |
| 629 | + * is set by try_to_block_task() and cleared by ttwu_do_wakeup() and tracks | |
| 630 | + * if the task is blocked. Traditionally this would mirror p->on_rq, however | |
| 631 | + * due things like DELAY_DEQUEUE and PROXY_EXEC, this can diverge. | |
| 632 | + * | |
| 627 | 633 | * task_cpu(p): is changed by set_task_cpu(), the rules are: |
| 628 | 634 | * |
| 629 | 635 | * - 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) · 3089 unmodified lines | ||
| 3719 | 3725 | */ |
| 3720 | 3726 | static inline void ttwu_do_wakeup(struct task_struct *p) |
| 3721 | 3727 | { |
| 3728 | + p->is_blocked = 0; | |
| 3722 | 3729 | WRITE_ONCE(p->__state, TASK_RUNNING); |
| 3723 | 3730 | trace_sched_wakeup(p); |
| 3724 | 3731 | } |
| @@ -4252,6 +4259,7 @@ int try_to_wake_up(struct task_struct *p, unsigned int state, int wake_flags) · 527 unmodified lines | ||
| 4252 | 4259 | * it disabling IRQs (this allows not taking ->pi_lock). |
| 4253 | 4260 | */ |
| 4254 | 4261 | WARN_ON_ONCE(p->se.sched_delayed); |
| 4262 | + WARN_ON_ONCE(p->is_blocked); | |
| 4255 | 4263 | /* If p is current, we know we can run here, so clear blocked_on */ |
| 4256 | 4264 | clear_task_blocked_on(p, NULL); |
| 4257 | 4265 | if (!ttwu_state_match(p, state, &success)) |
| @@ -4563,6 +4571,7 @@ static void __sched_fork(u64 clone_flags, struct task_struct *p) · 305 unmodified lines | ||
| 4563 | 4571 | |
| 4564 | 4572 | /* A delayed task cannot be in clone(). */ |
| 4565 | 4573 | WARN_ON_ONCE(p->se.sched_delayed); |
| 4574 | + WARN_ON_ONCE(p->is_blocked); | |
| 4566 | 4575 | |
| 4567 | 4576 | #ifdef CONFIG_FAIR_GROUP_SCHED |
| 4568 | 4577 | p->se.cfs_rq = NULL; |
| @@ -6676,6 +6685,7 @@ static bool try_to_block_task(struct rq *rq, struct task_struct *p, · 2107 unmodified lines | ||
| 6676 | 6685 | unsigned long task_state = *task_state_p; |
| 6677 | 6686 | |
| 6678 | 6687 | if (signal_pending_state(task_state, p)) { |
| 6688 | + p->is_blocked = 0; | |
| 6679 | 6689 | WRITE_ONCE(p->__state, TASK_RUNNING); |
| 6680 | 6690 | *task_state_p = TASK_RUNNING; |
| 6681 | 6691 | clear_task_blocked_on(p, NULL); |
| @@ -6683,6 +6693,8 @@ static bool try_to_block_task(struct rq *rq, struct task_struct *p, · 1 unmodified line | ||
| 6683 | 6693 | return false; |
| 6684 | 6694 | } |
| 6685 | 6695 | |
| 6696 | + p->is_blocked = 1; | |
| 6697 | + | |
| 6686 | 6698 | /* |
| 6687 | 6699 | * We check should_block after signal_pending because we |
| 6688 | 6700 | * 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) · 154 unmodified lines | ||
| 6843 | 6855 | /* if its PROXY_WAKING, do return migration or run if current */ |
| 6844 | 6856 | if (mutex == PROXY_WAKING) { |
| 6845 | 6857 | if (task_current(rq, p)) { |
| 6858 | + p->is_blocked = 0; | |
| 6846 | 6859 | clear_task_blocked_on(p, PROXY_WAKING); |
| 6847 | 6860 | return p; |
| 6848 | 6861 | } |
| @@ -6878,6 +6891,7 @@ find_proxy_task(struct rq *rq, struct task_struct *donor, struct rq_flags *rf) · 29 unmodified lines | ||
| 6878 | 6891 | * just run on this rq), or return-migrate the task. |
| 6879 | 6892 | */ |
| 6880 | 6893 | if (task_current(rq, p)) { |
| 6894 | + p->is_blocked = 0; | |
| 6881 | 6895 | __clear_task_blocked_on(p, NULL); |
| 6882 | 6896 | return p; |
| 6883 | 6897 | } |
| @@ -7111,7 +7125,7 @@ pick_again: · 227 unmodified lines | ||
| 7111 | 7125 | clear_task_blocked_on(prev, NULL); |
| 7112 | 7126 | |
| 7113 | 7127 | rq_set_donor(rq, next); |
| 7114 | − if (unlikely(next->blocked_on)) { | |
| 7128 | + if (unlikely(next->is_blocked && next->blocked_on)) { | |
| 7115 | 7129 | next = find_proxy_task(rq, next, &rf); |
| 7116 | 7130 | if (!next) { |
| 7117 | 7131 | zap_balance_callbacks(rq); |