Estimate COW image size during OTA generation am: 7a2657510c am: 268f043480 am: 499fee83f9 am: d14f0c42c1

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

Change-Id: I25cce2d7293632fa3441de4d1999f365af41e6ca
diff --git a/Android.bp b/Android.bp
index d5213a8..ed704ff 100644
--- a/Android.bp
+++ b/Android.bp
@@ -533,6 +533,7 @@
         "payload_generator/block_mapping.cc",
         "payload_generator/boot_img_filesystem.cc",
         "payload_generator/bzip.cc",
+        "payload_generator/cow_size_estimator.cc",
         "payload_generator/deflate_utils.cc",
         "payload_generator/delta_diff_generator.cc",
         "payload_generator/delta_diff_utils.cc",
diff --git a/BUILD.gn b/BUILD.gn
index 30ba977..6e282f5 100644
--- a/BUILD.gn
+++ b/BUILD.gn
@@ -344,6 +344,7 @@
     "payload_generator/block_mapping.cc",
     "payload_generator/boot_img_filesystem_stub.cc",
     "payload_generator/bzip.cc",
+    "payload_generator/cow_size_estimator_stub.cc",
     "payload_generator/deflate_utils.cc",
     "payload_generator/delta_diff_generator.cc",
     "payload_generator/delta_diff_utils.cc",
diff --git a/payload_consumer/delta_performer_unittest.cc b/payload_consumer/delta_performer_unittest.cc
index 2f308c0..74f17d2 100644
--- a/payload_consumer/delta_performer_unittest.cc
+++ b/payload_consumer/delta_performer_unittest.cc
@@ -228,13 +228,13 @@
     new_part.path = "/dev/zero";
     new_part.size = 1234;
 
-    payload.AddPartition(*old_part, new_part, aops, {});
+    payload.AddPartition(*old_part, new_part, aops, {}, 0);
 
     // 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, {}, {});
+    payload.AddPartition(*old_part, new_part, {}, {}, 0);
 
     ScopedTempFile payload_file("Payload-XXXXXX");
     string private_key =
diff --git a/payload_consumer/vabc_partition_writer.cc b/payload_consumer/vabc_partition_writer.cc
index 9e4d9b8..73bf413 100644
--- a/payload_consumer/vabc_partition_writer.cc
+++ b/payload_consumer/vabc_partition_writer.cc
@@ -24,6 +24,7 @@
 #include "update_engine/common/cow_operation_convert.h"
 #include "update_engine/common/utils.h"
 #include "update_engine/payload_consumer/extent_writer.h"
+#include "update_engine/payload_consumer/file_descriptor.h"
 #include "update_engine/payload_consumer/install_plan.h"
 #include "update_engine/payload_consumer/partition_writer.h"
 #include "update_engine/payload_consumer/snapshot_extent_writer.h"
@@ -90,35 +91,44 @@
   // TODO(zhangkelvin) Rewrite this in C++20 coroutine once that's available.
   auto converted = ConvertToCowOperations(partition_update_.operations(),
                                           partition_update_.merge_operations());
