// SPDX-License-Identifier: GPL-2.0
#include "bcachefs.h"

#include "alloc/background.h"
#include "alloc/check.h"
#include "alloc/lru.h"

#include "btree/cache.h"
#include "btree/update.h"
#include "btree/write_buffer.h"

#include "data/ec/init.h"

#include "init/error.h"
#include "init/passes.h"
#include "init/progress.h"

/*
 * This synthesizes deleted extents for holes, similar to BTREE_ITER_slots for
 * extents style btrees, but works on non-extents btrees:
 */
static struct bkey_s_c bch2_get_key_or_hole(struct btree_iter *iter, struct bpos end, struct bkey *hole)
{
	struct bkey_s_c k = bch2_btree_iter_peek_slot(iter);

	if (bkey_err(k))
		return k;

	if (k.k->type) {
		return k;
	} else {
		CLASS(btree_iter_copy, iter2)(iter);

		struct btree_path *path = btree_iter_path(iter->trans, iter);
		if (!bpos_eq(path->l[0].b->key.k.p, SPOS_MAX))
			end = bkey_min(end, bpos_nosnap_successor(path->l[0].b->key.k.p));

		end = bkey_min(end, POS(iter->pos.inode, iter->pos.offset + U32_MAX - 1));

		/*
		 * btree node min/max is a closed interval, upto takes a half
		 * open interval:
		 */
		k = bch2_btree_iter_peek_max(&iter2, &end);
		if (bkey_err(k))
			return k;

		struct bpos next = iter2.pos;
		BUG_ON(next.offset >= iter->pos.offset + U32_MAX);

		bkey_init(hole);
		hole->p = iter->pos;

		bch2_key_resize(hole, next.offset - iter->pos.offset);
		return (struct bkey_s_c) { hole, NULL };
	}
}

static bool next_bucket(struct bch_fs *c, struct bch_dev **ca, struct bpos *bucket)
{
	if (*ca) {
		if (bucket->offset < (*ca)->mi.first_bucket)
			bucket->offset = (*ca)->mi.first_bucket;

		if (bucket->offset < (*ca)->mi.nbuckets)
			return true;

		bch2_dev_put(*ca);
		*ca = NULL;
		bucket->inode++;
		bucket->offset = 0;
	}

	guard(rcu)();
	*ca = __bch2_next_dev_idx(c, bucket->inode, NULL);
	if (*ca) {
		*bucket = POS((*ca)->dev_idx, (*ca)->mi.first_bucket);
		bch2_dev_get(*ca);
	}

	return *ca != NULL;
}

static struct bkey_s_c bch2_get_key_or_real_bucket_hole(struct btree_iter *iter,
					struct bch_dev **ca, struct bkey *hole)
{
	while (true) {
		struct bch_fs *c = iter->trans->c;
		struct bkey_s_c k = bch2_get_key_or_hole(iter, POS_MAX, hole);
		if (bkey_err(k))
			return k;

		*ca = bch2_dev_iterate_noerror(c, *ca, k.k->p.inode);

		if (!k.k->type) {
			struct bpos hole_start = bkey_start_pos(k.k);

			if (!*ca || !bucket_valid(*ca, hole_start.offset)) {
				if (!next_bucket(c, ca, &hole_start))
					return bkey_s_c_null;

				bch2_btree_iter_set_pos(iter, hole_start);
				continue;
			}

			if (k.k->p.offset > (*ca)->mi.nbuckets)
				bch2_key_resize(hole, (*ca)->mi.nbuckets - hole_start.offset);
		}

		return k;
	}
}

int bch2_need_discard_or_freespace_err(struct btree_trans *trans,
					 struct bkey_s_c alloc_k,
					 bool set, bool discard, bool repair)
{
	struct bch_fs *c = trans->c;
	enum bch_fsck_flags flags = FSCK_CAN_IGNORE|(repair ? FSCK_CAN_FIX : 0);
	enum bch_sb_error_id err_id = discard
		? BCH_FSCK_ERR_need_discard_key_wrong
		: BCH_FSCK_ERR_freespace_key_wrong;
	enum btree_id btree = discard ? BTREE_ID_need_discard : BTREE_ID_freespace;
	CLASS(printbuf, buf)();

	bch2_bkey_val_to_text(&buf, c, alloc_k);

	int ret = __bch2_fsck_err(NULL, trans, POS_MIN, flags, err_id,
				  "bucket incorrectly %sset in %s btree\n%s",
				  set ? "" : "un",
				  bch2_btree_id_str(btree),
				  buf.buf);
	if (bch2_err_matches(ret, BCH_ERR_fsck_ignore) ||
	    bch2_err_matches(ret, BCH_ERR_fsck_errors_not_fixed))
		ret = 0;
	return ret;
}

