libsnapshot: Use the COW size from the update manifest.

When Virtual A/B Compression is enabled, the manifest contains the
predicted COW size. Use this instead of the algorithm based on the
kernel COW format.

Bug: 168554689
Test: vts_libsnapshot_test
Change-Id: I545679b4834957ff80a930d91cb44afbadebb66c
diff --git a/fs_mgr/libsnapshot/Android.bp b/fs_mgr/libsnapshot/Android.bp
index f1b0031..910911e 100644
--- a/fs_mgr/libsnapshot/Android.bp
+++ b/fs_mgr/libsnapshot/Android.bp
@@ -46,6 +46,7 @@
     header_libs: [
         "libchrome",
         "libfiemap_headers",
+        "libstorage_literals_headers",
         "libupdate_engine_headers",
     ],
     export_static_lib_headers: [
diff --git a/fs_mgr/libsnapshot/partition_cow_creator.cpp b/fs_mgr/libsnapshot/partition_cow_creator.cpp
index 0df5664..0c81146 100644
--- a/fs_mgr/libsnapshot/partition_cow_creator.cpp
+++ b/fs_mgr/libsnapshot/partition_cow_creator.cpp
@@ -18,6 +18,7 @@
 
 #include <android-base/logging.h>
 #include <android/snapshot/snapshot.pb.h>
+#include <storage_literals/storage_literals.h>
 
 #include "dm_snapshot_internals.h"
 #include "utility.h"
@@ -34,6 +35,8 @@
 namespace android {
 namespace snapshot {
 
+using namespace android::storage_literals;
+
 // 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
@@ -138,6 +141,17 @@
 }
 
 uint64_t PartitionCowCreator::GetCowSize() {
+    if (compression_enabled) {
+        if (update == nullptr || !update->has_estimate_cow_size()) {
+            LOG(ERROR) << "Update manifest does not include a COW size";
+            return 0;
+        }
+
+        // Add an extra 2MB of wiggle room for any minor differences in labels/metadata
+        // that might come up.
+        return update->estimate_cow_size() + 2_MiB;
+    }
+
     // 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;
@@ -149,9 +163,9 @@
         WriteExtent(&sc, de, sectors_per_block);
     }
 
-    if (operations == nullptr) return sc.cow_size_bytes();
+    if (update == nullptr) return sc.cow_size_bytes();
 
-    for (const auto& iop : *operations) {
+    for (const auto& iop : update->operations()) {
         const InstallOperation* written_op = &iop;
         InstallOperation buf;
         // Do not allocate space for extents that are going to be skipped
@@ -213,6 +227,9 @@
 
     LOG(INFO) << "Remaining free space for COW: " << free_region_length << " bytes";
     auto cow_size = GetCowSize();
+    if (!cow_size) {
+        return {};
+    }
 
     // Compute the COW partition size.
     uint64_t cow_partition_size = std::min(cow_size, free_region_length);
diff --git a/fs_mgr/libsnapshot/partition_cow_creator.h b/fs_mgr/libsnapshot/partition_cow_creator.h
index 699f9a1..64d186b 100644
--- a/fs_mgr/libsnapshot/partition_cow_creator.h
+++ b/fs_mgr/libsnapshot/partition_cow_creator.h
@@ -36,6 +36,7 @@
     using MetadataBuilder = android::fs_mgr::MetadataBuilder;
     using Partition = android::fs_mgr::Partition;
     using InstallOperation = chromeos_update_engine::InstallOperation;
+    using PartitionUpdate = chromeos_update_engine::PartitionUpdate;
     template <typename T>
     using RepeatedPtrField = google::protobuf::RepeatedPtrField<T>;
 
@@ -50,11 +51,13 @@
     MetadataBuilder* current_metadata = nullptr;
     // The suffix of the current slot.
     std::string current_suffix;
-    // List of operations to be applied on the partition.
-    const RepeatedPtrField<InstallOperation>* operations = nullptr;
+    // Partition information from the OTA manifest.
+    const PartitionUpdate* update = nullptr;
     // Extra extents that are going to be invalidated during the update
     // process.
     std::vector<ChromeOSExtent> extra_extents = {};
+    // True if compression is enabled.
+    bool compression_enabled = false;
 
     struct Return {
         SnapshotStatus snapshot_status;
diff --git a/fs_mgr/libsnapshot/partition_cow_creator_test.cpp b/fs_mgr/libsnapshot/partition_cow_creator_test.cpp
index 2970222..e4b476f 100644
--- a/fs_mgr/libsnapshot/partition_cow_creator_test.cpp
+++ b/fs_mgr/libsnapshot/partition_cow_creator_test.cpp
@@ -145,13 +145,15 @@
 
     auto cow_device_size = [](const std::vector<InstallOperation>& iopv, MetadataBuilder* builder_a,
                               MetadataBuilder* builder_b, Partition* system_b) {
-        RepeatedInstallOperationPtr riop(iopv.begin(), iopv.end());
+        PartitionUpdate update;
+        *update.mutable_operations() = RepeatedInstallOperationPtr(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};
+                                    .update = &update};
 
         auto ret = creator.Run();
 
@@ -218,7 +220,7 @@
                                 .target_partition = system_b,
                                 .current_metadata = builder_a.get(),
                                 .current_suffix = "_a",
-                                .operations = nullptr};
+                                .update = nullptr};
 
     auto ret = creator.Run();
 
@@ -228,6 +230,58 @@
     ASSERT_EQ(0u, ret->snapshot_status.cow_partition_size());
 }
 
+TEST_F(PartitionCowCreatorTest, CompressionEnabled) {
+    constexpr uint64_t super_size = 1_MiB;
+    auto builder_a = MetadataBuilder::New(super_size, 1_KiB, 2);
+    ASSERT_NE(builder_a, nullptr);
+
+    auto builder_b = MetadataBuilder::New(super_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, 128_KiB));
+
+    PartitionUpdate update;
+    update.set_estimate_cow_size(256_KiB);
+
+    PartitionCowCreator creator{.target_metadata = builder_b.get(),
+                                .target_suffix = "_b",
+                                .target_partition = system_b,
+                                .current_metadata = builder_a.get(),
+                                .current_suffix = "_a",
+                                .compression_enabled = true,
+                                .update = &update};
+
+    auto ret = creator.Run();
+    ASSERT_TRUE(ret.has_value());
+    ASSERT_EQ(ret->snapshot_status.cow_file_size(), 1458176);
+}
+
+TEST_F(PartitionCowCreatorTest, CompressionWithNoManifest) {
+    constexpr uint64_t super_size = 1_MiB;
+    auto builder_a = MetadataBuilder::New(super_size, 1_KiB, 2);
+    ASSERT_NE(builder_a, nullptr);
+
+    auto builder_b = MetadataBuilder::New(super_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, 128_KiB));
+
+    PartitionUpdate update;
+
+    PartitionCowCreator creator{.target_metadata = builder_b.get(),
+                                .target_suffix = "_b",
+                                .target_partition = system_b,
+                                .current_metadata = builder_a.get(),
+                                .current_suffix = "_a",
+                                .compression_enabled = true,
+                                .update = nullptr};
+
+    auto ret = creator.Run();
+    ASSERT_FALSE(ret.has_value());
+}
+
 TEST(DmSnapshotInternals, CowSizeCalculator) {
     SKIP_IF_NON_VIRTUAL_AB();
 
diff --git a/fs_mgr/libsnapshot/snapshot.cpp b/fs_mgr/libsnapshot/snapshot.cpp
index f9bb0dd..793680b 100644
--- a/fs_mgr/libsnapshot/snapshot.cpp
+++ b/fs_mgr/libsnapshot/snapshot.cpp
@@ -75,7 +75,7 @@
 using chromeos_update_engine::DeltaArchiveManifest;
 using chromeos_update_engine::Extent;
 using chromeos_update_engine::FileDescriptor;
-using chromeos_update_engine::InstallOperation;
+using chromeos_update_engine::PartitionUpdate;
 template <typename T>
 using RepeatedPtrField = google::protobuf::RepeatedPtrField<T>;
 using std::chrono::duration_cast;
@@ -116,10 +116,6 @@
     metadata_dir_ = device_->GetMetadataDir();
 }
 
