bcachefs: Fix suspend when moving data faster than ratelimit

Signed-off-by: Kent Overstreet <kent.overstreet@linux.dev>

fs/bcachefs/move.c

@@ -470,7 +470,7 @@ int bch2_move_data(struct bch_fs *c,
 	struct bkey_s_c_extent e;
 	struct data_opts data_opts;
 	enum data_cmd data_cmd;
-	u64 cur_inum = U64_MAX;
+	u64 delay, cur_inum = U64_MAX;
 	int ret = 0, ret2;
 
 	closure_init_stack(&ctxt.cl);
@@ -484,12 +484,30 @@ int bch2_move_data(struct bch_fs *c,
 	if (rate)
 		bch2_ratelimit_reset(rate);
 
-	while (!kthread || !(ret = kthread_should_stop())) {
-		if (rate &&
-		    bch2_ratelimit_delay(rate) &&
-		    (bch2_btree_iter_unlock(&stats->iter),
-		     (ret = bch2_ratelimit_wait_freezable_stoppable(rate))))
-			break;
+	while (1) {
+		do {
+			delay = rate ? bch2_ratelimit_delay(rate) : 0;
+
+			if (delay) {
+				bch2_btree_iter_unlock(&stats->iter);
+				set_current_state(TASK_INTERRUPTIBLE);
+			}
+
+			if (kthread && (ret = kthread_should_stop())) {
+				__set_current_state(TASK_RUNNING);
+				goto out;
+			}
+
+			if (delay)
+				schedule_timeout(delay);
+
+			if (unlikely(freezing(current))) {
+				bch2_btree_iter_unlock(&stats->iter);
+				move_ctxt_wait_event(&ctxt, list_empty(&ctxt.reads));
+				closure_sync(&ctxt.cl);
+				try_to_freeze();
+			}
+		} while (delay);
 peek:
 		k = bch2_btree_iter_peek(&stats->iter);
 		if (!k.k)
@@ -560,7 +578,7 @@ int bch2_move_data(struct bch_fs *c,
 		bch2_btree_iter_next(&stats->iter);
 		bch2_btree_iter_cond_resched(&stats->iter);
 	}
-
+out:
 	bch2_btree_iter_unlock(&stats->iter);
 
 	move_ctxt_wait_event(&ctxt, list_empty(&ctxt.reads));

fs/bcachefs/util.c

@@ -424,27 +424,6 @@ void bch2_ratelimit_increment(struct bch_ratelimit *d, u64 done)
 		d->next = now - NSEC_PER_SEC * 2;
 }
 
-int bch2_ratelimit_wait_freezable_stoppable(struct bch_ratelimit *d)
-{
-	bool kthread = (current->flags & PF_KTHREAD) != 0;
-
-	while (1) {
-		u64 delay = bch2_ratelimit_delay(d);
-
-		if (delay)
-			set_current_state(TASK_INTERRUPTIBLE);
-
-		if (kthread && kthread_should_stop())
-			return 1;
-
-		if (!delay)
-			return 0;
-
-		schedule_timeout(delay);
-		try_to_freeze();
-	}
-}
-
 /* pd controller: */
 
 /*

fs/bcachefs/util.h

@@ -377,7 +377,6 @@ static inline void bch2_ratelimit_reset(struct bch_ratelimit *d)
 
 u64 bch2_ratelimit_delay(struct bch_ratelimit *);
 void bch2_ratelimit_increment(struct bch_ratelimit *, u64);
-int bch2_ratelimit_wait_freezable_stoppable(struct bch_ratelimit *);
 
 struct bch_pd_controller {
 	struct bch_ratelimit	rate;