update_engine: use cow_size_info am: 09d022d905

Original change: https://android-review.googlesource.com/c/platform/system/update_engine/+/2829130

Change-Id: I087c6300764daeaf7195532db15cf6d3c5b51e2c
Signed-off-by: Automerger Merge Worker <android-build-automerger-merge-worker@system.gserviceaccount.com>
diff --git a/download_action_android_unittest.cc b/download_action_android_unittest.cc
index 968f875..426d7fd 100644
--- a/download_action_android_unittest.cc
+++ b/download_action_android_unittest.cc
@@ -102,8 +102,8 @@
   ftruncate(partition_file.fd(), 4096);
   partition_config.size = 4096;
   partition_config.path = partition_file.path();
-  ASSERT_TRUE(
-      payload_file.AddPartition(partition_config, partition_config, {}, {}, 0));
+  ASSERT_TRUE(payload_file.AddPartition(
+      partition_config, partition_config, {}, {}, {}));
   ScopedTempFile blob_file("Blob-XXXXXX");
   ScopedTempFile manifest_file("Manifest-XXXXXX");
   uint64_t metadata_size;
diff --git a/payload_consumer/delta_performer_unittest.cc b/payload_consumer/delta_performer_unittest.cc
index e90bb7c..04cfaad 100644
--- a/payload_consumer/delta_performer_unittest.cc
+++ b/payload_consumer/delta_performer_unittest.cc
@@ -228,13 +228,14 @@
     new_part.path = "/dev/zero";
     new_part.size = 1234;
 
-    payload.AddPartition(*old_part, new_part, aops, {}, 0);
+    payload.AddPartition(*old_part, new_part, aops, {}, {});
 
     // We include a kernel partition without operations.
     old_part->name = kPartitionNameKernel;
     new_part.name = kPartitionNameKernel;
     new_part.size = 0;
-    payload.AddPartition(*old_part, new_part, {}, {}, 0);
+
+    payload.AddPartition(*old_part, new_part, {}, {}, {});
 
     ScopedTempFile payload_file("Payload-XXXXXX");
     string private_key =
diff --git a/payload_consumer/snapshot_extent_writer_unittest.cc b/payload_consumer/snapshot_extent_writer_unittest.cc
index 057dbda..31ec3cb 100644
--- a/payload_consumer/snapshot_extent_writer_unittest.cc
+++ b/payload_consumer/snapshot_extent_writer_unittest.cc
@@ -100,12 +100,15 @@
   }
 
   // Return number of bytes the cow image occupies on disk.
-  uint64_t GetCowSize() override {
+  uint64_t GetCowSize() const {
     return std::accumulate(
         operations_.begin(), operations_.end(), 0, [](auto&& acc, auto&& op) {
           return acc + op.second.data.size();
         });
   }
+  android::snapshot::CowSizeInfo GetCowSizeInfo() const override {
+    return android::snapshot::CowSizeInfo{GetCowSize(), 0};
+  }
   bool Contains(size_t block) {
     return operations_.find(block) != operations_.end();
   }
diff --git a/payload_generator/cow_size_estimator.cc b/payload_generator/cow_size_estimator.cc
index 81cfd0b..22ef609 100644
--- a/payload_generator/cow_size_estimator.cc
+++ b/payload_generator/cow_size_estimator.cc
@@ -24,6 +24,7 @@
 
 #include <android-base/unique_fd.h>
 #include <libsnapshot/cow_writer.h>
+#include <libsnapshot/cow_format.h>
 
 #include "update_engine/common/utils.h"
 #include "update_engine/payload_consumer/vabc_partition_writer.h"
@@ -135,7 +136,7 @@
   return cow_writer->Finalize();
 }
 
-size_t EstimateCowSize(
+android::snapshot::CowSizeInfo EstimateCowSizeInfo(
     FileDescriptorPtr source_fd,
     FileDescriptorPtr target_fd,
     const google::protobuf::RepeatedPtrField<InstallOperation>& operations,
@@ -147,7 +148,8 @@
     const bool xor_enabled) {
   android::snapshot::CowOptions options{
       .block_size = static_cast<uint32_t>(block_size),
-      .compression = std::move(compression)};
+      .compression = std::move(compression),
+      .max_blocks = (partition_size / block_size)};
   auto cow_writer =
       CreateCowEstimator(android::snapshot::kCowVersionManifest, options);
   CHECK_NE(cow_writer, nullptr) << "Could not create cow estimator";
@@ -159,7 +161,7 @@
                   cow_writer.get(),
                   partition_size,
                   xor_enabled));
