COW partition creator uses DmSnapshotCowSizeCalculator

Instead of using heuristics on the calculation of the COW size, use the
class representing the internal structure of the COW device as generated
by the Linux kernel.
This simplifies the forecast of the required COW device size depending
on the write operations performed on the associated snapshot device.

Created also an additional test case for the COW size calculation.

Change-Id: If147001544988716430b36a4c75dfa7f34b4c8d9
Bug: 140835698
Test: libsnapshot_test
Signed-off-by: Alessio Balsini <balsini@google.com>

fs_mgr/libsnapshot/partition_cow_creator.cpp

@@ -17,8 +17,9 @@
 #include <math.h>
 
 #include <android-base/logging.h>
-
 #include <android/snapshot/snapshot.pb.h>
+
+#include "dm_snapshot_internals.h"
 #include "utility.h"
 
 using android::dm::kSectorSize;
@@ -33,13 +34,6 @@ using RepeatedPtrField = google::protobuf::RepeatedPtrField<T>;
 namespace android {
 namespace snapshot {
 
-// Round |d| up to a multiple of |block_size|.
-static uint64_t RoundUp(double d, uint64_t block_size) {
-    uint64_t ret = ((uint64_t)ceil(d) + block_size - 1) / block_size * block_size;
-    CHECK(ret >= d) << "Can't round " << d << " up to a multiple of " << block_size;
-    return ret;
-}
-
 // Intersect two linear extents. If no intersection, return an extent with length 0.
 static std::unique_ptr<Extent> Intersect(Extent* target_extent, Extent* existing_extent) {
     // Convert target_extent and existing_extent to linear extents. Zero extents
@@ -68,33 +62,58 @@ bool PartitionCowCreator::HasExtent(Partition* p, Extent* e) {
     return false;
 }
 
-std::optional<uint64_t> PartitionCowCreator::GetCowSize(uint64_t snapshot_size) {
-    // TODO: Use |operations|. to determine a minimum COW size.
-    // kCowEstimateFactor is good for prototyping but we can't use that in production.
-    static constexpr double kCowEstimateFactor = 1.05;
-    auto cow_size = RoundUp(snapshot_size * kCowEstimateFactor, kDefaultBlockSize);
-    return cow_size;
+uint64_t PartitionCowCreator::GetCowSize() {
+    // WARNING: The origin partition should be READ-ONLY
+    const uint64_t logical_block_size = current_metadata->logical_block_size();
+    const unsigned int sectors_per_block = logical_block_size / kSectorSize;
+    DmSnapCowSizeCalculator sc(kSectorSize, kSnapshotChunkSize);
+
+    if (operations == nullptr) return sc.cow_size_bytes();
+
+    for (const auto& iop : *operations) {
+        for (const auto& de : iop.dst_extents()) {
+            // Skip if no blocks are written
+            if (de.num_blocks() == 0) continue;
+
+            // Flag all the blocks that were written
+            const auto block_boundary = de.start_block() + de.num_blocks();
+            for (auto b = de.start_block(); b < block_boundary; ++b) {
+                for (unsigned int s = 0; s < sectors_per_block; ++s) {
+                    const auto sector_id = b * sectors_per_block + s;
+                    sc.WriteSector(sector_id);
+                }
+            }
+        }
+    }
+
+    return sc.cow_size_bytes();
 }
 
 std::optional<PartitionCowCreator::Return> PartitionCowCreator::Run() {
     CHECK(current_metadata->GetBlockDevicePartitionName(0) == LP_METADATA_DEFAULT_PARTITION_NAME &&
           target_metadata->GetBlockDevicePartitionName(0) == LP_METADATA_DEFAULT_PARTITION_NAME);
 
-    uint64_t logical_block_size = current_metadata->logical_block_size();
+    const uint64_t logical_block_size = current_metadata->logical_block_size();
     CHECK(logical_block_size != 0 && !(logical_block_size & (logical_block_size - 1)))
             << "logical_block_size is not power of 2";
 
     Return ret;
     ret.snapshot_status.set_name(target_partition->name());
     ret.snapshot_status.set_device_size(target_partition->size());
-
-    // TODO(b/141889746): Optimize by using a smaller snapshot. Some ranges in target_partition
-    // may be written directly.
     ret.snapshot_status.set_snapshot_size(target_partition->size());
 
-    auto cow_size = GetCowSize(ret.snapshot_status.snapshot_size());
-    if (!cow_size.has_value()) return std::nullopt;
-
+    // Being the COW partition virtual, its size doesn't affect the storage
+    // memory that will be occupied by the target.
+    // The actual storage space is affected by the COW file, whose size depends
+    // on the chunks that diverged between |current| and |target|.
+    // If the |target| partition is bigger than |current|, the data that is
+    // modified outside of |current| can be written directly to |current|.
+    // This because the data that will be written outside of |current| would
+    // not invalidate any useful information of |current|, thus:
+    // - if the snapshot is accepted for merge, this data would be already at
+    // the right place and should not be copied;
+    // - in the unfortunate case of the snapshot to be discarded, the regions
+    // modified by this data can be set as free regions and reused.
     // Compute regions that are free in both current and target metadata. These are the regions
     // we can use for COW partition.
     auto target_free_regions = target_metadata->GetFreeRegions();
@@ -102,13 +121,15 @@ std::optional<PartitionCowCreator::Return> PartitionCowCreator::Run() {
     auto free_regions = Interval::Intersect(target_free_regions, current_free_regions);
     uint64_t free_region_length = 0;
     for (const auto& interval : free_regions) {
-        free_region_length += interval.length() * kSectorSize;
+        free_region_length += interval.length();
     }
+    free_region_length *= kSectorSize;
 
     LOG(INFO) << "Remaining free space for COW: " << free_region_length << " bytes";
+    auto cow_size = GetCowSize();
 
     // Compute the COW partition size.
-    uint64_t cow_partition_size = std::min(*cow_size, free_region_length);
+    uint64_t cow_partition_size = std::min(cow_size, free_region_length);
     // Round it down to the nearest logical block. Logical partitions must be a multiple
     // of logical blocks.
     cow_partition_size &= ~(logical_block_size - 1);
@@ -116,8 +137,7 @@ std::optional<PartitionCowCreator::Return> PartitionCowCreator::Run() {
     // Assign cow_partition_usable_regions to indicate what regions should the COW partition uses.
     ret.cow_partition_usable_regions = std::move(free_regions);
 
-    // The rest of the COW space is allocated on ImageManager.
-    uint64_t cow_file_size = (*cow_size) - ret.snapshot_status.cow_partition_size();
+    auto cow_file_size = cow_size - cow_partition_size;
     // Round it up to the nearest sector.
     cow_file_size += kSectorSize - 1;
     cow_file_size &= ~(kSectorSize - 1);

fs_mgr/libsnapshot/partition_cow_creator.h

@@ -60,7 +60,7 @@ struct PartitionCowCreator {
 
   private:
     bool HasExtent(Partition* p, Extent* e);
-    std::optional<uint64_t> GetCowSize(uint64_t snapshot_size);
+    uint64_t GetCowSize();
 };
 
 }  // namespace snapshot

fs_mgr/libsnapshot/partition_cow_creator_test.cpp

@@ -14,12 +14,14 @@
 
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
+#include <libdm/dm.h>
 #include <liblp/builder.h>
 #include <liblp/property_fetcher.h>
 
 #include "dm_snapshot_internals.h"
 #include "partition_cow_creator.h"
 #include "test_helpers.h"
+#include "utility.h"
 
 using namespace android::fs_mgr;
 
@@ -100,6 +102,90 @@ TEST_F(PartitionCowCreatorTest, Holes) {
     ASSERT_TRUE(ret.has_value());
 }
 
+TEST_F(PartitionCowCreatorTest, CowSize) {
+    using InstallOperation = chromeos_update_engine::InstallOperation;
+    using RepeatedInstallOperationPtr = google::protobuf::RepeatedPtrField<InstallOperation>;
+    using Extent = chromeos_update_engine::Extent;
+
+    constexpr uint64_t initial_size = 50_MiB;
+    constexpr uint64_t final_size = 40_MiB;
+
+    auto builder_a = MetadataBuilder::New(initial_size, 1_KiB, 2);
+    ASSERT_NE(builder_a, nullptr);
+    auto system_a = builder_a->AddPartition("system_a", LP_PARTITION_ATTR_READONLY);
+    ASSERT_NE(system_a, nullptr);
+    ASSERT_TRUE(builder_a->ResizePartition(system_a, final_size));
+
+    auto builder_b = MetadataBuilder::New(initial_size, 1_KiB, 2);
+    ASSERT_NE(builder_b, nullptr);
+    auto system_b = builder_b->AddPartition("system_b", LP_PARTITION_ATTR_READONLY);
+    ASSERT_NE(system_b, nullptr);
+    ASSERT_TRUE(builder_b->ResizePartition(system_b, final_size));
+
+    const uint64_t block_size = builder_b->logical_block_size();
+    const uint64_t chunk_size = kSnapshotChunkSize * dm::kSectorSize;
+    ASSERT_EQ(chunk_size, block_size);
+
+    auto cow_device_size = [](const std::vector<InstallOperation>& iopv, MetadataBuilder* builder_a,
+                              MetadataBuilder* builder_b, Partition* system_b) {
+        RepeatedInstallOperationPtr riop(iopv.begin(), iopv.end());
+        PartitionCowCreator creator{.target_metadata = builder_b,
+                                    .target_suffix = "_b",
+                                    .target_partition = system_b,
+                                    .current_metadata = builder_a,
+                                    .current_suffix = "_a",
+                                    .operations = &riop};
+
+        auto ret = creator.Run();
+
+        if (ret.has_value()) {
+            return ret->snapshot_status.cow_file_size() + ret->snapshot_status.cow_partition_size();
+        }
+        return std::numeric_limits<uint64_t>::max();
+    };
+
+    std::vector<InstallOperation> iopv;
+    InstallOperation iop;
+    Extent* e;
+
+    // No data written, no operations performed
+    ASSERT_EQ(2 * chunk_size, cow_device_size(iopv, builder_a.get(), builder_b.get(), system_b));
+
+    // No data written
+    e = iop.add_dst_extents();
+    e->set_start_block(0);
+    e->set_num_blocks(0);
+    iopv.push_back(iop);
+    ASSERT_EQ(2 * chunk_size, cow_device_size(iopv, builder_a.get(), builder_b.get(), system_b));
+
+    e = iop.add_dst_extents();
+    e->set_start_block(1);
+    e->set_num_blocks(0);
+    iopv.push_back(iop);
+    ASSERT_EQ(2 * chunk_size, cow_device_size(iopv, builder_a.get(), builder_b.get(), system_b));
+
+    // Fill the first block
+    e = iop.add_dst_extents();
+    e->set_start_block(0);
+    e->set_num_blocks(1);
+    iopv.push_back(iop);
+    ASSERT_EQ(3 * chunk_size, cow_device_size(iopv, builder_a.get(), builder_b.get(), system_b));
+
+    // Fill the second block
+    e = iop.add_dst_extents();
+    e->set_start_block(1);
+    e->set_num_blocks(1);
+    iopv.push_back(iop);
+    ASSERT_EQ(4 * chunk_size, cow_device_size(iopv, builder_a.get(), builder_b.get(), system_b));
+
+    // Jump to 5th block and write 2
+    e = iop.add_dst_extents();
+    e->set_start_block(5);
+    e->set_num_blocks(2);
+    iopv.push_back(iop);
+    ASSERT_EQ(6 * chunk_size, cow_device_size(iopv, builder_a.get(), builder_b.get(), system_b));
+}
+
 TEST(DmSnapshotInternals, CowSizeCalculator) {
     DmSnapCowSizeCalculator cc(512, 8);
     unsigned long int b;

fs_mgr/libsnapshot/snapshot_test.cpp

@@ -778,9 +778,17 @@ TEST_F(SnapshotUpdateTest, FullUpdateFlow) {
     }
 
     // Grow all partitions.
-    SetSize(sys_, 3788_KiB);
-    SetSize(vnd_, 3788_KiB);
-    SetSize(prd_, 3788_KiB);
+    constexpr uint64_t partition_size = 3788_KiB;
+    SetSize(sys_, partition_size);
+    SetSize(vnd_, partition_size);
+    SetSize(prd_, partition_size);
+
+    // Create fake install operations to grow the COW device size.
+    for (auto& partition : {sys_, vnd_, prd_}) {
+        auto e = partition->add_operations()->add_dst_extents();
+        e->set_start_block(0);
+        e->set_num_blocks(GetSize(partition) / manifest_.block_size());
+    }
 
     // Execute the update.
     ASSERT_TRUE(sm->BeginUpdate());
@@ -943,6 +951,13 @@ TEST_F(SnapshotUpdateTest, TestRollback) {
     ASSERT_TRUE(sm->BeginUpdate());
     ASSERT_TRUE(sm->UnmapUpdateSnapshot("sys_b"));
 
+    // Create fake install operations to grow the COW device size.
+    for (auto& partition : {sys_, vnd_, prd_}) {
+        auto e = partition->add_operations()->add_dst_extents();
+        e->set_start_block(0);
+        e->set_num_blocks(GetSize(partition) / manifest_.block_size());
+    }
+
     ASSERT_TRUE(sm->CreateUpdateSnapshots(manifest_));
 
     // Write some data to target partitions.

fs_mgr/libsnapshot/test_helpers.cpp

@@ -140,5 +140,9 @@ void SetSize(PartitionUpdate* partition_update, uint64_t size) {
     partition_update->mutable_new_partition_info()->set_size(size);
 }
 
+uint64_t GetSize(PartitionUpdate* partition_update) {
+    return partition_update->mutable_new_partition_info()->size();
+}
+
 }  // namespace snapshot
 }  // namespace android

fs_mgr/libsnapshot/test_helpers.h

@@ -138,5 +138,8 @@ AssertionResult FillFakeMetadata(MetadataBuilder* builder, const DeltaArchiveMan
 // In the update package metadata, set a partition with the given size.
 void SetSize(PartitionUpdate* partition_update, uint64_t size);
 
+// Get partition size from update package metadata.
+uint64_t GetSize(PartitionUpdate* partition_update);
+
 }  // namespace snapshot
 }  // namespace android