3 files changed, 26 insertions, 30 deletions

diff --git a/fs/bcachefs/move.c b/fs/bcachefs/move.c
index 4a5e435bfe4b..65188f1aa966 100644
--- a/fs/bcachefs/move.c
+++ b/fs/bcachefs/move.c
@@ -471,7 +471,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);
@@ -485,12 +485,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)
@@ -561,7 +579,7 @@ next_nondata:
 		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));
diff --git a/fs/bcachefs/util.c b/fs/bcachefs/util.c
index 1272ea7a7a28..5cfaed5b305c 100644
--- a/fs/bcachefs/util.c
+++ b/fs/bcachefs/util.c
@@ -421,27 +421,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: */
 
 /*
diff --git a/fs/bcachefs/util.h b/fs/bcachefs/util.h
index 487591c4c198..178bf98308e8 100644
--- a/fs/bcachefs/util.h
+++ b/fs/bcachefs/util.h
@@ -383,7 +383,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;