Add CowMergeOperations as a hint for snapshot write

As proposed in http://go/vabc, we want to reduce the cow size
for VAB. One nature apporach is to skip writing the idential
blocks to snapshot; instead we can read from the souce blocks.

Similiar to the non-A/B update schema, we need to compute a
sequence for snapshot merge to avoid the read after write problem.
If there is a circular dependency, we will omit some blocks in the
result sequence to break the cycles. So libsnapshot will write
the raw data of these blocks to cow.

All extents in the CowMergeOperations are subsets of a particular
OTA SOURCE_COPY InstallOperation. Also, these src & ext extents
will be contiguous to improve the libsnapshot read performance
before merge completes, as well as to simplify the sequence
generation.

Bug: 162274240
Test: unittest pass, genertes an OTA
Change-Id: I12c952593d83a8e34a0a6cff5a2066c9103a0d30
diff --git a/Android.bp b/Android.bp
index b6ee476..751e355 100644
--- a/Android.bp
+++ b/Android.bp
@@ -491,6 +491,7 @@
         "payload_generator/extent_utils.cc",
         "payload_generator/full_update_generator.cc",
         "payload_generator/mapfile_filesystem.cc",
+        "payload_generator/merge_sequence_generator.cc",
         "payload_generator/payload_file.cc",
         "payload_generator/payload_generation_config_android.cc",
         "payload_generator/payload_generation_config.cc",
diff --git a/payload_consumer/delta_performer_unittest.cc b/payload_consumer/delta_performer_unittest.cc
index fbd754f..449201c 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, {});
 
     // 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, {}, {});
 
     test_utils::ScopedTempFile payload_file("Payload-XXXXXX");
     string private_key =
diff --git a/payload_generator/delta_diff_generator.cc b/payload_generator/delta_diff_generator.cc
index aa49252..c2b35ee 100644
--- a/payload_generator/delta_diff_generator.cc
+++ b/payload_generator/delta_diff_generator.cc
@@ -39,6 +39,7 @@
 #include "update_engine/payload_generator/blob_file_writer.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"
 
 using std::string;
@@ -59,12 +60,14 @@
       const PartitionConfig& new_part,
       BlobFileWriter* file_writer,
       std::vector<AnnotatedOperation>* aops,
+      std::vector<CowMergeOperation>* cow_merge_sequence,
       std::unique_ptr<chromeos_update_engine::OperationsGenerator> strategy)
       : config_(config),
         old_part_(old_part),
         new_part_(new_part),
         file_writer_(file_writer),
         aops_(aops),
+        cow_merge_sequence_(cow_merge_sequence),
         strategy_(std::move(strategy)) {}
   PartitionProcessor(PartitionProcessor&&) noexcept = default;
   void Run() override {
@@ -78,6 +81,17 @@
       LOG(FATAL) << "GenerateOperations(" << old_part_.name << ", "
                  << new_part_.name << ") failed";
     }
+
+    bool snapshot_enabled =
+        config_.target.dynamic_partition_metadata &&
+        config_.target.dynamic_partition_metadata->snapshot_enabled();
+    if (old_part_.path.empty() || !snapshot_enabled) {
+      return;
+    }
+    auto generator = MergeSequenceGenerator::Create(*aops_);
+    if (!generator || !generator->Generate(cow_merge_sequence_)) {
+      LOG(FATAL) << "Failed to generate merge sequence";
+    }
   }
 
  private:
@@ -86,6 +100,7 @@
   const PartitionConfig& new_part_;
   BlobFileWriter* file_writer_;
   std::vector<AnnotatedOperation>* aops_;
+  std::vector<CowMergeOperation>* cow_merge_sequence_;
   std::unique_ptr<chromeos_update_engine::OperationsGenerator> strategy_;
   DISALLOW_COPY_AND_ASSIGN(PartitionProcessor);
 };
@@ -123,6 +138,8 @@
     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<PartitionProcessor> partition_tasks{};
     auto thread_count = std::min<int>(diff_utils::GetMaxThreads(),
                                       config.target.partitions.size());
