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/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 {