Merge "COW partition creator uses DmSnapshotCowSizeCalculator"
diff --git a/fs_mgr/libsnapshot/partition_cow_creator.cpp b/fs_mgr/libsnapshot/partition_cow_creator.cpp
index eedc1cd..77315b4 100644
--- a/fs_mgr/libsnapshot/partition_cow_creator.cpp
+++ b/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 @@
 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 @@
     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 @@
     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 @@
     // 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);
diff --git a/fs_mgr/libsnapshot/partition_cow_creator.h b/fs_mgr/libsnapshot/partition_cow_creator.h
index 8888f78..d3d186b 100644
--- a/fs_mgr/libsnapshot/partition_cow_creator.h
+++ b/fs_mgr/libsnapshot/partition_cow_creator.h
@@ -60,7 +60,7 @@
 
   private:
     bool HasExtent(Partition* p, Extent* e);
-    std::optional<uint64_t> GetCowSize(uint64_t snapshot_size);
+    uint64_t GetCowSize();
 };
 
 }  // namespace snapshot
diff --git a/fs_mgr/libsnapshot/partition_cow_creator_test.cpp b/fs_mgr/libsnapshot/partition_cow_creator_test.cpp
index f683f5b..eae6c35 100644
--- a/fs_mgr/libsnapshot/partition_cow_creator_test.cpp
+++ b/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 @@
     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;
diff --git a/fs_mgr/libsnapshot/snapshot_test.cpp b/fs_mgr/libsnapshot/snapshot_test.cpp
index 2161e5b..5728582 100644
--- a/fs_mgr/libsnapshot/snapshot_test.cpp
+++ b/fs_mgr/libsnapshot/snapshot_test.cpp
@@ -784,9 +784,17 @@
     }
 
     // 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());
@@ -949,6 +957,13 @@
     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.
diff --git a/fs_mgr/libsnapshot/test_helpers.cpp b/fs_mgr/libsnapshot/test_helpers.cpp
index 539c5c5..312fa3e 100644
--- a/fs_mgr/libsnapshot/test_helpers.cpp
+++ b/fs_mgr/libsnapshot/test_helpers.cpp
@@ -140,5 +140,9 @@
     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
diff --git a/fs_mgr/libsnapshot/test_helpers.h b/fs_mgr/libsnapshot/test_helpers.h
index 5fe701f..0f70afe 100644
--- a/fs_mgr/libsnapshot/test_helpers.h
+++ b/fs_mgr/libsnapshot/test_helpers.h
@@ -141,5 +141,8 @@
 // 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