diff --git a/fs/bcachefs/alloc_background.c b/fs/bcachefs/alloc_background.c
index 472466d16b68..ce0ea4886288 100644
--- a/fs/bcachefs/alloc_background.c
+++ b/fs/bcachefs/alloc_background.c
@@ -2175,21 +2175,25 @@ void bch2_dev_allocator_remove(struct bch_fs *c, struct bch_dev *ca)
 	}
 	mutex_unlock(&c->btree_reserve_cache_lock);
 
-	while (1) {
-		struct open_bucket *ob;
+	spin_lock(&c->freelist_lock);
+	i = 0;
+	while (i < c->open_buckets_partial_nr) {
+		struct open_bucket *ob =
+			c->open_buckets + c->open_buckets_partial[i];
 
-		spin_lock(&c->freelist_lock);
-		if (!ca->open_buckets_partial_nr) {
+		if (ob->dev == ca->dev_idx) {
+			--c->open_buckets_partial_nr;
+			swap(c->open_buckets_partial[i],
+			     c->open_buckets_partial[c->open_buckets_partial_nr]);
+			ob->on_partial_list = false;
 			spin_unlock(&c->freelist_lock);
-			break;
+			bch2_open_bucket_put(c, ob);
+			spin_lock(&c->freelist_lock);
+		} else {
+			i++;
 		}
-		ob = c->open_buckets +
-			ca->open_buckets_partial[--ca->open_buckets_partial_nr];
-		ob->on_partial_list = false;
-		spin_unlock(&c->freelist_lock);
-
-		bch2_open_bucket_put(c, ob);
 	}
+	spin_unlock(&c->freelist_lock);
 
 	bch2_ec_stop_dev(c, ca);
 
diff --git a/fs/bcachefs/alloc_foreground.c b/fs/bcachefs/alloc_foreground.c
index 0b0fe4fea6cc..1405d6b8cbc5 100644
--- a/fs/bcachefs/alloc_foreground.c
+++ b/fs/bcachefs/alloc_foreground.c
@@ -154,18 +154,15 @@ static void open_bucket_free_unused(struct bch_fs *c,
 				    struct write_point *wp,
 				    struct open_bucket *ob)
 {
-	struct bch_dev *ca = bch_dev_bkey_exists(c, ob->dev);
 	bool may_realloc = wp->data_type == BCH_DATA_user;
 
-	BUG_ON(ca->open_buckets_partial_nr >
-	       ARRAY_SIZE(ca->open_buckets_partial));
+	BUG_ON(c->open_buckets_partial_nr >=
+	       ARRAY_SIZE(c->open_buckets_partial));
 
-	if (ca->open_buckets_partial_nr <
-	    ARRAY_SIZE(ca->open_buckets_partial) &&
-	    may_realloc) {
+	if (may_realloc) {
 		spin_lock(&c->freelist_lock);
 		ob->on_partial_list = true;
-		ca->open_buckets_partial[ca->open_buckets_partial_nr++] =
+		c->open_buckets_partial[c->open_buckets_partial_nr++] =
 			ob - c->open_buckets;
 		spin_unlock(&c->freelist_lock);
 
@@ -394,12 +391,13 @@ static struct open_bucket *try_alloc_partial_bucket(struct bch_fs *c, struct bch
 
 	spin_lock(&c->freelist_lock);
 
-	for (i = ca->open_buckets_partial_nr - 1; i >= 0; --i) {
-		ob = c->open_buckets + ca->open_buckets_partial[i];
+	for (i = c->open_buckets_partial_nr - 1; i >= 0; --i) {
+		ob = c->open_buckets + c->open_buckets_partial[i];
 
-		if (reserve <= ob->alloc_reserve) {
-			array_remove_item(ca->open_buckets_partial,
-					  ca->open_buckets_partial_nr,
+		if (ob->dev == ca->dev_idx &&
+		    reserve <= ob->alloc_reserve) {
+			array_remove_item(c->open_buckets_partial,
+					  c->open_buckets_partial_nr,
 					  i);
 			ob->on_partial_list = false;
 			ob->alloc_reserve = reserve;
diff --git a/fs/bcachefs/bcachefs.h b/fs/bcachefs/bcachefs.h
index 85a815cdf586..509a16469613 100644
--- a/fs/bcachefs/bcachefs.h
+++ b/fs/bcachefs/bcachefs.h
@@ -516,9 +516,6 @@ struct bch_dev {
 	unsigned		nr_open_buckets;
 	unsigned		nr_btree_reserve;
 
-	open_bucket_idx_t	open_buckets_partial[OPEN_BUCKETS_COUNT];
-	open_bucket_idx_t	open_buckets_partial_nr;
-
 	size_t			inc_gen_needs_gc;
 	size_t			inc_gen_really_needs_gc;
 	size_t			buckets_waiting_on_journal;
@@ -859,6 +856,9 @@ struct bch_fs {
 	struct open_bucket	open_buckets[OPEN_BUCKETS_COUNT];
 	open_bucket_idx_t	open_buckets_hash[OPEN_BUCKETS_COUNT];
 
+	open_bucket_idx_t	open_buckets_partial[OPEN_BUCKETS_COUNT];
+	open_bucket_idx_t	open_buckets_partial_nr;
+
 	struct write_point	btree_write_point;
 	struct write_point	rebalance_write_point;