static noinline_for_stack
int bch2_check_alloc_key(struct btree_trans *trans,
			 struct bkey_s_c alloc_k,
			 struct btree_iter *alloc_iter,
			 struct btree_iter *discard_iter,
			 struct btree_iter *freespace_iter,
			 struct btree_iter *bucket_gens_iter,
			 struct wb_maybe_flush *last_flushed)
{
	struct bch_fs *c = trans->c;
	unsigned gens_offset;
	struct bkey_s_c k;
	CLASS(printbuf, buf)();
	int ret = 0;

	CLASS(bch2_dev_bucket_tryget_noerror, ca)(c, alloc_k.k->p);
	if (ret_fsck_err_on(!ca,
			trans, alloc_key_to_missing_dev_bucket,
			"alloc key for invalid device:bucket %llu:%llu",
			alloc_k.k->p.inode, alloc_k.k->p.offset))
		try(bch2_btree_delete_at(trans, alloc_iter, 0));
	if (!ca)
		return 0;

	if (!ca->mi.freespace_initialized)
		return 0;

	struct bkey_i_alloc_v4 *a = errptr_try(bch2_alloc_to_v4_mut(trans, alloc_k));

	if (a->v.data_type == BCH_DATA_need_discard) {
		bch2_btree_iter_set_pos(discard_iter,
			POS(a->v.journal_seq_empty, bucket_to_u64(alloc_k.k->p)));
		k = bkey_try(bch2_btree_iter_peek_slot(discard_iter));

		if (k.k->type != KEY_TYPE_set)
			try(bch2_btree_write_buffer_maybe_flush(trans, alloc_k, last_flushed));

		if (need_discard_or_freespace_err_on(k.k->type != KEY_TYPE_set,
						     trans, alloc_k, false, true, true))
			try(bch2_btree_bit_mod_buffered(trans, BTREE_ID_need_discard,
							discard_iter->pos, true));
	}

	bch2_btree_iter_set_pos(freespace_iter, alloc_freespace_pos(alloc_k.k->p, a->v));
	k = bkey_try(bch2_btree_iter_peek_slot(freespace_iter));

	bool is_free = a->v.data_type == BCH_DATA_free;
	if (need_discard_or_freespace_err_on(!!k.k->type != is_free,
					     trans, alloc_k, !is_free, false, true))
		try(bch2_btree_bit_mod_iter(trans, freespace_iter, is_free));

	bch2_btree_iter_set_pos(bucket_gens_iter, alloc_gens_pos(alloc_k.k->p, &gens_offset));
	k = bkey_try(bch2_btree_iter_peek_slot(bucket_gens_iter));

	if (ret_fsck_err_on(a->v.generation != alloc_gen(k, gens_offset),
			trans, bucket_gens_key_wrong,
			"incorrect gen in bucket_gens btree (got %u should be %u)\n%s",
			alloc_gen(k, gens_offset), a->v.generation,
			(printbuf_reset(&buf),
			 bch2_bkey_val_to_text(&buf, c, alloc_k), buf.buf))) {
		struct bkey_i_bucket_gens *g =
			errptr_try(bch2_trans_kmalloc(trans, sizeof(*g)));

		if (k.k->type == KEY_TYPE_bucket_gens) {
			bkey_reassemble(&g->k_i, k);
		} else {
			bkey_bucket_gens_init(&g->k_i);
			g->k.p = alloc_gens_pos(alloc_k.k->p, &gens_offset);
		}

		g->v.gens[gens_offset] = a->v.generation;

		try(bch2_trans_update(trans, bucket_gens_iter, &g->k_i, 0));
	}

	ret = bch2_bucket_nr_stripes(trans, alloc_k.k->p);
	if (ret < 0)
		return ret;
	unsigned nr_stripes = ret;
	ret = 0;