-  return cow_writer->GetCowSize();
+  return cow_writer->GetCowSizeInfo();
 }
 
 }  // namespace chromeos_update_engine
diff --git a/payload_generator/cow_size_estimator.h b/payload_generator/cow_size_estimator.h
index 3bf6880..c9fab09 100644
--- a/payload_generator/cow_size_estimator.h
+++ b/payload_generator/cow_size_estimator.h
@@ -28,7 +28,7 @@
 // generators to put an estimate cow size in OTA payload. When installing an OTA
 // update, libsnapshot will take this estimate as a hint to allocate spaces.
 // If |xor_enabled| is true, then |source_fd| must be non-null.
-size_t EstimateCowSize(
+android::snapshot::CowSizeInfo EstimateCowSizeInfo(
     FileDescriptorPtr source_fd,
     FileDescriptorPtr target_fd,
     const google::protobuf::RepeatedPtrField<InstallOperation>& operations,
diff --git a/payload_generator/delta_diff_generator.cc b/payload_generator/delta_diff_generator.cc
index cfec1b9..b1dc693 100644
--- a/payload_generator/delta_diff_generator.cc
+++ b/payload_generator/delta_diff_generator.cc
@@ -75,7 +75,7 @@
       BlobFileWriter* file_writer,
       std::vector<AnnotatedOperation>* aops,
       std::vector<CowMergeOperation>* cow_merge_sequence,
-      size_t* cow_size,
+      android::snapshot::CowSizeInfo* cow_info,
       std::unique_ptr<chromeos_update_engine::OperationsGenerator> strategy)
       : config_(config),
         old_part_(old_part),
@@ -83,7 +83,7 @@
         file_writer_(file_writer),
         aops_(aops),
         cow_merge_sequence_(cow_merge_sequence),
-        cow_size_(cow_size),
+        cow_info_(cow_info),
         strategy_(std::move(strategy)) {}
   PartitionProcessor(PartitionProcessor&&) noexcept = default;
 
@@ -134,7 +134,7 @@
       source_fd->Open(old_part_.path.c_str(), O_RDONLY);
     }
 
-    *cow_size_ = EstimateCowSize(
+    *cow_info_ = EstimateCowSizeInfo(
         std::move(source_fd),
         std::move(target_fd),
         std::move(operations),
@@ -143,8 +143,11 @@
         config_.target.dynamic_partition_metadata->vabc_compression_param(),
         new_part_.size,
         config_.enable_vabc_xor);
+
+    // ops buffer size == 0 for v2 of cow format
     LOG(INFO) << "Estimated COW size for partition: " << new_part_.name << " "
-              << *cow_size_;
+              << cow_info_->cow_size
+              << " ops buffer size: " << cow_info_->op_count_max;
   }
 
  private:
@@ -154,7 +157,7 @@
   BlobFileWriter* file_writer_;
   std::vector<AnnotatedOperation>* aops_;
   std::vector<CowMergeOperation>* cow_merge_sequence_;
-  size_t* cow_size_;
+  android::snapshot::CowSizeInfo* cow_info_;
   std::unique_ptr<chromeos_update_engine::OperationsGenerator> strategy_;
   DISALLOW_COPY_AND_ASSIGN(PartitionProcessor);
 };
@@ -188,7 +191,8 @@
     std::vector<std::vector<CowMergeOperation>> all_merge_sequences;
     all_merge_sequences.resize(config.target.partitions.size());
 
-    std::vector<size_t> all_cow_sizes(config.target.partitions.size(), 0);
+    std::vector<android::snapshot::CowSizeInfo> all_cow_info(
+        config.target.partitions.size());
 
     std::vector<PartitionProcessor> partition_tasks{};
     auto thread_count = std::min<int>(diff_utils::GetMaxThreads(),
@@ -223,7 +227,7 @@
                                                    &blob_file,
                                                    &all_aops[i],
                                                    &all_merge_sequences[i],
-                                                   &all_cow_sizes[i],
+                                                   &all_cow_info[i],
                                                    std::move(strategy)));
     }
     thread_pool.Start();
