From 30bdee991027d1380459fd91e0a98ae21bc13c19 Mon Sep 17 00:00:00 2001
From: Yifan Hong <elsk@google.com>
Date: Mon, 23 Sep 2019 18:31:01 -0700
Subject: [PATCH] libsnapshot: CreateUpdateSnapshots reuse COW space

At the second update for Virtual A/B device,
PartitionCowCreator thinks the COW partitions in
the existing super metadata as occupied. In fact, these
partitions aren't used after the merge.

Now, unmap these partitions from the device mapper. If
something bad happens and the previous update has not been
merged yet, the unmap will fail. Then, delete these old
COW partitions from the device.

Test: Virtual A/B update twice
Bug: 135752105
Change-Id: Iab867ded19755089e6e0800e553a10fbcddbb931
---
 fs_mgr/libsnapshot/snapshot.cpp      | 20 +++++++++++
 fs_mgr/libsnapshot/snapshot_test.cpp | 53 ++++++++++++++++++++++++++++
 2 files changed, 73 insertions(+)

diff --git a/fs_mgr/libsnapshot/snapshot.cpp b/fs_mgr/libsnapshot/snapshot.cpp
index 9d0272b33..2e689bd9f 100644
--- a/fs_mgr/libsnapshot/snapshot.cpp
+++ b/fs_mgr/libsnapshot/snapshot.cpp
@@ -1741,6 +1741,22 @@ bool SnapshotManager::ForceLocalImageManager() {
     return true;
 }
 
+static void UnmapAndDeleteCowPartition(MetadataBuilder* current_metadata) {
+    auto& dm = DeviceMapper::Instance();
+    std::vector<std::string> to_delete;
+    for (auto* existing_cow_partition : current_metadata->ListPartitionsInGroup(kCowGroupName)) {
+        if (!dm.DeleteDeviceIfExists(existing_cow_partition->name())) {
+            LOG(WARNING) << existing_cow_partition->name()
+                         << " cannot be unmapped and its space cannot be reclaimed";
+            continue;
+        }
+        to_delete.push_back(existing_cow_partition->name());
+    }
+    for (const auto& name : to_delete) {
+        current_metadata->RemovePartition(name);
+    }
+}
+
 bool SnapshotManager::CreateUpdateSnapshots(const DeltaArchiveManifest& manifest) {
     auto lock = LockExclusive();
     if (!lock) return false;
@@ -1788,6 +1804,10 @@ bool SnapshotManager::CreateUpdateSnapshots(const DeltaArchiveManifest& manifest
         return false;
     }
 
+    // Delete previous COW partitions in current_metadata so that PartitionCowCreator marks those as
+    // free regions.
+    UnmapAndDeleteCowPartition(current_metadata.get());
+
     // Check that all these metadata is not retrofit dynamic partitions. Snapshots on
     // devices with retrofit dynamic partitions does not make sense.
     // This ensures that current_metadata->GetFreeRegions() uses the same device
diff --git a/fs_mgr/libsnapshot/snapshot_test.cpp b/fs_mgr/libsnapshot/snapshot_test.cpp
index c7b2b4cfb..f67e76e8e 100644
--- a/fs_mgr/libsnapshot/snapshot_test.cpp
+++ b/fs_mgr/libsnapshot/snapshot_test.cpp
@@ -47,8 +47,10 @@ using android::fiemap::IImageManager;
 using android::fs_mgr::BlockDeviceInfo;
 using android::fs_mgr::CreateLogicalPartitionParams;
 using android::fs_mgr::DestroyLogicalPartition;
+using android::fs_mgr::Extent;
 using android::fs_mgr::GetPartitionGroupName;
 using android::fs_mgr::GetPartitionName;
+using android::fs_mgr::Interval;
 using android::fs_mgr::MetadataBuilder;
 using chromeos_update_engine::DeltaArchiveManifest;
 using chromeos_update_engine::PartitionUpdate;
@@ -983,6 +985,57 @@ TEST_F(SnapshotUpdateTest, CancelAfterApply) {
     ASSERT_TRUE(sm->CancelUpdate());
 }
 
+static std::vector<Interval> ToIntervals(const std::vector<std::unique_ptr<Extent>>& extents) {
+    std::vector<Interval> ret;
+    std::transform(extents.begin(), extents.end(), std::back_inserter(ret),
+                   [](const auto& extent) { return extent->AsLinearExtent()->AsInterval(); });
+    return ret;
+}
+
+// Test that at the second update, old COW partition spaces are reclaimed.
+TEST_F(SnapshotUpdateTest, ReclaimCow) {
+    // Execute the first update.
+    ASSERT_TRUE(sm->BeginUpdate());
+    ASSERT_TRUE(sm->CreateUpdateSnapshots(manifest_));
+    ASSERT_TRUE(sm->FinishedSnapshotWrites());
+
+    // Simulate shutting down the device.
+    ASSERT_TRUE(UnmapAll());
+
+    // After reboot, init does first stage mount.
+    auto init = SnapshotManager::NewForFirstStageMount(new TestDeviceInfo(fake_super, "_b"));
+    ASSERT_NE(init, nullptr);
+    ASSERT_TRUE(init->NeedSnapshotsInFirstStageMount());
+    ASSERT_TRUE(init->CreateLogicalAndSnapshotPartitions("super"));
+    init = nullptr;
+
+    // Initiate the merge and wait for it to be completed.
+    auto new_sm = SnapshotManager::New(new TestDeviceInfo(fake_super, "_b"));
+    ASSERT_TRUE(new_sm->InitiateMerge());
+    ASSERT_EQ(UpdateState::MergeCompleted, new_sm->ProcessUpdateState());
+
+    // Execute the second update.
+    ASSERT_TRUE(new_sm->BeginUpdate());
+    ASSERT_TRUE(new_sm->CreateUpdateSnapshots(manifest_));
+
+    // Check that the old COW space is reclaimed and does not occupy space of mapped partitions.
+    auto src = MetadataBuilder::New(*opener_, "super", 1);
+    auto tgt = MetadataBuilder::New(*opener_, "super", 0);
+    for (const auto& cow_part_name : {"sys_a-cow", "vnd_a-cow", "prd_a-cow"}) {
+        auto* cow_part = tgt->FindPartition(cow_part_name);
+        ASSERT_NE(nullptr, cow_part) << cow_part_name << " does not exist in target metadata";
+        auto cow_intervals = ToIntervals(cow_part->extents());
+        for (const auto& old_part_name : {"sys_b", "vnd_b", "prd_b"}) {
+            auto* old_part = src->FindPartition(old_part_name);
+            ASSERT_NE(nullptr, old_part) << old_part_name << " does not exist in source metadata";
+            auto old_intervals = ToIntervals(old_part->extents());
+
+            auto intersect = Interval::Intersect(cow_intervals, old_intervals);
+            ASSERT_TRUE(intersect.empty()) << "COW uses space of source partitions";
+        }
+    }
+}
+
 }  // namespace snapshot
 }  // namespace android