ShouldSkipOperation -> OptimizeOperation

For SOURCE_COPY operations like
  563412 -> 123456
OptimizeOperation optimizes it to
  5612 -> 1256
and skip writing blocks that does not need to be written for snapshot
partitions.

Bug: 148623880
Test: update_engine_unittests
Test: apply incremental OTA

Change-Id: Ifd2c3851f703f272a74c8f0e9a1c9a82dbcce3e3
diff --git a/common/dynamic_partition_control_interface.h b/common/dynamic_partition_control_interface.h
index 952e8d4..a878f2e 100644
--- a/common/dynamic_partition_control_interface.h
+++ b/common/dynamic_partition_control_interface.h
@@ -56,12 +56,17 @@
   // Return the feature flags of Virtual A/B on this device.
   virtual FeatureFlag GetVirtualAbFeatureFlag() = 0;
 
-  // Checks if |operation| can be skipped on the given partition.
+  // Attempt to optimize |operation|.
+  // If successful, |optimized| contains an operation with extents that
+  // needs to be written.
+  // If failed, no optimization is available, and caller should perform
+  // |operation| directly.
   // |partition_name| should not have the slot suffix; implementation of
   // DynamicPartitionControlInterface checks partition at the target slot
   // previously set with PreparePartitionsForUpdate().
-  virtual bool ShouldSkipOperation(const std::string& partition_name,
-                                   const InstallOperation& operation) = 0;
+  virtual bool OptimizeOperation(const std::string& partition_name,
+                                 const InstallOperation& operation,
+                                 InstallOperation* optimized) = 0;
 
   // Do necessary cleanups before destroying the object.
   virtual void Cleanup() = 0;
diff --git a/common/dynamic_partition_control_stub.cc b/common/dynamic_partition_control_stub.cc
index 3d31e20..dee4555 100644
--- a/common/dynamic_partition_control_stub.cc
+++ b/common/dynamic_partition_control_stub.cc
@@ -33,8 +33,10 @@
   return FeatureFlag(FeatureFlag::Value::NONE);
 }
 