	if (fsck_err_on(a->v.stripe_refcount != nr_stripes,
			trans, alloc_key_stripe_refcount_wrong,
			"alloc key with incorrect stripe_refcount, should be %u\n%s",
			nr_stripes,
			(printbuf_reset(&buf),
			 bch2_bkey_val_to_text(&buf, c, alloc_k), buf.buf))) {
		a->v.stripe_refcount = nr_stripes;
		try(bch2_trans_update(trans, alloc_iter, &a->k_i, 0));
	}
fsck_err:
	return ret;
}

static noinline_for_stack
int bch2_check_alloc_hole_freespace(struct btree_trans *trans,
				    struct bch_dev *ca,
				    struct bpos start,
				    struct bpos *end,
				    struct btree_iter *freespace_iter)
{
	CLASS(printbuf, buf)();

	if (!ca->mi.freespace_initialized)
		return 0;

	bch2_btree_iter_set_pos(freespace_iter, start);

	struct bkey_s_c k = bkey_try(bch2_btree_iter_peek_slot(freespace_iter));

	*end = bkey_min(k.k->p, *end);

	if (ret_fsck_err_on(k.k->type != KEY_TYPE_set,
			trans, freespace_hole_missing,
			"hole in alloc btree missing in freespace btree\n"
			"device %llu buckets %llu-%llu",
			freespace_iter->pos.inode,
			freespace_iter->pos.offset,
			end->offset)) {
		struct bkey_i *update =
			errptr_try(bch2_trans_kmalloc(trans, sizeof(*update)));

		bkey_init(&update->k);
		update->k.type	= KEY_TYPE_set;
		update->k.p	= freespace_iter->pos;
		bch2_key_resize(&update->k,
				min_t(u64, U32_MAX, end->offset -
				      freespace_iter->pos.offset));

		try(bch2_trans_update(trans, freespace_iter, update, 0));
	}

	return 0;
}

static noinline_for_stack
int bch2_check_alloc_hole_bucket_gens(struct btree_trans *trans,
				      struct bpos start,
				      struct bpos *end,
				      struct btree_iter *bucket_gens_iter)
{
	CLASS(printbuf, buf)();
	unsigned gens_offset, gens_end_offset;

	bch2_btree_iter_set_pos(bucket_gens_iter, alloc_gens_pos(start, &gens_offset));

	struct bkey_s_c k = bkey_try(bch2_btree_iter_peek_slot(bucket_gens_iter));

	if (bkey_cmp(alloc_gens_pos(start, &gens_offset),
		     alloc_gens_pos(*end,  &gens_end_offset)))
		gens_end_offset = KEY_TYPE_BUCKET_GENS_NR;

	if (k.k->type == KEY_TYPE_bucket_gens) {
		struct bkey_i_bucket_gens g;
		bool need_update = false;

		bkey_reassemble(&g.k_i, k);

		for (unsigned i = gens_offset; i < gens_end_offset; i++) {
			if (ret_fsck_err_on(g.v.gens[i], trans,
					bucket_gens_hole_wrong,
					"hole in alloc btree at %llu:%llu with nonzero gen in bucket_gens btree (%u)",
					bucket_gens_pos_to_alloc(k.k->p, i).inode,
					bucket_gens_pos_to_alloc(k.k->p, i).offset,
					g.v.gens[i])) {
				g.v.gens[i] = 0;
				need_update = true;
			}
		}

		if (need_update) {
			struct bkey_i *u = errptr_try(bch2_trans_kmalloc(trans, sizeof(g)));

			memcpy(u, &g, sizeof(g));

			try(bch2_trans_update(trans, bucket_gens_iter, u, 0));
		}
	}

	*end = bkey_min(*end, bucket_gens_pos_to_alloc(bpos_nosnap_successor(k.k->p), 0));
	return 0;
}

struct check_freespace_key_async {
	struct work_struct	work;
	struct bch_fs		*c;
	struct bbpos		pos;
};

static int bch2_recheck_freespace_key(struct btree_trans *trans, struct bbpos pos)
{
	CLASS(btree_iter, iter)(trans, pos.btree, pos.pos, 0);
	struct bkey_s_c k = bkey_try(bch2_btree_iter_peek_slot(&iter));

	u8 gen;
	return k.k->type != KEY_TYPE_set
		? __bch2_check_freespace_key(trans, &iter, &gen, NULL, FSCK_ERR_SILENT, NULL)
		: 0;
}