-static inline bool IsCompressionEnabled() {
-    return android::base::GetBoolProperty("ro.virtual_ab.compression.enabled", false);
-}
-
 static std::string GetCowName(const std::string& snapshot_name) {
     return snapshot_name + "-cow";
 }
@@ -2527,8 +2523,9 @@
             .target_partition = nullptr,
             .current_metadata = current_metadata.get(),
             .current_suffix = current_suffix,
-            .operations = nullptr,
+            .update = nullptr,
             .extra_extents = {},
+            .compression_enabled = IsCompressionEnabled(),
     };
 
     auto ret = CreateUpdateSnapshotsInternal(lock.get(), manifest, &cow_creator, &created_devices,
@@ -2572,12 +2569,11 @@
         return Return::Error();
     }
 
-    std::map<std::string, const RepeatedPtrField<InstallOperation>*> install_operation_map;
+    std::map<std::string, const PartitionUpdate*> partition_map;
     std::map<std::string, std::vector<Extent>> extra_extents_map;
     for (const auto& partition_update : manifest.partitions()) {
         auto suffixed_name = partition_update.partition_name() + target_suffix;
-        auto&& [it, inserted] =
-                install_operation_map.emplace(suffixed_name, &partition_update.operations());
+        auto&& [it, inserted] = partition_map.emplace(suffixed_name, &partition_update);
         if (!inserted) {
             LOG(ERROR) << "Duplicated partition " << partition_update.partition_name()
                        << " in update manifest.";
@@ -2595,10 +2591,10 @@
 
     for (auto* target_partition : ListPartitionsWithSuffix(target_metadata, target_suffix)) {
         cow_creator->target_partition = target_partition;
-        cow_creator->operations = nullptr;
-        auto operations_it = install_operation_map.find(target_partition->name());
-        if (operations_it != install_operation_map.end()) {
-            cow_creator->operations = operations_it->second;
+        cow_creator->update = nullptr;
+        auto iter = partition_map.find(target_partition->name());
+        if (iter != partition_map.end()) {
+            cow_creator->update = iter->second;
         } else {
             LOG(INFO) << target_partition->name()
                       << " isn't included in the payload, skipping the cow creation.";
@@ -2614,6 +2610,7 @@
         // Compute the device sizes for the partition.
         auto cow_creator_ret = cow_creator->Run();
         if (!cow_creator_ret.has_value()) {
+            LOG(ERROR) << "PartitionCowCreator returned no value for " << target_partition->name();
             return Return::Error();
         }
 
diff --git a/fs_mgr/libsnapshot/snapshot_test.cpp b/fs_mgr/libsnapshot/snapshot_test.cpp
index 7fc64e5..74558ab 100644
--- a/fs_mgr/libsnapshot/snapshot_test.cpp
+++ b/fs_mgr/libsnapshot/snapshot_test.cpp
@@ -82,7 +82,6 @@
 std::string fake_super;
 
 void MountMetadata();
-bool IsCompressionEnabled();
 
 class SnapshotTest : public ::testing::Test {
   public:
@@ -796,12 +795,15 @@
         group_->add_partition_names("prd");
         sys_ = manifest_.add_partitions();
         sys_->set_partition_name("sys");
+        sys_->set_estimate_cow_size(6_MiB);
         SetSize(sys_, 3_MiB);
         vnd_ = manifest_.add_partitions();
         vnd_->set_partition_name("vnd");
+        vnd_->set_estimate_cow_size(6_MiB);
         SetSize(vnd_, 3_MiB);
         prd_ = manifest_.add_partitions();
         prd_->set_partition_name("prd");
+        prd_->set_estimate_cow_size(6_MiB);
         SetSize(prd_, 3_MiB);
 
         // Initialize source partition metadata using |manifest_|.
@@ -1045,7 +1047,11 @@
     auto tgt = MetadataBuilder::New(*opener_, "super", 1);
     ASSERT_NE(tgt, nullptr);
     ASSERT_NE(nullptr, tgt->FindPartition("sys_b-cow"));
-    ASSERT_NE(nullptr, tgt->FindPartition("vnd_b-cow"));
+    if (IsCompressionEnabled()) {
+        ASSERT_EQ(nullptr, tgt->FindPartition("vnd_b-cow"));
+    } else {
+        ASSERT_NE(nullptr, tgt->FindPartition("vnd_b-cow"));
+    }
     ASSERT_EQ(nullptr, tgt->FindPartition("prd_b-cow"));
 
     // Write some data to target partitions.
@@ -1453,10 +1459,6 @@
     MetadataMountedTest().TearDown();
 }
 
-bool IsCompressionEnabled() {
-    return android::base::GetBoolProperty("ro.virtual_ab.compression.enabled", false);
-}
-
 TEST_F(MetadataMountedTest, Android) {
     auto device = sm->EnsureMetadataMounted();
     EXPECT_NE(nullptr, device);
diff --git a/fs_mgr/libsnapshot/update_engine/update_metadata.proto b/fs_mgr/libsnapshot/update_engine/update_metadata.proto
index 202e39b..dda214e 100644
--- a/fs_mgr/libsnapshot/update_engine/update_metadata.proto
+++ b/fs_mgr/libsnapshot/update_engine/update_metadata.proto
@@ -62,6 +62,7 @@
     repeated InstallOperation operations = 8;
     optional Extent hash_tree_extent = 11;
     optional Extent fec_extent = 15;
+    optional uint64 estimate_cow_size = 19;
 }
 
 message DynamicPartitionGroup {
diff --git a/fs_mgr/libsnapshot/utility.cpp b/fs_mgr/libsnapshot/utility.cpp
index 4cae83a..7342fd4 100644
--- a/fs_mgr/libsnapshot/utility.cpp
+++ b/fs_mgr/libsnapshot/utility.cpp
@@ -22,6 +22,7 @@
 
 #include <android-base/file.h>
 #include <android-base/logging.h>
+#include <android-base/properties.h>
 #include <android-base/strings.h>
 #include <fs_mgr/roots.h>
 
@@ -182,5 +183,9 @@
     new_extent->set_num_blocks(num_blocks);
 }
 
+bool IsCompressionEnabled() {
+    return android::base::GetBoolProperty("ro.virtual_ab.compression.enabled", false);
+}
+
 }  // namespace snapshot
 }  // namespace android
diff --git a/fs_mgr/libsnapshot/utility.h b/fs_mgr/libsnapshot/utility.h
index 482888a..3e6873b 100644
--- a/fs_mgr/libsnapshot/utility.h
+++ b/fs_mgr/libsnapshot/utility.h
@@ -129,5 +129,7 @@
 void AppendExtent(google::protobuf::RepeatedPtrField<chromeos_update_engine::Extent>* extents,
                   uint64_t start_block, uint64_t num_blocks);
 
+bool IsCompressionEnabled();
+
 }  // namespace snapshot
 }  // namespace android