struct rq_flags rf;
 
        rq_lock_irqsave(rq, &rf);
-       if (on)
+       if (on) {
+               WARN_ON_ONCE(rq->balance_callback);
                rq->balance_callback = &balance_push_callback;
-       else
+       } else if (rq->balance_callback == &balance_push_callback) {
                rq->balance_callback = NULL;
+       }
        rq_unlock_irqrestore(rq, &rf);
 }
 
        struct rq *rq = cpu_rq(cpu);
        struct rq_flags rf;
 
+       /*
+        * Make sure that when the hotplug state machine does a roll-back
+        * we clear balance_push. Ideally that would happen earlier...
+        */
        balance_push_set(cpu, false);
 
 #ifdef CONFIG_SCHED_SMT
        }
        rq_unlock_irqrestore(rq, &rf);
 
+       /*
+        * Now that the CPU is offline, make sure we're welcome
+        * to new tasks once we come back up.
+        */
+       balance_push_set(cpu, false);
+
        calc_load_migrate(rq);
        update_max_interval();
        nohz_balance_exit_idle(rq);