5 files changed, 32 insertions, 6 deletions

diff --git a/fs/bcachefs/btree_cache.c b/fs/bcachefs/btree_cache.c
index 76cad6109297..9b331c319acc 100644
--- a/fs/bcachefs/btree_cache.c
+++ b/fs/bcachefs/btree_cache.c
@@ -119,8 +119,7 @@ static struct btree *__btree_node_mem_alloc(struct bch_fs *c, gfp_t gfp)
 		return NULL;
 
 	bkey_btree_ptr_init(&b->key);
-	six_lock_init(&b->c.lock);
-	lockdep_set_novalidate_class(&b->c.lock);
+	bch2_btree_lock_init(&b->c);
 	INIT_LIST_HEAD(&b->list);
 	INIT_LIST_HEAD(&b->write_blocked);
 	b->byte_order = ilog2(btree_bytes(c));
diff --git a/fs/bcachefs/btree_io.c b/fs/bcachefs/btree_io.c
index 7ffdce97214e..358a285c3bcf 100644
--- a/fs/bcachefs/btree_io.c
+++ b/fs/bcachefs/btree_io.c
@@ -33,6 +33,8 @@ void bch2_btree_node_io_unlock(struct btree *b)
 
 void bch2_btree_node_io_lock(struct btree *b)
 {
+	bch2_assert_btree_nodes_not_locked();
+
 	wait_on_bit_lock_io(&b->flags, BTREE_NODE_write_in_flight,
 			    TASK_UNINTERRUPTIBLE);
 }
@@ -51,12 +53,16 @@ void __bch2_btree_node_wait_on_write(struct btree *b)
 
 void bch2_btree_node_wait_on_read(struct btree *b)
 {
+	bch2_assert_btree_nodes_not_locked();
+
 	wait_on_bit_io(&b->flags, BTREE_NODE_read_in_flight,
 		       TASK_UNINTERRUPTIBLE);
 }
 
 void bch2_btree_node_wait_on_write(struct btree *b)
 {
+	bch2_assert_btree_nodes_not_locked();
+
 	wait_on_bit_io(&b->flags, BTREE_NODE_write_in_flight,
 		       TASK_UNINTERRUPTIBLE);
 }
diff --git a/fs/bcachefs/btree_key_cache.c b/fs/bcachefs/btree_key_cache.c
index 67db6b9d8e10..1e692c0a2f3a 100644
--- a/fs/bcachefs/btree_key_cache.c
+++ b/fs/bcachefs/btree_key_cache.c
@@ -282,8 +282,7 @@ bkey_cached_alloc(struct btree_trans *trans, struct btree_path *path,
 		return NULL;
 init:
 	INIT_LIST_HEAD(&ck->list);
-	__six_lock_init(&ck->c.lock, "b->c.lock", &bch2_btree_node_lock_key);
-	lockdep_set_novalidate_class(&ck->c.lock);
+	bch2_btree_lock_init(&ck->c);
 	if (pcpu_readers)
 		six_lock_pcpu_alloc(&ck->c.lock);
 
diff --git a/fs/bcachefs/btree_locking.c b/fs/bcachefs/btree_locking.c
index 14a0614af436..9e097ab668a6 100644
--- a/fs/bcachefs/btree_locking.c
+++ b/fs/bcachefs/btree_locking.c
@@ -4,7 +4,23 @@
 #include "btree_locking.h"
 #include "btree_types.h"
 
-struct lock_class_key bch2_btree_node_lock_key;
+static struct lock_class_key bch2_btree_node_lock_key;
+
+void bch2_btree_lock_init(struct btree_bkey_cached_common *b)
+{
+	__six_lock_init(&b->lock, "b->c.lock", &bch2_btree_node_lock_key);
+	lockdep_set_novalidate_class(&b->lock);
+}
+
+#ifdef CONFIG_LOCKDEP
+void bch2_assert_btree_nodes_not_locked(void)
+{
+#if 0
+	//Re-enable when lock_class_is_held() is merged:
+	BUG_ON(lock_class_is_held(&bch2_btree_node_lock_key));
+#endif
+}
+#endif
 
 /* Btree node locking: */
 
diff --git a/fs/bcachefs/btree_locking.h b/fs/bcachefs/btree_locking.h
index 30c89daa5009..76aac49966fe 100644
--- a/fs/bcachefs/btree_locking.h
+++ b/fs/bcachefs/btree_locking.h
@@ -13,7 +13,13 @@
 #include "btree_iter.h"
 #include "six.h"
 
-extern struct lock_class_key bch2_btree_node_lock_key;
+void bch2_btree_lock_init(struct btree_bkey_cached_common *);
+
+#ifdef CONFIG_LOCKDEP
+void bch2_assert_btree_nodes_not_locked(void);
+#else
+static inline void bch2_assert_btree_nodes_not_locked(void) {}
+#endif
 
 static inline bool is_btree_node(struct btree_path *path, unsigned l)
 {