static void check_discard_freespace_key_work(struct work_struct *work)
{
	struct check_freespace_key_async *w =
		container_of(work, struct check_freespace_key_async, work);

	bch2_trans_do(w->c, bch2_recheck_freespace_key(trans, w->pos));
	enumerated_ref_put(&w->c->writes, BCH_WRITE_REF_check_discard_freespace_key);
	kfree(w);
}

static int delete_freespace_key(struct btree_trans *trans,
				struct btree_iter *iter,
				bool async_repair)
{
	struct bch_fs *c = trans->c;

	if (!async_repair) {
		try(bch2_btree_bit_mod_iter(trans, iter, false));
		try(bch2_trans_commit(trans, NULL, NULL, BCH_TRANS_COMMIT_no_enospc));
		return bch_err_throw(c, transaction_restart_commit);
	} else {
		/*
		 * We can't repair here when called from the allocator path: the
		 * commit will recurse back into the allocator
		 *
		 * Returning 1 indicates to the caller of
		 * check_discard_freespace_key() - "don't allocate this bucket"
		 */
		struct check_freespace_key_async *w = kzalloc(sizeof(*w), GFP_KERNEL);
		if (!w)
			return 1;

		if (!enumerated_ref_tryget(&c->writes, BCH_WRITE_REF_check_discard_freespace_key)) {
			kfree(w);
			return 1;
		}

		INIT_WORK(&w->work, check_discard_freespace_key_work);
		w->c = c;
		w->pos = BBPOS(iter->btree_id, iter->pos);
		queue_work(c->write_ref_wq, &w->work);
		return 1;
	}
}

static int bch2_check_discard_key(struct btree_trans *trans, struct btree_iter *iter,
				  struct wb_maybe_flush *last_flushed)
{
	struct bch_fs *c = trans->c;
	struct bpos bucket = u64_to_bucket(iter->pos.offset);

	if (!bch2_dev_bucket_exists(c, bucket)) {
		if (ret_fsck_err(trans,
			     need_discard_freespace_key_to_invalid_dev_bucket,
			     "entry in %s btree for nonexistant dev:bucket %llu:%llu",
			     bch2_btree_id_str(iter->btree_id), bucket.inode, bucket.offset))
			try(bch2_btree_bit_mod_buffered(trans, BTREE_ID_need_discard, iter->pos, false));
		return 0;
	}

	CLASS(btree_iter, alloc_iter)(trans, BTREE_ID_alloc, bucket, 0);
	struct bkey_s_c alloc_k = bkey_try(bch2_btree_iter_peek_slot(&alloc_iter));

	struct bch_alloc_v4 a_convert;
	const struct bch_alloc_v4 *a = bch2_alloc_to_v4(alloc_k, &a_convert);

	if (a->data_type != BCH_DATA_need_discard ||
	    a->journal_seq_empty != iter->pos.inode) {
		try(bch2_btree_write_buffer_maybe_flush(trans, alloc_k, last_flushed));

		CLASS(printbuf, buf)();
		if (ret_fsck_err(trans, need_discard_freespace_key_bad,
			     "%s\nincorrectly set at %s:%llu:%llu:0",
			     (bch2_bkey_val_to_text(&buf, c, alloc_k), buf.buf),
			     bch2_btree_id_str(iter->btree_id),
			     iter->pos.inode,
			     iter->pos.offset))
			try(bch2_btree_bit_mod_buffered(trans, BTREE_ID_need_discard, iter->pos, false));
	}

	return 0;
}

int __bch2_check_freespace_key(struct btree_trans *trans, struct btree_iter *iter, u8 *gen,
			       u64 *journal_seq_empty, enum bch_fsck_flags fsck_flags,
			       struct wb_maybe_flush *last_flushed)
{
	struct bch_fs *c = trans->c;
	CLASS(printbuf, buf)();
	int ret = 0;

	bool async_repair = fsck_flags & FSCK_ERR_NO_LOG;
	fsck_flags |= FSCK_CAN_FIX|FSCK_CAN_IGNORE;

	struct bpos bucket = iter->pos;
	bucket.offset &= ~(~0ULL << 56);
	u64 genbits = iter->pos.offset & (~0ULL << 56);