-  std::vector<uint8_t> buffer(block_size_);
+
+  WriteAllCowOps(block_size_, converted, cow_writer_.get(), source_fd_);
+  // Emit label 0 to mark end of all SOURCE_COPY operations
+  cow_writer_->AddLabel(0);
+  TEST_AND_RETURN_FALSE(
+      prefs_->SetInt64(kPrefsUpdateStatePartitionNextOperation, 0));
+  return true;
+}
+
+bool VABCPartitionWriter::WriteAllCowOps(
+    size_t block_size,
+    const std::vector<CowOperation>& converted,
+    android::snapshot::ICowWriter* cow_writer,
+    FileDescriptorPtr source_fd) {
+  std::vector<uint8_t> buffer(block_size);
 
   for (const auto& cow_op : converted) {
     switch (cow_op.op) {
       case CowOperation::CowCopy:
         TEST_AND_RETURN_FALSE(
-            cow_writer_->AddCopy(cow_op.dst_block, cow_op.src_block));
+            cow_writer->AddCopy(cow_op.dst_block, cow_op.src_block));
         break;
       case CowOperation::CowReplace:
         ssize_t bytes_read = 0;
-        TEST_AND_RETURN_FALSE(utils::PReadAll(source_fd_,
+        TEST_AND_RETURN_FALSE(utils::PReadAll(source_fd,
                                               buffer.data(),
-                                              block_size_,
-                                              cow_op.src_block * block_size_,
+                                              block_size,
+                                              cow_op.src_block * block_size,
                                               &bytes_read));
-        if (bytes_read <= 0 || static_cast<size_t>(bytes_read) != block_size_) {
+        if (bytes_read <= 0 || static_cast<size_t>(bytes_read) != block_size) {
           LOG(ERROR) << "source_fd->Read failed: " << bytes_read;
           return false;
         }
-        TEST_AND_RETURN_FALSE(cow_writer_->AddRawBlocks(
-            cow_op.dst_block, buffer.data(), block_size_));
+        TEST_AND_RETURN_FALSE(cow_writer->AddRawBlocks(
+            cow_op.dst_block, buffer.data(), block_size));
         break;
     }
   }
-
-  // Emit label 0 to mark end of all SOURCE_COPY operations
-  cow_writer_->AddLabel(0);
-  TEST_AND_RETURN_FALSE(
-      prefs_->SetInt64(kPrefsUpdateStatePartitionNextOperation, 0));
   return true;
 }
 
diff --git a/payload_consumer/vabc_partition_writer.h b/payload_consumer/vabc_partition_writer.h
index 3fc97ce..ddade70 100644
--- a/payload_consumer/vabc_partition_writer.h
+++ b/payload_consumer/vabc_partition_writer.h
@@ -18,9 +18,11 @@
 #define UPDATE_ENGINE_VABC_PARTITION_WRITER_H_
 
 #include <memory>
+#include <vector>
 
 #include <libsnapshot/snapshot_writer.h>
 
+#include "update_engine/common/cow_operation_convert.h"
 #include "update_engine/payload_consumer/install_plan.h"
 #include "update_engine/payload_consumer/partition_writer.h"
 
@@ -44,6 +46,11 @@
   [[nodiscard]] bool Flush() override;
   void CheckpointUpdateProgress(size_t next_op_index) override;
 
+  static bool WriteAllCowOps(size_t block_size,
+                             const std::vector<CowOperation>& converted,
+                             android::snapshot::ICowWriter* cow_writer,
+                             FileDescriptorPtr source_fd);
+
  private:
   std::unique_ptr<android::snapshot::ISnapshotWriter> cow_writer_;
 };
diff --git a/payload_generator/cow_size_estimator.cc b/payload_generator/cow_size_estimator.cc
new file mode 100644
index 0000000..662a3ca
--- /dev/null
+++ b/payload_generator/cow_size_estimator.cc
@@ -0,0 +1,110 @@
+//
+// Copyright (C) 2020 The Android Open Source Project
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+//
+
+#include "update_engine/payload_generator/cow_size_estimator.h"
+
+#include <utility>
+#include <vector>
+
+#include <libsnapshot/cow_writer.h>
+
+#include "android-base/unique_fd.h"
+#include "update_engine/common/cow_operation_convert.h"
+#include "update_engine/payload_consumer/vabc_partition_writer.h"
+#include "update_engine/update_metadata.pb.h"
+
+namespace chromeos_update_engine {
+using android::snapshot::CowWriter;
+
+void PerformReplaceOp(const InstallOperation& op,
+                      CowWriter* writer,
+                      FileDescriptorPtr target_fd,
+                      size_t block_size) {
+  std::vector<unsigned char> buffer;
+  for (const auto& extent : op.dst_extents()) {
+    buffer.resize(extent.num_blocks() * block_size);
+    // No need to read from payload.bin then decompress, just read from target
+    // directly.
+    ssize_t bytes_read = 0;
+    auto success = utils::PReadAll(target_fd,
+                                   buffer.data(),
+                                   buffer.size(),
+                                   extent.start_block() * block_size,
+                                   &bytes_read);
+    CHECK(success);
+    CHECK_EQ(static_cast<size_t>(bytes_read), buffer.size());
+    writer->AddRawBlocks(extent.start_block(), buffer.data(), buffer.size());
+  }
+}
+
+void PerformZeroOp(const InstallOperation& op,
+                   CowWriter* writer,
+                   size_t block_size) {
+  for (const auto& extent : op.dst_extents()) {
+    writer->AddZeroBlocks(extent.start_block(), extent.num_blocks());
+  }
+}
+
+size_t EstimateCowSize(
+    FileDescriptorPtr source_fd,
+    FileDescriptorPtr target_fd,
+    const google::protobuf::RepeatedPtrField<InstallOperation>& operations,
+    const google::protobuf::RepeatedPtrField<CowMergeOperation>&
+        merge_operations,
+    size_t block_size) {
+  android::snapshot::CowWriter cow_writer{
+      {.block_size = static_cast<uint32_t>(block_size), .compression = "gz"}};
+  // CowWriter treats -1 as special value, will discard all the data but still
+  // reports Cow size. Good for estimation purposes
+  cow_writer.Initialize(android::base::borrowed_fd{-1});
+
+  const auto converted = ConvertToCowOperations(operations, merge_operations);
+  VABCPartitionWriter::WriteAllCowOps(
+      block_size, converted, &cow_writer, source_fd);
+  cow_writer.AddLabel(0);
+  for (const auto& op : operations) {
+    switch (op.type()) {
+      case InstallOperation::REPLACE:
+      case InstallOperation::REPLACE_BZ:
+      case InstallOperation::REPLACE_XZ:
+        PerformReplaceOp(op, &cow_writer, target_fd, block_size);
+        break;
+      case InstallOperation::ZERO:
+      case InstallOperation::DISCARD:
+        PerformZeroOp(op, &cow_writer, block_size);
+        break;
+      case InstallOperation::SOURCE_COPY:
+      case InstallOperation::MOVE:
+        // Already handeled by WriteAllCowOps,
+        break;
+      case InstallOperation::SOURCE_BSDIFF:
+      case InstallOperation::BROTLI_BSDIFF:
+      case InstallOperation::PUFFDIFF:
+      case InstallOperation::BSDIFF:
+        // We might do something special by adding CowBsdiff to CowWriter.
+        // For now proceed the same way as normal REPLACE operation.
+        PerformReplaceOp(op, &cow_writer, target_fd, block_size);
+        break;
+    }
+    // Arbitrary label number, we won't be resuming use these labels here.
+    // They are emitted just to keep size estimates accurate. As update_engine
+    // emits 1 label for every op.
+    cow_writer.AddLabel(2);
+  }
+  // TODO(zhangkelvin) Take FEC extents into account once VABC stabilizes
+  return cow_writer.GetCowSize();
+}
+}  // namespace chromeos_update_engine
diff --git a/payload_generator/cow_size_estimator.h b/payload_generator/cow_size_estimator.h
new file mode 100644
index 0000000..cba89b5
--- /dev/null
+++ b/payload_generator/cow_size_estimator.h
@@ -0,0 +1,36 @@
+//
+// Copyright (C) 2020 The Android Open Source Project
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+//
+#include <cstddef>
+
+#include <update_engine/update_metadata.pb.h>
+
+#include "update_engine/payload_consumer/file_descriptor.h"
+
+namespace chromeos_update_engine {
+// Given file descriptor to the source image, target image, and list of
+// operations, estimate the size of COW image if the operations are applied on
+// Virtual AB Compression enabled device. This is intended to be used by update
+// 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.
+size_t EstimateCowSize(
+    FileDescriptorPtr source_fd,
+    FileDescriptorPtr target_fd,
+    const google::protobuf::RepeatedPtrField<InstallOperation>& operations,
+    const google::protobuf::RepeatedPtrField<CowMergeOperation>&
+        merge_operations,
+    size_t block_size);
+
+}  // namespace chromeos_update_engine
diff --git a/payload_generator/cow_size_estimator_stub.cc b/payload_generator/cow_size_estimator_stub.cc
new file mode 100644
index 0000000..9d94d63
--- /dev/null
+++ b/payload_generator/cow_size_estimator_stub.cc
@@ -0,0 +1,31 @@
+//
+// Copyright (C) 2020 The Android Open Source Project
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+//
+
+#include "update_engine/payload_generator/cow_size_estimator.h"
+
+namespace chromeos_update_engine {
+
+size_t EstimateCowSize(
+    FileDescriptorPtr source_fd,
+    FileDescriptorPtr target_fd,
+    const google::protobuf::RepeatedPtrField<InstallOperation>& operations,
+    const google::protobuf::RepeatedPtrField<CowMergeOperation>&
+        merge_operations,
+    size_t block_size) {
+  return 0;
+}
+
+}  // namespace chromeos_update_engine
diff --git a/payload_generator/delta_diff_generator.cc b/payload_generator/delta_diff_generator.cc
index ff8b0da..47c92e0 100644
--- a/payload_generator/delta_diff_generator.cc
+++ b/payload_generator/delta_diff_generator.cc
@@ -33,14 +33,17 @@
 
 #include "update_engine/common/utils.h"
 #include "update_engine/payload_consumer/delta_performer.h"
+#include "update_engine/payload_consumer/file_descriptor.h"
 #include "update_engine/payload_consumer/payload_constants.h"
 #include "update_engine/payload_generator/ab_generator.h"
 #include "update_engine/payload_generator/annotated_operation.h"
 #include "update_engine/payload_generator/blob_file_writer.h"
+#include "update_engine/payload_generator/cow_size_estimator.h"
 #include "update_engine/payload_generator/delta_diff_utils.h"
 #include "update_engine/payload_generator/full_update_generator.h"
 #include "update_engine/payload_generator/merge_sequence_generator.h"
 #include "update_engine/payload_generator/payload_file.h"
+#include "update_engine/update_metadata.pb.h"
 
 using std::string;
 using std::unique_ptr;
@@ -53,6 +56,18 @@
 const size_t kBlockSize = 4096;  // bytes
 
 class PartitionProcessor : public base::DelegateSimpleThread::Delegate {
+  bool IsDynamicPartition(const std::string& partition_name) {
+    for (const auto& group :
+         config_.target.dynamic_partition_metadata->groups()) {
+      const auto& names = group.partition_names();
+      if (std::find(names.begin(), names.end(), partition_name) !=
+          names.end()) {
+        return true;
+      }
+    }
+    return false;
+  }
+
  public:
   explicit PartitionProcessor(
       const PayloadGenerationConfig& config,
@@ -61,6 +76,7 @@
       BlobFileWriter* file_writer,
       std::vector<AnnotatedOperation>* aops,
       std::vector<CowMergeOperation>* cow_merge_sequence,
+      size_t* cow_size,
       std::unique_ptr<chromeos_update_engine::OperationsGenerator> strategy)
       : config_(config),
         old_part_(old_part),
@@ -68,11 +84,13 @@
         file_writer_(file_writer),
         aops_(aops),
         cow_merge_sequence_(cow_merge_sequence),
+        cow_size_(cow_size),
         strategy_(std::move(strategy)) {}
   PartitionProcessor(PartitionProcessor&&) noexcept = default;
+
   void Run() override {
     LOG(INFO) << "Started an async task to process partition "
-              << old_part_.name;
+              << new_part_.name;
     bool success = strategy_->GenerateOperations(
         config_, old_part_, new_part_, file_writer_, aops_);
     if (!success) {
@@ -85,13 +103,38 @@
     bool snapshot_enabled =
         config_.target.dynamic_partition_metadata &&
         config_.target.dynamic_partition_metadata->snapshot_enabled();
-    if (old_part_.path.empty() || !snapshot_enabled) {
+    if (!snapshot_enabled || !IsDynamicPartition(new_part_.name)) {
       return;
     }
-    auto generator = MergeSequenceGenerator::Create(*aops_);
-    if (!generator || !generator->Generate(cow_merge_sequence_)) {
-      LOG(FATAL) << "Failed to generate merge sequence";
+    if (!old_part_.path.empty()) {
+      auto generator = MergeSequenceGenerator::Create(*aops_);
+      if (!generator || !generator->Generate(cow_merge_sequence_)) {
+        LOG(FATAL) << "Failed to generate merge sequence";
+      }
     }
+
+    LOG(INFO) << "Estimating COW size for partition: " << new_part_.name;
+    // Need the contents of source/target image bytes when doing
+    // dry run.
+    FileDescriptorPtr source_fd{new EintrSafeFileDescriptor()};
+    source_fd->Open(old_part_.path.c_str(), O_RDONLY);
+
+    auto target_fd = std::make_unique<EintrSafeFileDescriptor>();
+    target_fd->Open(new_part_.path.c_str(), O_RDONLY);
+
+    google::protobuf::RepeatedPtrField<InstallOperation> operations;
+
+    for (const AnnotatedOperation& aop : *aops_) {
+      *operations.Add() = aop.op;
+    }
+    *cow_size_ = EstimateCowSize(
+        source_fd,
+        std::move(target_fd),
+        operations,
+        {cow_merge_sequence_->begin(), cow_merge_sequence_->end()},
+        config_.block_size);
+    LOG(INFO) << "Estimated COW size for partition: " << new_part_.name << " "
+              << *cow_size_;
   }
 
  private:
@@ -101,6 +144,7 @@
   BlobFileWriter* file_writer_;
   std::vector<AnnotatedOperation>* aops_;
   std::vector<CowMergeOperation>* cow_merge_sequence_;
+  size_t* cow_size_;
   std::unique_ptr<chromeos_update_engine::OperationsGenerator> strategy_;
   DISALLOW_COPY_AND_ASSIGN(PartitionProcessor);
 };
@@ -130,8 +174,12 @@
     PartitionConfig empty_part("");
     std::vector<std::vector<AnnotatedOperation>> all_aops;
     all_aops.resize(config.target.partitions.size());
+
     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<PartitionProcessor> partition_tasks{};
     auto thread_count = std::min<int>(diff_utils::GetMaxThreads(),
                                       config.target.partitions.size());
@@ -163,6 +211,7 @@
                                                    &blob_file,
                                                    &all_aops[i],
                                                    &all_merge_sequences[i],
+                                                   &all_cow_sizes[i],
                                                    std::move(strategy)));
     }
     thread_pool.Start();
@@ -179,7 +228,8 @@
           payload.AddPartition(old_part,
                                new_part,
                                std::move(all_aops[i]),
-                               std::move(all_merge_sequences[i])));
+                               std::move(all_merge_sequences[i]),
+                               all_cow_sizes[i]));
     }
   }
   data_file.CloseFd();
