update_engine: use cow_size_info

Updating update_engine to use a struct called CowSizeInfo which contains
extra parameters needed in v3 of cow format. These fields should be set
in delta archive manifest and passed to snapshot.cpp which will then
use the fields to correctly configure the cow with a correctly sized
operation buffer space.

Test: ota on cuttlefish
Change-Id: I8306bf3e4bbc3fcd8f2e31be116fa0623022c227

6 files changed, 30 insertions, 19 deletions

diff --git a/payload_generator/cow_size_estimator.cc b/payload_generator/cow_size_estimator.cc
index 81cfd0bd..22ef6096 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 @@ bool CowDryRun(
   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 @@ size_t EstimateCowSize(
     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 @@ size_t EstimateCowSize(
                   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 3bf6880f..c9fab09b 100644
--- a/payload_generator/cow_size_estimator.h
+++ b/payload_generator/cow_size_estimator.h
@@ -28,7 +28,7 @@ namespace chromeos_update_engine {
 // 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 cfec1b94..b1dc6931 100644
--- a/payload_generator/delta_diff_generator.cc
+++ b/payload_generator/delta_diff_generator.cc
@@ -75,7 +75,7 @@ class PartitionProcessor : public base::DelegateSimpleThread::Delegate {
       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 @@ class PartitionProcessor : public base::DelegateSimpleThread::Delegate {
         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 @@ class PartitionProcessor : public base::DelegateSimpleThread::Delegate {
       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 @@ class PartitionProcessor : public base::DelegateSimpleThread::Delegate {
         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 @@ class PartitionProcessor : public base::DelegateSimpleThread::Delegate {
   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 @@ bool GenerateUpdatePayloadFile(const PayloadGenerationConfig& config,
     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 @@ bool GenerateUpdatePayloadFile(const PayloadGenerationConfig& config,
                                                    &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 @@ bool GenerateUpdatePayloadFile(const PayloadGenerationConfig& config,
                                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 f404c790..8f5b826b 100644
--- a/payload_generator/payload_file.cc
+++ b/payload_generator/payload_file.cc
@@ -98,15 +98,15 @@ bool PayloadFile::AddPartition(const PartitionConfig& old_conf,
                                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 @@ bool PayloadFile::WritePayload(const string& payload_file,
     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 86bf2439..cb414945 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 @@ class PayloadFile {
                     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 @@ class PayloadFile {
     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 f7ac7430..b4bfb810 100644
--- a/payload_generator/payload_properties_unittest.cc
+++ b/payload_generator/payload_properties_unittest.cc
@@ -88,7 +88,7 @@ class PayloadPropertiesTest : public ::testing::Test {
     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(