@@ -153,6 +170,7 @@
                                                    new_part,
                                                    &blob_file,
                                                    &all_aops[i],
+                                                   &all_merge_sequences[i],
                                                    std::move(strategy)));
     }
     thread_pool.Start();
@@ -166,7 +184,10 @@
           config.is_delta ? config.source.partitions[i] : empty_part;
       const PartitionConfig& new_part = config.target.partitions[i];
       TEST_AND_RETURN_FALSE(
-          payload.AddPartition(old_part, new_part, std::move(all_aops[i])));
+          payload.AddPartition(old_part,
+                               new_part,
+                               std::move(all_aops[i]),
+                               std::move(all_merge_sequences[i])));
     }
   }
 
diff --git a/payload_generator/merge_sequence_generator.cc b/payload_generator/merge_sequence_generator.cc
new file mode 100644
index 0000000..cf73ba3
--- /dev/null
+++ b/payload_generator/merge_sequence_generator.cc
@@ -0,0 +1,43 @@
+//
+// 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/merge_sequence_generator.h"
+
+namespace chromeos_update_engine {
+
+std::unique_ptr<MergeSequenceGenerator> MergeSequenceGenerator::Create(
+    const std::vector<AnnotatedOperation>& aops) {
+  return std::unique_ptr<MergeSequenceGenerator>(
+      new MergeSequenceGenerator({}));
+}
+
+bool MergeSequenceGenerator::FindDependency(
+    std::map<CowMergeOperation, std::set<CowMergeOperation>>* result) const {
+  CHECK(result);
+  return true;
+}
+
+bool MergeSequenceGenerator::Generate(
+    std::vector<CowMergeOperation>* sequence) const {
+  return true;
+}
+
+bool MergeSequenceGenerator::ValidateSequence(
+    const std::vector<CowMergeOperation>& sequence) {
+  return true;
+}
+
+}  // namespace chromeos_update_engine
diff --git a/payload_generator/merge_sequence_generator.h b/payload_generator/merge_sequence_generator.h
new file mode 100644
index 0000000..bfc04d9
--- /dev/null
+++ b/payload_generator/merge_sequence_generator.h
@@ -0,0 +1,63 @@
+//
+// 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.
+//
+
+#ifndef UPDATE_ENGINE_PAYLOAD_GENERATOR_MERGE_SEQUENCE_GENERATOR_H_
+#define UPDATE_ENGINE_PAYLOAD_GENERATOR_MERGE_SEQUENCE_GENERATOR_H_
+
+#include <map>
+#include <memory>
+#include <set>
+#include <utility>
+#include <vector>
+
+#include "update_engine/payload_generator/annotated_operation.h"
+#include "update_engine/payload_generator/extent_ranges.h"
+#include "update_engine/payload_generator/extent_utils.h"
+#include "update_engine/update_metadata.pb.h"
+
+namespace chromeos_update_engine {
+
+// This class takes a list of CowMergeOperations; and sorts them so that no
+// read after write will happen by following the sequence. When there is a
+// cycle, we will omit some operations in the list. Therefore, the result
+// sequence may not contain all blocks in the input list.
+class MergeSequenceGenerator {
+ public:
+  // Creates an object from a list of OTA InstallOperations. Returns nullptr on
+  // failure.
+  static std::unique_ptr<MergeSequenceGenerator> Create(
+      const std::vector<AnnotatedOperation>& aops);
+  // Checks that no read after write happens in the given sequence.
+  static bool ValidateSequence(const std::vector<CowMergeOperation>& sequence);
+
+  // Generates a merge sequence from |operations_|, puts the result in
+  // |sequence|. Returns false on failure.
+  bool Generate(std::vector<CowMergeOperation>* sequence) const;
+
+ private:
+  explicit MergeSequenceGenerator(std::vector<CowMergeOperation> transfers)
+      : operations_(std::move(transfers)) {}
+
+  // For a given merge operation, finds all the operations that should merge
+  // after myself. Put the result in |merge_after|.
+  bool FindDependency(std::map<CowMergeOperation, std::set<CowMergeOperation>>*
+                          merge_after) const;
+  // The list of CowMergeOperations to sort.
+  std::vector<CowMergeOperation> operations_;
+};
+
+}  // namespace chromeos_update_engine
+#endif
diff --git a/payload_generator/payload_file.cc b/payload_generator/payload_file.cc
index 1388f2d..49dff4e 100644
--- a/payload_generator/payload_file.cc
+++ b/payload_generator/payload_file.cc
@@ -86,10 +86,12 @@
 
 bool PayloadFile::AddPartition(const PartitionConfig& old_conf,
                                const PartitionConfig& new_conf,
-                               vector<AnnotatedOperation> aops) {
+                               vector<AnnotatedOperation> aops,
+                               vector<CowMergeOperation> merge_sequence) {
   Partition part;
   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;
@@ -163,6 +165,10 @@
     for (const AnnotatedOperation& aop : part.aops) {
       *partition->add_operations() = aop.op;
     }
+    for (const auto& merge_op : part.cow_merge_sequence) {
+      *partition->add_merge_operations() = merge_op;
+    }
+
     if (part.old_info.has_size() || part.old_info.has_hash())
       *(partition->mutable_old_partition_info()) = part.old_info;
     if (part.new_info.has_size() || part.new_info.has_hash())
diff --git a/payload_generator/payload_file.h b/payload_generator/payload_file.h
index 3dce00f..8b17956 100644
--- a/payload_generator/payload_file.h
+++ b/payload_generator/payload_file.h
@@ -43,7 +43,8 @@
   // reference a blob stored in the file provided to WritePayload().
   bool AddPartition(const PartitionConfig& old_conf,
                     const PartitionConfig& new_conf,
-                    std::vector<AnnotatedOperation> aops);
+                    std::vector<AnnotatedOperation> aops,
+                    std::vector<CowMergeOperation> merge_sequence);
 
   // 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