	/*
	 * async_repair means we're the allocator: committed-only reads, see
	 * bucket_alloc_scan(). fsck's synchronous use keeps normal semantics:
	 */
	CLASS(btree_iter, alloc_iter)(trans, BTREE_ID_alloc, bucket,
						     async_repair
						     ? BTREE_ITER_cached|BTREE_ITER_committed
						     : 0);
	struct bkey_s_c alloc_k = bkey_try(bch2_btree_iter_peek_slot(&alloc_iter));

	if (!bch2_dev_bucket_exists(c, bucket)) {
		if (__fsck_err(trans, fsck_flags,
			       need_discard_freespace_key_to_invalid_dev_bucket,
			       "entry in %s btree for nonexistant dev:bucket %llu:%llu",
			       bch2_btree_id_str(iter->btree_id), bucket.inode, bucket.offset))
			ret = delete_freespace_key(trans, iter, async_repair);
		else
			ret = 1;
		goto out;
	}

	struct bch_alloc_v4 a_convert;
	const struct bch_alloc_v4 *a = bch2_alloc_to_v4(alloc_k, &a_convert);

	if (a->data_type != BCH_DATA_free ||
	    genbits != alloc_freespace_genbits(*a)) {
		if (last_flushed)
			try(bch2_btree_write_buffer_maybe_flush(trans, alloc_k, last_flushed));

		if (__fsck_err(trans, fsck_flags,
			       need_discard_freespace_key_bad,
			     "%s\nincorrectly set at %s:%llu:%llu:0 (genbits %llu should be %llu)",
			     (bch2_bkey_val_to_text(&buf, c, alloc_k), buf.buf),
			     bch2_btree_id_str(iter->btree_id),
			     iter->pos.inode,
			     iter->pos.offset,
			     genbits >> 56, alloc_freespace_genbits(*a) >> 56))
			ret = delete_freespace_key(trans, iter, async_repair);
		else
			ret = 1;
		goto out;
	}

	*gen = a->generation;
	if (journal_seq_empty)
		*journal_seq_empty = a->journal_seq_empty;
out:
fsck_err:
	bch2_set_btree_iter_dontneed(&alloc_iter);
	return ret;
}

static int bch2_check_freespace_key(struct btree_trans *trans, struct btree_iter *iter)
{
	u8 gen;
	int ret = __bch2_check_freespace_key(trans, iter, &gen, NULL, 0, NULL);
	return ret < 0 ? ret : 0;
}

/*
 * We've already checked that generation numbers in the bucket_gens btree are
 * valid for buckets that exist; this just checks for keys for nonexistent
 * buckets.
 */
static noinline_for_stack
int bch2_check_bucket_gens_key(struct btree_trans *trans,
			       struct btree_iter *iter,
			       struct bkey_s_c k)
{
	struct bch_fs *c = trans->c;
	struct bkey_i_bucket_gens g;
	u64 start = bucket_gens_pos_to_alloc(k.k->p, 0).offset;
	u64 end = bucket_gens_pos_to_alloc(bpos_nosnap_successor(k.k->p), 0).offset;
	u64 b;
	bool need_update = false;
	CLASS(printbuf, buf)();

	BUG_ON(k.k->type != KEY_TYPE_bucket_gens);
	bkey_reassemble(&g.k_i, k);

	CLASS(bch2_dev_tryget_noerror, ca)(c, k.k->p.inode);
	if (!ca) {
		if (ret_fsck_err(trans, bucket_gens_to_invalid_dev,
			     "bucket_gens key for invalid device:\n%s",
			     (bch2_bkey_val_to_text(&buf, c, k), buf.buf)))
			return bch2_btree_delete_at(trans, iter, 0);
		return 0;
	}

	if (ret_fsck_err_on(end <= ca->mi.first_bucket ||
			start >= ca->mi.nbuckets,
			trans, bucket_gens_to_invalid_buckets,
			"bucket_gens key for invalid buckets:\n%s",
			(bch2_bkey_val_to_text(&buf, c, k), buf.buf))) {
		return bch2_btree_delete_at(trans, iter, 0);
	}

	for (b = start; b < ca->mi.first_bucket; b++)
		if (ret_fsck_err_on(g.v.gens[b & KEY_TYPE_BUCKET_GENS_MASK],
				trans, bucket_gens_nonzero_for_invalid_buckets,
				"bucket_gens key has nonzero gen for invalid bucket")) {
			g.v.gens[b & KEY_TYPE_BUCKET_GENS_MASK] = 0;
			need_update = true;
		}

