}
 
 /*
- * Stop waiting if either state is not BP_EAGAIN and ballooning action is
- * needed, or if the credit has changed while state is BP_EAGAIN.
+ * Stop waiting if either state is BP_DONE and ballooning action is
+ * needed, or if the credit has changed while state is not BP_DONE.
  */
 static bool balloon_thread_cond(enum bp_state state, long credit)
 {
-       if (state != BP_EAGAIN)
+       if (state == BP_DONE)
                credit = 0;
 
        return current_credit() != credit || kthread_should_stop();
 
        set_freezable();
        for (;;) {
-               if (state == BP_EAGAIN)
-                       timeout = balloon_stats.schedule_delay * HZ;
-               else
+               switch (state) {
+               case BP_DONE:
+               case BP_ECANCELED:
                        timeout = 3600 * HZ;
+                       break;
+               case BP_EAGAIN:
+                       timeout = balloon_stats.schedule_delay * HZ;
+                       break;
+               case BP_WAIT:
+                       timeout = HZ;
+                       break;
+               }
+
                credit = current_credit();
 
                wait_event_freezable_timeout(balloon_thread_wq,