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
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(