	for (b = ca->mi.nbuckets; b < end; b++)
		if (ret_fsck_err_on(g.v.gens[b & KEY_TYPE_BUCKET_GENS_MASK],
				trans, bucket_gens_nonzero_for_invalid_buckets,
				"bucket_gens key has nonzero gen for invalid bucket")) {
			g.v.gens[b & KEY_TYPE_BUCKET_GENS_MASK] = 0;
			need_update = true;
		}

	if (need_update) {
		struct bkey_i *u = errptr_try(bch2_trans_kmalloc(trans, sizeof(g)));

		memcpy(u, &g, sizeof(g));
		try(bch2_trans_update(trans, iter, u, 0));
	}

	return 0;
}

static int check_btree_alloc_iter(struct btree_trans *trans,
				  struct bch_dev **ca,
				  struct btree_iter *iter,
				  struct btree_iter *discard_iter,
				  struct btree_iter *freespace_iter,
				  struct btree_iter *bucket_gens_iter,
				  struct progress_indicator *progress,
				  struct wb_maybe_flush *last_flushed)
{
	try(bch2_recovery_cancelled(trans->c));

	struct bkey hole;
	struct bkey_s_c k = bkey_try(bch2_get_key_or_real_bucket_hole(iter, ca, &hole));

	if (!k.k)
		return 1;

	try(bch2_progress_update_iter(trans, progress, iter));

	struct bpos next;
	if (k.k->type) {
		next = bpos_nosnap_successor(k.k->p);

		try(bch2_check_alloc_key(trans, k, iter,
					 discard_iter,
					 freespace_iter,
					 bucket_gens_iter,
					 last_flushed));
	} else {
		next = k.k->p;

		try(bch2_check_alloc_hole_freespace(trans, *ca,
						    bkey_start_pos(k.k),
						    &next,
						    freespace_iter));
		try(bch2_check_alloc_hole_bucket_gens(trans,
						      bkey_start_pos(k.k),
						      &next,
						      bucket_gens_iter));
	}

	try(bch2_trans_commit(trans, NULL, NULL, BCH_TRANS_COMMIT_no_enospc));

	bch2_btree_iter_set_pos(iter, next);
	return 0;
}

static int check_btree_alloc(struct btree_trans *trans)
{
	struct progress_indicator progress;
	bch2_progress_init(&progress, __func__, trans->c, BIT_ULL(BTREE_ID_alloc), 0);

	CLASS(btree_iter, iter)(trans, BTREE_ID_alloc, POS_MIN, BTREE_ITER_prefetch);
	CLASS(btree_iter, discard_iter)(trans, BTREE_ID_need_discard, POS_MIN, BTREE_ITER_prefetch);
	CLASS(btree_iter, freespace_iter)(trans, BTREE_ID_freespace, POS_MIN, BTREE_ITER_prefetch);
	CLASS(btree_iter, bucket_gens_iter)(trans, BTREE_ID_bucket_gens, POS_MIN, BTREE_ITER_prefetch);

	struct wb_maybe_flush last_flushed __cleanup(wb_maybe_flush_exit);
	wb_maybe_flush_init(&last_flushed);

	struct bch_dev *ca __free(bch2_dev_put) = NULL;
	int ret = 0;

	while (!(ret = lockrestart_do(trans,
				check_btree_alloc_iter(trans, &ca, &iter,
						       &discard_iter,
						       &freespace_iter,
						       &bucket_gens_iter,
						       &progress,
						       &last_flushed))))
		;

	return min(0, ret);
}

