Snap for 6227479 from 132f621ff8ca4ddd60e897c3145a1acbb00eda14 to rvc-release

Change-Id: I41f96d6eae851db67eabdc6997fcbf021a30866a

2 files changed, 131 insertions, 25 deletions

diff --git a/contrib/android/basefs_allocator.c b/contrib/android/basefs_allocator.c
index 32305319..4f9f5c15 100644
--- a/contrib/android/basefs_allocator.c
+++ b/contrib/android/basefs_allocator.c
@@ -9,6 +9,8 @@
 struct base_fs_allocator {
 	struct ext2fs_hashmap *entries;
 	struct basefs_entry *cur_entry;
+	/* The next expected logical block to allocate for cur_entry. */
+	blk64_t next_lblk;
 	/* Blocks which are definitely owned by a single inode in BaseFS. */
 	ext2fs_block_bitmap exclusive_block_map;
 	/* Blocks which are available to the first inode that requests it. */
@@ -235,29 +237,66 @@ out:
 }
 
 /* Try and acquire the next usable block from the Base FS map. */
-static int get_next_block(ext2_filsys fs, struct base_fs_allocator *allocator,
-			  struct block_range_list* list, blk64_t *ret)
+static errcode_t get_next_block(ext2_filsys fs, struct base_fs_allocator *allocator,
+				struct block_range_list* list, blk64_t *ret)
 {
 	blk64_t block;
 	ext2fs_block_bitmap exclusive_map = allocator->exclusive_block_map;
 	ext2fs_block_bitmap dedup_map = allocator->dedup_block_map;
 
-	while (list->head) {
+	if (!list->head)
+		return EXT2_ET_BLOCK_ALLOC_FAIL;
+
+	block = consume_next_block(list);
+	if (block >= ext2fs_blocks_count(fs->super))
+		return EXT2_ET_BLOCK_ALLOC_FAIL;
+	if (ext2fs_test_block_bitmap2(exclusive_map, block)) {
+		ext2fs_unmark_block_bitmap2(exclusive_map, block);
+		*ret = block;
+		return 0;
+	}
+	if (ext2fs_test_block_bitmap2(dedup_map, block)) {
+		ext2fs_unmark_block_bitmap2(dedup_map, block);
+		*ret = block;
+		return 0;
+	}
+	return EXT2_ET_BLOCK_ALLOC_FAIL;
+}
+
+/*
+ * BaseFS lists blocks in logical block order. However, the allocator hook is
+ * only called if a block needs to be allocated. In the case of a deduplicated
+ * block, or a hole, the hook is not invoked. This means the next block
+ * allocation request will be out of sequence. For example, consider if BaseFS
+ * specifies the following (0 being a hole):
+ *     1 2 3 0 4 5
+ *
+ * If the new file has a hole at logical block 0, we could accidentally
+ * shift the entire expected block list as follows:
+ *     0 1 2 0 3 4
+ *
+ * To account for this, we track the next expected logical block in the
+ * allocator. If the current request is for a later logical block, we skip and
+ * free the intermediate physical blocks that would have been allocated. This
+ * ensures the original block assignment is respected.
+ */
+static void skip_blocks(ext2_filsys fs, struct base_fs_allocator *allocator,
+			struct blk_alloc_ctx *ctx)
+{
+	blk64_t block;
+	struct block_range_list *list = &allocator->cur_entry->blocks;
+	ext2fs_block_bitmap exclusive_map = allocator->exclusive_block_map;
+
+	while (list->head && allocator->next_lblk < ctx->lblk) {
 		block = consume_next_block(list);
 		if (block >= ext2fs_blocks_count(fs->super))
 			continue;
 		if (ext2fs_test_block_bitmap2(exclusive_map, block)) {
 			ext2fs_unmark_block_bitmap2(exclusive_map, block);
-			*ret = block;
-			return 0;
-		}
-		if (ext2fs_test_block_bitmap2(dedup_map, block)) {
-			ext2fs_unmark_block_bitmap2(dedup_map, block);
-			*ret = block;
-			return 0;
+			ext2fs_unmark_block_bitmap2(fs->block_map, block);
 		}
+		allocator->next_lblk++;
 	}
-	return -1;
 }
 
 static errcode_t basefs_block_allocator(ext2_filsys fs, blk64_t goal,
@@ -269,6 +308,10 @@ static errcode_t basefs_block_allocator(ext2_filsys fs, blk64_t goal,
 	ext2fs_block_bitmap dedup_map = allocator->dedup_block_map;
 
 	if (e && ctx && (ctx->flags & BLOCK_ALLOC_DATA)) {
+		if (allocator->next_lblk < ctx->lblk)
+			skip_blocks(fs, allocator, ctx);
+		allocator->next_lblk = ctx->lblk + 1;
+
 		if (!get_next_block(fs, allocator, &e->blocks, ret))
 			return 0;
 	}
@@ -278,17 +321,28 @@ static errcode_t basefs_block_allocator(ext2_filsys fs, blk64_t goal,
 		ext2fs_mark_block_bitmap2(fs->block_map, *ret);
 		return 0;
 	}
-	if (retval == EXT2_ET_BLOCK_ALLOC_FAIL) {
-		/* Try to steal a block from the dedup pool. */
-		retval = ext2fs_find_first_set_block_bitmap2(dedup_map,
-			fs->super->s_first_data_block,
-			ext2fs_blocks_count(fs->super) - 1, ret);
-		if (!retval) {
-			ext2fs_unmark_block_bitmap2(dedup_map, *ret);
+	if (retval != EXT2_ET_BLOCK_ALLOC_FAIL)
+		return retval;
+
+	/* Try to steal a block from the dedup pool. */
+	retval = ext2fs_find_first_set_block_bitmap2(dedup_map,
+		fs->super->s_first_data_block,
+		ext2fs_blocks_count(fs->super) - 1, ret);
+	if (!retval) {
+		ext2fs_unmark_block_bitmap2(dedup_map, *ret);
+		return 0;
+	}
+
+	/*
+	 * As a last resort, take any block from our file's list. This
+	 * risks bloating the diff, but means we are more likely to
+	 * successfully build an image.
+	 */
+	while (e->blocks.head) {
+		if (!get_next_block(fs, allocator, &e->blocks, ret))
 			return 0;
-		}
 	}
-	return retval;
+	return EXT2_ET_BLOCK_ALLOC_FAIL;
 }
 
 void base_fs_alloc_cleanup(ext2_filsys fs)
@@ -308,10 +362,12 @@ errcode_t base_fs_alloc_set_target(ext2_filsys fs, const char *target_path,
 	if (mode != S_IFREG)
 		return 0;
 
-	if (allocator)
+	if (allocator) {
 		allocator->cur_entry = ext2fs_hashmap_lookup(allocator->entries,
 						      target_path,
 						      strlen(target_path));
+		allocator->next_lblk = 0;
+	}
 	return 0;
 }
 
diff --git a/contrib/android/fsmap.c b/contrib/android/fsmap.c
index 36adb7f0..9ee8472d 100644
--- a/contrib/android/fsmap.c
+++ b/contrib/android/fsmap.c
@@ -23,11 +23,51 @@ static int walk_block(ext2_filsys fs  EXT2FS_ATTR((unused)), blk64_t *blocknr,
 	return format->add_block(fs, *blocknr, blockcnt < 0, format->private);
 }
 
+static errcode_t ino_iter_extents(ext2_filsys fs, ext2_ino_t ino,
+				  ext2_extent_handle_t extents,
+				  struct walk_ext_priv_data *pdata)
+{
+	blk64_t block;
+	errcode_t retval;
+	blk64_t next_lblk = 0;
+	int op = EXT2_EXTENT_ROOT;
+	struct ext2fs_extent extent;
+	struct fsmap_format *format = pdata->format;
+
+	for (;;) {
+		retval = ext2fs_extent_get(extents, op, &extent);
+		if (retval)
+			break;
+
+		op = EXT2_EXTENT_NEXT;
+
+		if ((extent.e_flags & EXT2_EXTENT_FLAGS_SECOND_VISIT) ||
+		    !(extent.e_flags & EXT2_EXTENT_FLAGS_LEAF))
+			continue;
+
+		for (; next_lblk < extent.e_lblk; next_lblk++)
+			format->add_block(fs, 0, 0, format->private);
+
+		block = extent.e_pblk;
+		for (; next_lblk < extent.e_lblk + extent.e_len; next_lblk++)
+			format->add_block(fs, block++, 0, format->private);
+	}
+
+	if (retval == EXT2_ET_EXTENT_NO_NEXT)
+		retval = 0;
+	if (retval) {
+		com_err(__func__, retval, ("getting extents of ino \"%u\""),
+			ino);
+	}
+	return retval;
+}
+
 static errcode_t ino_iter_blocks(ext2_filsys fs, ext2_ino_t ino,
 				 struct walk_ext_priv_data *pdata)
 {
 	errcode_t retval;
 	struct ext2_inode inode;
+	ext2_extent_handle_t extents;
 	struct fsmap_format *format = pdata->format;
 
 	retval = ext2fs_read_inode(fs, ino, &inode);
@@ -38,10 +78,20 @@ static errcode_t ino_iter_blocks(ext2_filsys fs, ext2_ino_t ino,
 		return format->inline_data(&(inode.i_block[0]),
 					   format->private);
 
-	retval = ext2fs_block_iterate3(fs, ino, 0, NULL, walk_block, pdata);
-	if (retval)
-		com_err(__func__, retval, _("listing blocks of ino \"%u\""),
-			ino);
+	retval = ext2fs_extent_open(fs, ino, &extents);
+	if (retval == EXT2_ET_INODE_NOT_EXTENT) {
+		retval = ext2fs_block_iterate3(fs, ino, BLOCK_FLAG_READ_ONLY,
+			NULL, walk_block, pdata);
+		if (retval) {
+			com_err(__func__, retval, _("listing blocks of ino \"%u\""),
+				ino);
+		}
+		return retval;
+	}
+
+	retval = ino_iter_extents(fs, ino, extents, pdata);
+
+	ext2fs_extent_free(extents);
 	return retval;
 }