Add CowMergeOperations as a hint for snapshot write am: e9156ec8de

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

Change-Id: I61f25f9d339da6171cc0ea4c65844d143c39972e
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 {