@@ -90,6 +91,7 @@
 
     // The operations to be performed to this partition.
     std::vector<AnnotatedOperation> aops;
+    std::vector<CowMergeOperation> cow_merge_sequence;
 
     PartitionInfo old_info;
     PartitionInfo new_info;
diff --git a/payload_generator/payload_properties_unittest.cc b/payload_generator/payload_properties_unittest.cc
index db3902c..e0072fc 100644
--- a/payload_generator/payload_properties_unittest.cc
+++ b/payload_generator/payload_properties_unittest.cc
@@ -98,7 +98,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, {});
 
     uint64_t metadata_size;
     EXPECT_TRUE(payload.WritePayload(
diff --git a/update_metadata.proto b/update_metadata.proto
index f79e38b..373ee5e 100644
--- a/update_metadata.proto
+++ b/update_metadata.proto
@@ -225,6 +225,22 @@
   optional bytes src_sha256_hash = 9;
 }
 
+// Hints to VAB snapshot to skip writing some blocks if these blocks are
+// identical to the ones on the source image. The src & dst extents for each
+// CowMergeOperation should be contiguous, and they're a subset of an OTA
+// InstallOperation.
+// During merge time, we need to follow the pre-computed sequence to avoid
+// read after write, similar to the inplace update schema.
+message CowMergeOperation {
+  enum Type {
+    COW_COPY = 0;  // identical blocks
+  }
+  optional Type type = 1;
+
+  optional Extent src_extent = 2;
+  optional Extent dst_extent = 3;
+}
+
 // Describes the update to apply to a single partition.
 message PartitionUpdate {
   // A platform-specific name to identify the partition set being updated. For
@@ -293,6 +309,11 @@
   // as an effort to support partial updates. For most partitions,
   // this is the build timestamp.
   optional string version = 17;
+
+  // A sorted list of CowMergeOperation. When writing cow, we can choose to
+  // skip writing the raw bytes for these extents. During snapshot merge, the
+  // bytes will read from the source partitions instead.
+  repeated CowMergeOperation merge_operations = 18;
 }
 
 message DynamicPartitionGroup {