diff --git a/payload_generator/payload_file.cc b/payload_generator/payload_file.cc
index 74423d1..33c0749 100644
--- a/payload_generator/payload_file.cc
+++ b/payload_generator/payload_file.cc
@@ -80,8 +80,10 @@
 bool PayloadFile::AddPartition(const PartitionConfig& old_conf,
                                const PartitionConfig& new_conf,
                                vector<AnnotatedOperation> aops,
-                               vector<CowMergeOperation> merge_sequence) {
+                               vector<CowMergeOperation> merge_sequence,
+                               size_t cow_size) {
   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);
@@ -129,6 +131,9 @@
     if (!part.version.empty()) {
       partition->set_version(part.version);
     }
+    if (part.cow_size > 0) {
+      partition->set_estimate_cow_size(part.cow_size);
+    }
     if (part.postinstall.run) {
       partition->set_run_postinstall(true);
       if (!part.postinstall.path.empty())
diff --git a/payload_generator/payload_file.h b/payload_generator/payload_file.h
index 8b17956..3a45793 100644
--- a/payload_generator/payload_file.h
+++ b/payload_generator/payload_file.h
@@ -44,7 +44,8 @@
   bool AddPartition(const PartitionConfig& old_conf,
                     const PartitionConfig& new_conf,
                     std::vector<AnnotatedOperation> aops,
-                    std::vector<CowMergeOperation> merge_sequence);
+                    std::vector<CowMergeOperation> merge_sequence,
+                    size_t cow_size);
 
   // 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
@@ -100,6 +101,7 @@
     VerityConfig verity;
     // Per partition timestamp.
     std::string version;
+    size_t cow_size;
   };
 
   std::vector<Partition> part_vec_;
diff --git a/payload_generator/payload_properties_unittest.cc b/payload_generator/payload_properties_unittest.cc
index ed936ff..0ff364f 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, {});
+    payload.AddPartition(old_part, new_part, aops, {}, 0);
 
     uint64_t metadata_size;
     EXPECT_TRUE(payload.WritePayload(
diff --git a/scripts/payload_info.py b/scripts/payload_info.py
index 7625ee8..8343d21 100755
--- a/scripts/payload_info.py
+++ b/scripts/payload_info.py
@@ -75,8 +75,11 @@
       DisplayValue('  Number of "%s" ops' % partition.partition_name,
                    len(partition.operations))
     for partition in manifest.partitions:
-      DisplayValue("Timestamp for " +
+      DisplayValue("  Timestamp for " +
                    partition.partition_name, partition.version)
+    for partition in manifest.partitions:
+      DisplayValue("  COW Size for " +
+                   partition.partition_name, partition.estimate_cow_size)
     DisplayValue('Block size', manifest.block_size)
     DisplayValue('Minor version', manifest.minor_version)