int bch2_check_alloc_info(struct bch_fs *c)
{
	CLASS(btree_trans, trans)(c);

	try(check_btree_alloc(trans));

	{
		struct wb_maybe_flush last_flushed __cleanup(wb_maybe_flush_exit);
		wb_maybe_flush_init(&last_flushed);

		try(for_each_btree_key_commit(trans, iter,
				BTREE_ID_need_discard, POS_MIN,
				BTREE_ITER_prefetch, k,
				NULL, NULL, BCH_TRANS_COMMIT_no_enospc,
			bch2_recovery_cancelled(c) ?:
			bch2_check_discard_key(trans, &iter, &last_flushed)));
	}

	{
		/*
		 * Check freespace btree: we're iterating over every individual
		 * pos of the freespace keys
		 */
		CLASS(btree_iter, iter)(trans, BTREE_ID_freespace, POS_MIN,
				     BTREE_ITER_prefetch);
		while (1) {
			try(bch2_recovery_cancelled(c));

			bch2_trans_begin(trans);
			struct bkey_s_c k = bch2_btree_iter_peek(&iter);
			if (!k.k)
				break;

			int ret = bkey_err(k) ?:
				bch2_check_freespace_key(trans, &iter);
			if (bch2_err_matches(ret, BCH_ERR_transaction_restart)) {
				ret = 0;
				continue;
			}
			if (ret) {
				CLASS(printbuf, buf)();
				bch2_bkey_val_to_text(&buf, c, k);
				bch_err(c, "while checking %s", buf.buf);
				return ret;
			}

			bch2_btree_iter_set_pos(&iter, bpos_nosnap_successor(iter.pos));
		}
	}

	try(for_each_btree_key_commit(trans, iter,
			BTREE_ID_bucket_gens, POS_MIN,
			BTREE_ITER_prefetch, k,
			NULL, NULL, BCH_TRANS_COMMIT_no_enospc,
		bch2_recovery_cancelled(c) ?:
		bch2_check_bucket_gens_key(trans, &iter, k)));

	return 0;
}

static int bch2_check_alloc_to_lru_ref(struct btree_trans *trans,
				       struct btree_iter *alloc_iter,
				       struct wb_maybe_flush *last_flushed)
{
	struct bch_fs *c = trans->c;
	struct bch_alloc_v4 a_convert;
	const struct bch_alloc_v4 *a;
	CLASS(printbuf, buf)();

	struct bkey_s_c alloc_k = bkey_try(bch2_btree_iter_peek(alloc_iter));
	if (!alloc_k.k)
		return 0;

	CLASS(bch2_dev_tryget_noerror, ca)(c, alloc_k.k->p.inode);
	if (!ca)
		return 0;

	a = bch2_alloc_to_v4(alloc_k, &a_convert);

	u64 lru_idx = alloc_lru_idx_fragmentation(*a, ca);
	if (lru_idx) {
		/*
		 * per_dev_fragmentation_lru upgrade: the per-device entry is
		 * either missing or was just written by check_lrus' migration
		 * and is sitting unflushed in the write buffer, which reads as
		 * missing - either way the lookup can't win, and the write
		 * buffer flush it triggers is far more expensive than just
		 * recreating the entry unconditionally:
		 */
		if (c->sb.version_upgrade_complete < bcachefs_metadata_version_per_dev_fragmentation_lru)
			try(bch2_lru_set(trans,
					 bucket_fragmentation_lru(alloc_k.k->p.inode),
					 bucket_to_u64(alloc_k.k->p),
					 lru_idx));
		else
			try(bch2_lru_check_set(trans,
					       bucket_fragmentation_lru(alloc_k.k->p.inode),
					       bucket_to_u64(alloc_k.k->p),
					       lru_idx, alloc_k, last_flushed));
	}

	if (a->data_type == BCH_DATA_cached) {
		if (ret_fsck_err_on(!a->io_time[READ],
				trans, alloc_key_cached_but_read_time_zero,
				"cached bucket with read_time 0\n%s",
			(printbuf_reset(&buf),
			 bch2_bkey_val_to_text(&buf, c, alloc_k), buf.buf))) {
			struct bkey_i_alloc_v4 *a_mut =
				errptr_try(bch2_alloc_to_v4_mut(trans, alloc_k));

			a_mut->v.io_time[READ] = bch2_current_io_time(c, READ);
			try(bch2_trans_update(trans, alloc_iter,
					      &a_mut->k_i, BTREE_TRIGGER_norun));

			a = &a_mut->v;
		}

		try(bch2_lru_check_set(trans, alloc_k.k->p.inode,
				       bucket_to_u64(alloc_k.k->p),
				       a->io_time[READ],
				       alloc_k, last_flushed));
	}

	return 0;
}

