Add OptimizeSourceCopyOperation

... so that an operation can be skipped partially. For example, if
an operation contains blocks:
    563412 -> 123456
... then optimized operation is:
    5612 -> 1256

Test: update_engine_unittests
Test: apply incremental OTA
Bug: 148623880

In an experiment, this reduces CoW size of an incremental update
package by 200MB (out of 700MB).

Change-Id: I86ca23fd589ddbc84c81318283b5f4e71782a759

fs_mgr/libsnapshot/include/libsnapshot/snapshot.h

@@ -74,7 +74,8 @@ class SnapshotStatus;
 
 static constexpr const std::string_view kCowGroupName = "cow";
 
-bool SourceCopyOperationIsClone(const chromeos_update_engine::InstallOperation& operation);
+bool OptimizeSourceCopyOperation(const chromeos_update_engine::InstallOperation& operation,
+                                 chromeos_update_engine::InstallOperation* optimized);
 
 enum class CreateResult : unsigned int {
     ERROR,

fs_mgr/libsnapshot/partition_cow_creator.cpp

@@ -62,17 +62,68 @@ bool PartitionCowCreator::HasExtent(Partition* p, Extent* e) {
     return false;
 }
 
-bool SourceCopyOperationIsClone(const InstallOperation& operation) {
-    using ChromeOSExtent = chromeos_update_engine::Extent;
-    if (operation.src_extents().size() != operation.dst_extents().size()) {
+bool OptimizeSourceCopyOperation(const InstallOperation& operation, InstallOperation* optimized) {
+    if (operation.type() != InstallOperation::SOURCE_COPY) {
         return false;
     }
-    return std::equal(operation.src_extents().begin(), operation.src_extents().end(),
-                      operation.dst_extents().begin(),
-                      [](const ChromeOSExtent& src, const ChromeOSExtent& dst) {
-                          return src.start_block() == dst.start_block() &&
-                                 src.num_blocks() == dst.num_blocks();
-                      });
+
+    optimized->Clear();
+    optimized->set_type(InstallOperation::SOURCE_COPY);
+
+    const auto& src_extents = operation.src_extents();
+    const auto& dst_extents = operation.dst_extents();
+
+    // If input is empty, skip by returning an empty result.
+    if (src_extents.empty() && dst_extents.empty()) {
+        return true;
+    }
+
+    auto s_it = src_extents.begin();
+    auto d_it = dst_extents.begin();
+    uint64_t s_offset = 0;  // offset within *s_it
+    uint64_t d_offset = 0;  // offset within *d_it
+    bool is_optimized = false;
+
+    while (s_it != src_extents.end() || d_it != dst_extents.end()) {
+        if (s_it == src_extents.end() || d_it == dst_extents.end()) {
+            LOG(ERROR) << "number of blocks do not equal in src_extents and dst_extents";
+            return false;
+        }
+        if (s_it->num_blocks() <= s_offset || d_it->num_blocks() <= d_offset) {
+            LOG(ERROR) << "Offset goes out of bounds.";
+            return false;
+        }
+
+        // Check the next |step| blocks, where |step| is the min of remaining blocks in the current
+        // source extent and current destination extent.
+        auto s_step = s_it->num_blocks() - s_offset;
+        auto d_step = d_it->num_blocks() - d_offset;
+        auto step = std::min(s_step, d_step);
+
+        bool moved = s_it->start_block() + s_offset != d_it->start_block() + d_offset;
+        if (moved) {
+            // If the next |step| blocks are not copied to the same location, add them to result.
+            AppendExtent(optimized->mutable_src_extents(), s_it->start_block() + s_offset, step);
+            AppendExtent(optimized->mutable_dst_extents(), d_it->start_block() + d_offset, step);
+        } else {
+            // The next |step| blocks are optimized out.
+            is_optimized = true;
+        }
+
+        // Advance offsets by |step|, and go to the next non-empty extent if the current extent is
+        // depleted.
+        s_offset += step;
+        d_offset += step;
+        while (s_it != src_extents.end() && s_offset >= s_it->num_blocks()) {
+            ++s_it;
+            s_offset = 0;
+        }
+        while (d_it != dst_extents.end() && d_offset >= d_it->num_blocks()) {
+            ++d_it;
+            d_offset = 0;
+        }
+    }
+    return is_optimized;
 }
 
 void WriteExtent(DmSnapCowSizeCalculator* sc, const chromeos_update_engine::Extent& de,
@@ -101,12 +152,15 @@ uint64_t PartitionCowCreator::GetCowSize() {
     if (operations == nullptr) return sc.cow_size_bytes();
 
     for (const auto& iop : *operations) {
-        // Do not allocate space for operations that are going to be skipped
+        const InstallOperation* written_op = &iop;
+        InstallOperation buf;
+        // Do not allocate space for extents that are going to be skipped
         // during OTA application.
-        if (iop.type() == InstallOperation::SOURCE_COPY && SourceCopyOperationIsClone(iop))
-            continue;
+        if (iop.type() == InstallOperation::SOURCE_COPY && OptimizeSourceCopyOperation(iop, &buf)) {
+            written_op = &buf;
+        }
 
-        for (const auto& de : iop.dst_extents()) {
+        for (const auto& de : written_op->dst_extents()) {
             WriteExtent(&sc, de, sectors_per_block);
         }
     }

fs_mgr/libsnapshot/partition_cow_creator_test.cpp

@@ -12,6 +12,9 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+#include <optional>
+#include <tuple>
+
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
 #include <libdm/dm.h>
@@ -26,6 +29,13 @@
 
 using namespace android::fs_mgr;
 
+using chromeos_update_engine::InstallOperation;
+using UeExtent = chromeos_update_engine::Extent;
+using google::protobuf::RepeatedPtrField;
+using ::testing::Matches;
+using ::testing::Pointwise;
+using ::testing::Truly;
+
 namespace android {
 namespace snapshot {
 
@@ -213,5 +223,76 @@ TEST(DmSnapshotInternals, CowSizeCalculator) {
     }
 }
 
+void BlocksToExtents(const std::vector<uint64_t>& blocks,
+                     google::protobuf::RepeatedPtrField<UeExtent>* extents) {
+    for (uint64_t block : blocks) {
+        AppendExtent(extents, block, 1);
+    }
+}
+
+template <typename T>
+std::vector<uint64_t> ExtentsToBlocks(const T& extents) {
+    std::vector<uint64_t> blocks;
+    for (const auto& extent : extents) {
+        for (uint64_t offset = 0; offset < extent.num_blocks(); ++offset) {
+            blocks.push_back(extent.start_block() + offset);
+        }
+    }
+    return blocks;
+}
+
+InstallOperation CreateCopyOp(const std::vector<uint64_t>& src_blocks,
+                              const std::vector<uint64_t>& dst_blocks) {
+    InstallOperation op;
+    op.set_type(InstallOperation::SOURCE_COPY);
+    BlocksToExtents(src_blocks, op.mutable_src_extents());
+    BlocksToExtents(dst_blocks, op.mutable_dst_extents());
+    return op;
+}
+
+// ExtentEqual(tuple<UeExtent, UeExtent>)
+MATCHER(ExtentEqual, "") {
+    auto&& [a, b] = arg;
+    return a.start_block() == b.start_block() && a.num_blocks() == b.num_blocks();
+}
+
+struct OptimizeOperationTestParam {
+    InstallOperation input;
+    std::optional<InstallOperation> expected_output;
+};
+
+class OptimizeOperationTest : public ::testing::TestWithParam<OptimizeOperationTestParam> {};
+TEST_P(OptimizeOperationTest, Test) {
+    InstallOperation actual_output;
+    EXPECT_EQ(GetParam().expected_output.has_value(),
+              OptimizeSourceCopyOperation(GetParam().input, &actual_output))
+            << "OptimizeSourceCopyOperation should "
+            << (GetParam().expected_output.has_value() ? "succeed" : "fail");
+    if (!GetParam().expected_output.has_value()) return;
+    EXPECT_THAT(actual_output.src_extents(),
+                Pointwise(ExtentEqual(), GetParam().expected_output->src_extents()));
+    EXPECT_THAT(actual_output.dst_extents(),
+                Pointwise(ExtentEqual(), GetParam().expected_output->dst_extents()));
+}
+
+std::vector<OptimizeOperationTestParam> GetOptimizeOperationTestParams() {
+    return {
+            {CreateCopyOp({}, {}), CreateCopyOp({}, {})},
+            {CreateCopyOp({1, 2, 4}, {1, 2, 4}), CreateCopyOp({}, {})},
+            {CreateCopyOp({1, 2, 3}, {4, 5, 6}), std::nullopt},
+            {CreateCopyOp({3, 2}, {1, 2}), CreateCopyOp({3}, {1})},
+            {CreateCopyOp({5, 6, 3, 4, 1, 2}, {1, 2, 3, 4, 5, 6}),
+             CreateCopyOp({5, 6, 1, 2}, {1, 2, 5, 6})},
+            {CreateCopyOp({1, 2, 3, 5, 5, 6}, {5, 6, 3, 4, 1, 2}),
+             CreateCopyOp({1, 2, 5, 5, 6}, {5, 6, 4, 1, 2})},
+            {CreateCopyOp({1, 2, 5, 6, 9, 10}, {1, 4, 5, 6, 7, 8}),
+             CreateCopyOp({2, 9, 10}, {4, 7, 8})},
+            {CreateCopyOp({2, 3, 3, 4, 4}, {1, 2, 3, 4, 5}), CreateCopyOp({2, 3, 4}, {1, 2, 5})},
+    };
+}
+
+INSTANTIATE_TEST_CASE_P(Snapshot, OptimizeOperationTest,
+                        ::testing::ValuesIn(GetOptimizeOperationTestParams()));
+
 }  // namespace snapshot
 }  // namespace android

fs_mgr/libsnapshot/utility.cpp

@@ -34,6 +34,7 @@ using android::fs_mgr::GetEntryForPath;
 using android::fs_mgr::MetadataBuilder;
 using android::fs_mgr::Partition;
 using android::fs_mgr::ReadDefaultFstab;
+using google::protobuf::RepeatedPtrField;
 
 namespace android {
 namespace snapshot {
@@ -166,5 +167,20 @@ std::ostream& operator<<(std::ostream& os, const Now&) {
     return os << std::put_time(&now, "%Y%m%d-%H%M%S");
 }
 
+void AppendExtent(RepeatedPtrField<chromeos_update_engine::Extent>* extents, uint64_t start_block,
+                  uint64_t num_blocks) {
+    if (extents->size() > 0) {
+        auto last_extent = extents->rbegin();
+        auto next_block = last_extent->start_block() + last_extent->num_blocks();
+        if (start_block == next_block) {
+            last_extent->set_num_blocks(last_extent->num_blocks() + num_blocks);
+            return;
+        }
+    }
+    auto* new_extent = extents->Add();
+    new_extent->set_start_block(start_block);
+    new_extent->set_num_blocks(num_blocks);
+}
+
 }  // namespace snapshot
 }  // namespace android

fs_mgr/libsnapshot/utility.h

@@ -125,5 +125,9 @@ bool WriteStringToFileAtomic(const std::string& content, const std::string& path
 struct Now {};
 std::ostream& operator<<(std::ostream& os, const Now&);
 
+// Append to |extents|. Merged into the last element if possible.
+void AppendExtent(google::protobuf::RepeatedPtrField<chromeos_update_engine::Extent>* extents,
+                  uint64_t start_block, uint64_t num_blocks);
+
 }  // namespace snapshot
 }  // namespace android