-bool DynamicPartitionControlStub::ShouldSkipOperation(
-    const std::string& partition_name, const InstallOperation& operation) {
+bool DynamicPartitionControlStub::OptimizeOperation(
+    const std::string& partition_name,
+    const InstallOperation& operation,
+    InstallOperation* optimized) {
   return false;
 }
 
diff --git a/common/dynamic_partition_control_stub.h b/common/dynamic_partition_control_stub.h
index 9be988b..ddef36d 100644
--- a/common/dynamic_partition_control_stub.h
+++ b/common/dynamic_partition_control_stub.h
@@ -30,8 +30,9 @@
  public:
   FeatureFlag GetDynamicPartitionsFeatureFlag() override;
   FeatureFlag GetVirtualAbFeatureFlag() override;
-  bool ShouldSkipOperation(const std::string& partition_name,
-                           const InstallOperation& operation) override;
+  bool OptimizeOperation(const std::string& partition_name,
+                         const InstallOperation& operation,
+                         InstallOperation* optimized) override;
   void Cleanup() override;
   bool PreparePartitionsForUpdate(uint32_t source_slot,
                                   uint32_t target_slot,
diff --git a/dynamic_partition_control_android.cc b/dynamic_partition_control_android.cc
index 1993661..d75a8fc 100644
--- a/dynamic_partition_control_android.cc
+++ b/dynamic_partition_control_android.cc
@@ -51,9 +51,9 @@
 using android::fs_mgr::Partition;
 using android::fs_mgr::PartitionOpener;
 using android::fs_mgr::SlotSuffixForSlotNumber;
+using android::snapshot::OptimizeSourceCopyOperation;
 using android::snapshot::Return;
 using android::snapshot::SnapshotManager;
-using android::snapshot::SourceCopyOperationIsClone;
 using android::snapshot::UpdateState;
 
 namespace chromeos_update_engine {
@@ -115,15 +115,17 @@
   return virtual_ab_;
 }
 
-bool DynamicPartitionControlAndroid::ShouldSkipOperation(
-    const std::string& partition_name, const InstallOperation& operation) {
+bool DynamicPartitionControlAndroid::OptimizeOperation(
+    const std::string& partition_name,
+    const InstallOperation& operation,
+    InstallOperation* optimized) {
   switch (operation.type()) {
     case InstallOperation::SOURCE_COPY:
       return target_supports_snapshot_ &&
              GetVirtualAbFeatureFlag().IsEnabled() &&
              mapped_devices_.count(partition_name +
                                    SlotSuffixForSlotNumber(target_slot_)) > 0 &&
-             SourceCopyOperationIsClone(operation);
+             OptimizeSourceCopyOperation(operation, optimized);
       break;
     default:
       break;
diff --git a/dynamic_partition_control_android.h b/dynamic_partition_control_android.h
index c842efc..1ff59a9 100644
--- a/dynamic_partition_control_android.h
+++ b/dynamic_partition_control_android.h
@@ -35,8 +35,9 @@
   ~DynamicPartitionControlAndroid();
   FeatureFlag GetDynamicPartitionsFeatureFlag() override;
   FeatureFlag GetVirtualAbFeatureFlag() override;
-  bool ShouldSkipOperation(const std::string& partition_name,
-                           const InstallOperation& operation) override;
+  bool OptimizeOperation(const std::string& partition_name,
+                         const InstallOperation& operation,
+                         InstallOperation* optimized) override;
   void Cleanup() override;
 
   bool PreparePartitionsForUpdate(uint32_t source_slot,
diff --git a/dynamic_partition_control_android_unittest.cc b/dynamic_partition_control_android_unittest.cc
index 3e8375c..4b1870d 100644
--- a/dynamic_partition_control_android_unittest.cc
+++ b/dynamic_partition_control_android_unittest.cc
@@ -625,7 +625,7 @@
       << "Should not be able to apply to current slot.";
 }
 
-TEST_P(DynamicPartitionControlAndroidTestP, ShouldSkipOperationTest) {
+TEST_P(DynamicPartitionControlAndroidTestP, OptimizeOperationTest) {
   ASSERT_TRUE(dynamicControl().PreparePartitionsForUpdate(
       source(),
       target(),
@@ -635,16 +635,17 @@
   dynamicControl().set_fake_mapped_devices({T("foo")});
 
   InstallOperation iop;
+  InstallOperation optimized;
   Extent *se, *de;
 
   // Not a SOURCE_COPY operation, cannot skip.
   iop.set_type(InstallOperation::REPLACE);
-  EXPECT_FALSE(dynamicControl().ShouldSkipOperation("foo", iop));
+  EXPECT_FALSE(dynamicControl().OptimizeOperation("foo", iop, &optimized));
 
   iop.set_type(InstallOperation::SOURCE_COPY);
 
   // By default GetVirtualAbFeatureFlag is disabled. Cannot skip operation.
-  EXPECT_FALSE(dynamicControl().ShouldSkipOperation("foo", iop));
+  EXPECT_FALSE(dynamicControl().OptimizeOperation("foo", iop, &optimized));
 
   // Enable GetVirtualAbFeatureFlag in the mock interface.
   ON_CALL(dynamicControl(), GetVirtualAbFeatureFlag())
@@ -652,19 +653,21 @@
 
   // By default target_supports_snapshot_ is set to false. Cannot skip
   // operation.
-  EXPECT_FALSE(dynamicControl().ShouldSkipOperation("foo", iop));
+  EXPECT_FALSE(dynamicControl().OptimizeOperation("foo", iop, &optimized));
 
   SetSnapshotEnabled(true);
 
   // Empty source and destination. Skip.
-  EXPECT_TRUE(dynamicControl().ShouldSkipOperation("foo", iop));
+  EXPECT_TRUE(dynamicControl().OptimizeOperation("foo", iop, &optimized));
+  EXPECT_TRUE(optimized.src_extents().empty());
+  EXPECT_TRUE(optimized.dst_extents().empty());
 
   se = iop.add_src_extents();
   se->set_start_block(0);
   se->set_num_blocks(1);
 
   // There is something in sources, but destinations are empty. Cannot skip.
-  EXPECT_FALSE(dynamicControl().ShouldSkipOperation("foo", iop));
+  EXPECT_FALSE(dynamicControl().OptimizeOperation("foo", iop, &optimized));
 
   InstallOperation iop2;
 
@@ -673,45 +676,79 @@
   de->set_num_blocks(1);
 
   // There is something in destinations, but sources are empty. Cannot skip.
-  EXPECT_FALSE(dynamicControl().ShouldSkipOperation("foo", iop2));
+  EXPECT_FALSE(dynamicControl().OptimizeOperation("foo", iop2, &optimized));
 
   de = iop.add_dst_extents();
   de->set_start_block(0);
   de->set_num_blocks(1);
 
   // Sources and destinations are identical. Skip.
-  EXPECT_TRUE(dynamicControl().ShouldSkipOperation("foo", iop));
+  EXPECT_TRUE(dynamicControl().OptimizeOperation("foo", iop, &optimized));
+  EXPECT_TRUE(optimized.src_extents().empty());
+  EXPECT_TRUE(optimized.dst_extents().empty());
 
   se = iop.add_src_extents();
   se->set_start_block(1);
   se->set_num_blocks(5);
 
   // There is something in source, but not in destination. Cannot skip.
-  EXPECT_FALSE(dynamicControl().ShouldSkipOperation("foo", iop));
+  EXPECT_FALSE(dynamicControl().OptimizeOperation("foo", iop, &optimized));
 
   de = iop.add_dst_extents();
   de->set_start_block(1);
   de->set_num_blocks(5);
 
   // There is source and destination are equal. Skip.
-  EXPECT_TRUE(dynamicControl().ShouldSkipOperation("foo", iop));
+  EXPECT_TRUE(dynamicControl().OptimizeOperation("foo", iop, &optimized));
+  EXPECT_TRUE(optimized.src_extents().empty());
+  EXPECT_TRUE(optimized.dst_extents().empty());
 
   de = iop.add_dst_extents();
   de->set_start_block(6);
   de->set_num_blocks(5);
 
   // There is something extra in dest. Cannot skip.
-  EXPECT_FALSE(dynamicControl().ShouldSkipOperation("foo", iop));
+  EXPECT_FALSE(dynamicControl().OptimizeOperation("foo", iop, &optimized));
 
   se = iop.add_src_extents();
   se->set_start_block(6);
   se->set_num_blocks(5);
 
   // Source and dest are identical again. Skip.
-  EXPECT_TRUE(dynamicControl().ShouldSkipOperation("foo", iop));
+  EXPECT_TRUE(dynamicControl().OptimizeOperation("foo", iop, &optimized));
+  EXPECT_TRUE(optimized.src_extents().empty());
+  EXPECT_TRUE(optimized.dst_extents().empty());
+
+  iop.Clear();
+  iop.set_type(InstallOperation::SOURCE_COPY);
+  se = iop.add_src_extents();
+  se->set_start_block(1);
+  se->set_num_blocks(1);
+  se = iop.add_src_extents();
+  se->set_start_block(3);
+  se->set_num_blocks(2);
+  se = iop.add_src_extents();
+  se->set_start_block(7);
+  se->set_num_blocks(2);
+  de = iop.add_dst_extents();
+  de->set_start_block(2);
+  de->set_num_blocks(5);
+
+  // [1, 3, 4, 7, 8] -> [2, 3, 4, 5, 6] should return [1, 7, 8] -> [2, 5, 6]
+  EXPECT_TRUE(dynamicControl().OptimizeOperation("foo", iop, &optimized));
+  ASSERT_EQ(2, optimized.src_extents_size());
+  ASSERT_EQ(2, optimized.dst_extents_size());
+  EXPECT_EQ(1u, optimized.src_extents(0).start_block());
+  EXPECT_EQ(1u, optimized.src_extents(0).num_blocks());
+  EXPECT_EQ(2u, optimized.dst_extents(0).start_block());
+  EXPECT_EQ(1u, optimized.dst_extents(0).num_blocks());
+  EXPECT_EQ(7u, optimized.src_extents(1).start_block());
+  EXPECT_EQ(2u, optimized.src_extents(1).num_blocks());
+  EXPECT_EQ(5u, optimized.dst_extents(1).start_block());
+  EXPECT_EQ(2u, optimized.dst_extents(1).num_blocks());
 
   // Don't skip for static partitions.
-  EXPECT_FALSE(dynamicControl().ShouldSkipOperation("bar", iop));
+  EXPECT_FALSE(dynamicControl().OptimizeOperation("bar", iop, &optimized));
 }
 
 }  // namespace chromeos_update_engine
diff --git a/payload_consumer/delta_performer.cc b/payload_consumer/delta_performer.cc
index d8e58b4..4c4ff04 100644
--- a/payload_consumer/delta_performer.cc
+++ b/payload_consumer/delta_performer.cc
@@ -1200,8 +1200,11 @@
   // decide it the operation should be skipped.
   const PartitionUpdate& partition = partitions_[current_partition_];
   const auto& partition_control = boot_control_->GetDynamicPartitionControl();
-  bool should_skip = partition_control->ShouldSkipOperation(
-      partition.partition_name(), operation);
+
+  InstallOperation buf;
+  bool should_optimize = partition_control->OptimizeOperation(
+      partition.partition_name(), operation, &buf);
+  const InstallOperation& optimized = should_optimize ? buf : operation;
 
   if (operation.has_src_sha256_hash()) {
     bool read_ok;
@@ -1213,9 +1216,18 @@
     // device doesn't match or there was an error reading the source partition.
     // Note that this code will also fall back if writing the target partition
     // fails.
-    if (should_skip) {
-      read_ok = fd_utils::ReadAndHashExtents(
-          source_fd_, operation.src_extents(), block_size_, &source_hash);
+    if (should_optimize) {
+      // Hash operation.src_extents(), then copy optimized.src_extents to
+      // optimized.dst_extents.
+      read_ok =
+          fd_utils::ReadAndHashExtents(
+              source_fd_, operation.src_extents(), block_size_, &source_hash) &&
+          fd_utils::CopyAndHashExtents(source_fd_,
+                                       optimized.src_extents(),
+                                       target_fd_,
+                                       optimized.dst_extents(),
+                                       block_size_,
+                                       nullptr /* skip hashing */);
     } else {
       read_ok = fd_utils::CopyAndHashExtents(source_fd_,
                                              operation.src_extents(),
@@ -1240,9 +1252,16 @@
                  << base::HexEncode(expected_source_hash.data(),
                                     expected_source_hash.size());
 
-    if (should_skip) {
+    if (should_optimize) {
       TEST_AND_RETURN_FALSE(fd_utils::ReadAndHashExtents(
           source_ecc_fd_, operation.src_extents(), block_size_, &source_hash));
+      TEST_AND_RETURN_FALSE(
+          fd_utils::CopyAndHashExtents(source_ecc_fd_,
+                                       optimized.src_extents(),
+                                       target_fd_,
+                                       optimized.dst_extents(),
+                                       block_size_,
+                                       nullptr /* skip hashing */));
     } else {
       TEST_AND_RETURN_FALSE(
           fd_utils::CopyAndHashExtents(source_ecc_fd_,
@@ -1264,22 +1283,19 @@
     // at this point, but we fall back to the raw device since the error
     // corrected device can be shorter or not available.
 
-    if (should_skip)
-      return true;
-
     if (OpenCurrentECCPartition() &&
         fd_utils::CopyAndHashExtents(source_ecc_fd_,
-                                     operation.src_extents(),
+                                     optimized.src_extents(),
                                      target_fd_,
-                                     operation.dst_extents(),
+                                     optimized.dst_extents(),
                                      block_size_,
                                      nullptr)) {
       return true;
     }
     TEST_AND_RETURN_FALSE(fd_utils::CopyAndHashExtents(source_fd_,
-                                                       operation.src_extents(),
+                                                       optimized.src_extents(),
                                                        target_fd_,
-                                                       operation.dst_extents(),
+                                                       optimized.dst_extents(),
                                                        block_size_,
                                                        nullptr));
   }