int bch2_check_alloc_to_lru_refs(struct bch_fs *c)
{
	struct wb_maybe_flush last_flushed __cleanup(wb_maybe_flush_exit);
	wb_maybe_flush_init(&last_flushed);

	struct progress_indicator progress;
	bch2_progress_init(&progress, __func__, c, BIT_ULL(BTREE_ID_alloc), 0);

	CLASS(btree_trans, trans)(c);
	return for_each_btree_key_commit(trans, iter, BTREE_ID_alloc,
				POS_MIN, BTREE_ITER_prefetch, k,
				NULL, NULL, BCH_TRANS_COMMIT_no_enospc, ({
			bch2_progress_update_iter(trans, &progress, &iter) ?:
			wb_maybe_flush_inc(&last_flushed) ?:
			bch2_check_alloc_to_lru_ref(trans, &iter, &last_flushed);
	}))?: bch2_check_stripe_refs(trans);
}

static int dev_freespace_init_iter(struct btree_trans *trans, struct bch_dev *ca,
				   struct btree_iter *iter, struct bpos end)
{
	struct bkey hole;
	struct bkey_s_c k = bkey_try(bch2_get_key_or_hole(iter, end, &hole));

	if (k.k->type) {
		/*
		 * We process live keys in the alloc btree one at a
		 * time:
		 */
		struct bch_alloc_v4 a_convert;
		const struct bch_alloc_v4 *a = bch2_alloc_to_v4(k, &a_convert);

		try(bch2_bucket_do_freespace_index(trans, ca, k, a, true));
		try(bch2_trans_commit(trans, NULL, NULL, BCH_TRANS_COMMIT_no_enospc));

		bch2_btree_iter_advance(iter);
	} else {
		struct bkey_i *freespace = errptr_try(bch2_trans_kmalloc(trans, sizeof(*freespace)));

		bkey_init(&freespace->k);
		freespace->k.type	= KEY_TYPE_set;
		freespace->k.p		= k.k->p;
		freespace->k.size	= k.k->size;

		try(bch2_btree_insert_trans(trans, BTREE_ID_freespace, freespace, 0));
		try(bch2_trans_commit(trans, NULL, NULL, BCH_TRANS_COMMIT_no_enospc));

		bch2_btree_iter_set_pos(iter, k.k->p);
	}

	return 0;
}

int bch2_dev_freespace_init(struct bch_fs *c, struct bch_dev *ca,
			    u64 bucket_start, u64 bucket_end)
{
	struct bpos end = POS(ca->dev_idx, bucket_end);
	unsigned long last_updated = jiffies;

	BUG_ON(bucket_start > bucket_end);
	BUG_ON(bucket_end > ca->mi.nbuckets);

	CLASS(btree_trans, trans)(c);
	CLASS(btree_iter, iter)(trans, BTREE_ID_alloc,
		POS(ca->dev_idx, max_t(u64, ca->mi.first_bucket, bucket_start)),
		BTREE_ITER_prefetch);
	/*
	 * Scan the alloc btree for every bucket on @ca, and add buckets to the
	 * freespace/need_discard/need_gc_gens btrees as needed:
	 */
	while (bkey_lt(iter.pos, end)) {
		if (time_after(jiffies, last_updated + HZ * 10)) {
			bch_info_dev(ca, "%s: currently at %llu/%llu",
				     __func__, iter.pos.offset, ca->mi.nbuckets);
			last_updated = jiffies;
		}

		try(lockrestart_do(trans, dev_freespace_init_iter(trans, ca, &iter, end)));
	}

	scoped_guard(mutex_noio, &c->sb_lock) {
		struct bch_member *m = bch2_members_v2_get_mut(c->disk_sb.sb, ca->dev_idx);
		SET_BCH_MEMBER_FREESPACE_INITIALIZED(m, true);
	}

	return 0;
}

int bch2_fs_freespace_init(struct bch_fs *c)
{
	if (c->sb.features & BIT_ULL(BCH_FEATURE_small_image))
		return 0;

	/*
	 * We can crash during the device add path, so we need to check this on
	 * every mount:
	 */

	bool doing_init = false;
	for_each_member_device(c, ca) {
		if (ca->mi.freespace_initialized)
			continue;

		if (!doing_init) {
			bch_info(c, "initializing freespace");
			doing_init = true;
		}

		try(bch2_dev_freespace_init(c, ca, 0, ca->mi.nbuckets));
	}

	if (doing_init) {
		guard(mutex_noio)(&c->sb_lock);
		bch2_write_super(c);
		bch_verbose(c, "done initializing freespace");
	}

	return 0;
}
