struct rq_flags rf;
 
        rq_lock_irqsave(rq, &rf);
+       rq->balance_push = on;
        if (on) {
                WARN_ON_ONCE(rq->balance_callback);
                rq->balance_callback = &balance_push_callback;
        int ret;
 
        set_cpu_active(cpu, false);
+       balance_push_set(cpu, true);
+
        /*
-        * We've cleared cpu_active_mask, wait for all preempt-disabled and RCU
-        * users of this state to go away such that all new such users will
-        * observe it.
+        * We've cleared cpu_active_mask / set balance_push, wait for all
+        * preempt-disabled and RCU users of this state to go away such that
+        * all new such users will observe it.
         *
         * Do sync before park smpboot threads to take care the rcu boost case.
         */
        synchronize_rcu();
 
-       balance_push_set(cpu, true);
-
        rq_lock_irqsave(rq, &rf);
        if (rq->rd) {
                update_rq_clock(rq);
 
        unsigned long           cpu_capacity_orig;
 
        struct callback_head    *balance_callback;
+       unsigned char           balance_push;
 
        unsigned char           nohz_idle_balance;
        unsigned char           idle_balance;