diff --git a/fs/btrfs/extent_io.c b/fs/btrfs/extent_io.c
index 3721820687d78c3258efd0bbffb46ac3f335bf8a..01a1412458622bb096649fbef0176396be01f5ad 100644
--- a/fs/btrfs/extent_io.c
+++ b/fs/btrfs/extent_io.c
@@ -817,6 +817,7 @@ __set_extent_bit(struct extent_io_tree *tree, u64 start, u64 end,
 		if (err)
 			extent_io_tree_panic(tree, err);
 
+		cache_state(prealloc, cached_state);
 		prealloc = NULL;
 		goto out;
 	}
@@ -1040,9 +1041,10 @@ int convert_extent_bit(struct extent_io_tree *tree, u64 start, u64 end,
 			goto out;
 		}
 		err = insert_state(tree, prealloc, start, end, &bits);
-		prealloc = NULL;
 		if (err)
 			extent_io_tree_panic(tree, err);
+		cache_state(prealloc, cached_state);
+		prealloc = NULL;
 		goto out;
 	}
 	state = rb_entry(node, struct extent_state, rb_node);