bcachefs: bch2_trans_kmalloc no longer allocates memory with btree locks held

When allocating memory, gfp flags should generally be

 - GFP_NOWAIT|__GFP_NOWARN if btree locks are held
 - GFP_NOFS if in the IO path or otherwise holding resources needed for
   IO submission
 - GFP_KERNEL otherwise
Signed-off-by: Kent Overstreet <kent.overstreet@linux.dev>

fs/bcachefs/btree_iter.c

@@ -2790,6 +2790,7 @@ void *__bch2_trans_kmalloc(struct btree_trans *trans, size_t size)
 	unsigned new_top = trans->mem_top + size;
 	size_t old_bytes = trans->mem_bytes;
 	size_t new_bytes = roundup_pow_of_two(new_top);
+	int ret;
 	void *new_mem;
 	void *p;
 
@@ -2797,7 +2798,11 @@ void *__bch2_trans_kmalloc(struct btree_trans *trans, size_t size)
 
 	WARN_ON_ONCE(new_bytes > BTREE_TRANS_MEM_MAX);
 
-	new_mem = krealloc(trans->mem, new_bytes, GFP_NOFS);
+	new_mem = krealloc(trans->mem, new_bytes, GFP_NOWAIT|__GFP_NOWARN);
+	if (unlikely(!new_mem)) {
+		bch2_trans_unlock(trans);
+
+		new_mem = krealloc(trans->mem, new_bytes, GFP_KERNEL);
 		if (!new_mem && new_bytes <= BTREE_TRANS_MEM_MAX) {
 			new_mem = mempool_alloc(&trans->c->btree_trans_mem_pool, GFP_KERNEL);
 			new_bytes = BTREE_TRANS_MEM_MAX;
@@ -2810,6 +2815,14 @@ void *__bch2_trans_kmalloc(struct btree_trans *trans, size_t size)
 		trans->mem = new_mem;
 		trans->mem_bytes = new_bytes;
 
+		ret = bch2_trans_relock(trans);
+		if (ret)
+			return ERR_PTR(ret);
+	}
+
+	trans->mem = new_mem;
+	trans->mem_bytes = new_bytes;
+
 	if (old_bytes) {
 		trace_and_count(trans->c, trans_restart_mem_realloced, trans, _RET_IP_, new_bytes);
 		return ERR_PTR(btree_trans_restart(trans, BCH_ERR_transaction_restart_mem_realloced));