diff --git a/fs/btrfs/ctree.c b/fs/btrfs/ctree.c
index e1557e5..984f474 100644
--- a/fs/btrfs/ctree.c
+++ b/fs/btrfs/ctree.c
@@ -280,7 +280,6 @@
 		*last_ret = search_start;
 		if (parent_level == 1)
 			btrfs_clear_buffer_defrag(tmp);
-		btrfs_set_buffer_defrag_done(tmp);
 		free_extent_buffer(tmp);
 	}
 	return err;
diff --git a/fs/btrfs/tree-defrag.c b/fs/btrfs/tree-defrag.c
index d31c21a..f86eccf 100644
--- a/fs/btrfs/tree-defrag.c
+++ b/fs/btrfs/tree-defrag.c
@@ -87,8 +87,8 @@
 		if (cache_only) {
 			next = btrfs_find_tree_block(root, bytenr,
 					   btrfs_level_size(root, *level - 1));
-			/* FIXME, test for defrag */
-			if (!next || !btrfs_buffer_uptodate(next)) {
+			if (!next || !btrfs_buffer_uptodate(next) ||
+			    !btrfs_buffer_defrag(next)) {
 				free_extent_buffer(next);
 				path->slots[*level]++;
 				continue;
@@ -147,7 +147,8 @@
 			root->defrag_level = i;
 			return 0;
 		} else {
-			if (*level > 1 && path->nodes[*level] != root->node) {
+			if (*level > 1 && path->nodes[*level] != root->node &&
+			    btrfs_buffer_defrag(path->nodes[*level])) {
 				struct extent_buffer *next;
 				u64 last;
 				int ret;
@@ -168,7 +169,6 @@
 			}
 
 			btrfs_clear_buffer_defrag(path->nodes[*level]);
-			btrfs_clear_buffer_defrag_done(path->nodes[*level]);
 			free_extent_buffer(path->nodes[*level]);
 			path->nodes[*level] = NULL;
 			*level = i + 1;