@@ -241,7 +245,7 @@
                                new_part,
                                std::move(all_aops[i]),
                                std::move(all_merge_sequences[i]),
-                               all_cow_sizes[i]));
+                               all_cow_info[i]));
     }
   }
   data_file.CloseFd();
diff --git a/payload_generator/payload_file.cc b/payload_generator/payload_file.cc
index f404c79..8f5b826 100644
--- a/payload_generator/payload_file.cc
+++ b/payload_generator/payload_file.cc
@@ -98,15 +98,15 @@
                                const PartitionConfig& new_conf,
                                vector<AnnotatedOperation> aops,
                                vector<CowMergeOperation> merge_sequence,
-                               size_t cow_size) {
+                               const android::snapshot::CowSizeInfo& cow_info) {
   Partition part;
-  part.cow_size = cow_size;
   part.name = new_conf.name;
   part.aops = std::move(aops);
   part.cow_merge_sequence = std::move(merge_sequence);
   part.postinstall = new_conf.postinstall;
   part.verity = new_conf.verity;
   part.version = new_conf.version;
+  part.cow_info = cow_info;
   // Initialize the PartitionInfo objects if present.
   if (!old_conf.path.empty())
     TEST_AND_RETURN_FALSE(
@@ -148,8 +148,11 @@
     if (!part.version.empty()) {
       partition->set_version(part.version);
     }
-    if (part.cow_size > 0) {
-      partition->set_estimate_cow_size(part.cow_size);
+    if (part.cow_info.cow_size > 0) {
+      partition->set_estimate_cow_size(part.cow_info.cow_size);
+    }
+    if (part.cow_info.op_count_max > 0) {
+      partition->set_estimate_op_count_max(part.cow_info.op_count_max);
     }
     if (part.postinstall.run) {
       partition->set_run_postinstall(true);
diff --git a/payload_generator/payload_file.h b/payload_generator/payload_file.h
index 86bf243..cb41494 100644
--- a/payload_generator/payload_file.h
+++ b/payload_generator/payload_file.h
@@ -23,6 +23,8 @@
 #include <brillo/secure_blob.h>
 #include <gtest/gtest_prod.h>  // for FRIEND_TEST
 
+#include <libsnapshot/cow_writer.h>
+
 #include "update_engine/payload_generator/annotated_operation.h"
 #include "update_engine/payload_generator/payload_generation_config.h"
 #include "update_engine/update_metadata.pb.h"
@@ -45,7 +47,7 @@
                     const PartitionConfig& new_conf,
                     std::vector<AnnotatedOperation> aops,
                     std::vector<CowMergeOperation> merge_sequence,
-                    size_t cow_size);
+                    const android::snapshot::CowSizeInfo& cow_info);
 
   // Write the payload to the |payload_file| file. The operations reference
   // blobs in the |data_blobs_path| file and the blobs will be reordered in the
@@ -108,7 +110,7 @@
     VerityConfig verity;
     // Per partition timestamp.
     std::string version;
-    size_t cow_size;
+    android::snapshot::CowSizeInfo cow_info;
   };
 
   std::vector<Partition> part_vec_;
diff --git a/payload_generator/payload_properties_unittest.cc b/payload_generator/payload_properties_unittest.cc
index f7ac743..b4bfb81 100644
--- a/payload_generator/payload_properties_unittest.cc
+++ b/payload_generator/payload_properties_unittest.cc
@@ -88,7 +88,7 @@
     EXPECT_TRUE(strategy->GenerateOperations(
         config, old_part, new_part, &blob_file_writer, &aops));
 
-    payload.AddPartition(old_part, new_part, aops, {}, 0);
+    payload.AddPartition(old_part, new_part, aops, {}, {});
 
     uint64_t metadata_size;
     EXPECT_TRUE(payload.WritePayload(
diff --git a/update_metadata.proto b/update_metadata.proto
index 3881464..535f386 100644
--- a/update_metadata.proto
+++ b/update_metadata.proto
@@ -316,6 +316,10 @@
   // as a hint. If set to 0, libsnapshot should use alternative
   // methods for estimating size.
   optional uint64 estimate_cow_size = 19;
+
+  // Information about the cow used by Cow Writer to specify
+  // number of cow operations to be written
+  optional uint64 estimate_op_count_max = 20;
 }
 
 message DynamicPartitionGroup {