Merge changes Iab867ded,Id37e34b4
am: ab09a1764b

Change-Id: I159c26546a5929998fb424e31673d7e00b3bab62
diff --git a/fs_mgr/libsnapshot/snapshot.cpp b/fs_mgr/libsnapshot/snapshot.cpp
index 9d0272b..2e689bd 100644
--- a/fs_mgr/libsnapshot/snapshot.cpp
+++ b/fs_mgr/libsnapshot/snapshot.cpp
@@ -1741,6 +1741,22 @@
     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 @@
         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 bd040b4..c94fde5 100644
--- a/fs_mgr/libsnapshot/snapshot_test.cpp
+++ b/fs_mgr/libsnapshot/snapshot_test.cpp
@@ -46,8 +46,10 @@
 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;
@@ -505,10 +507,7 @@
     ASSERT_TRUE(sm->FinishedSnapshotWrites());
     ASSERT_TRUE(DestroyLogicalPartition("test_partition_b-base"));
 
-    auto rebooted = new TestDeviceInfo(fake_super);
-    rebooted->set_slot_suffix("_b");
-
-    auto init = SnapshotManager::NewForFirstStageMount(rebooted);
+    auto init = SnapshotManager::NewForFirstStageMount(new TestDeviceInfo(fake_super, "_b"));
     ASSERT_NE(init, nullptr);
     ASSERT_TRUE(init->NeedSnapshotsInFirstStageMount());
     ASSERT_TRUE(init->CreateLogicalAndSnapshotPartitions("super"));
@@ -548,10 +547,7 @@
     FormatFakeSuper();
     ASSERT_TRUE(CreatePartition("test_partition_b", kDeviceSize));
 
-    auto rebooted = new TestDeviceInfo(fake_super);
-    rebooted->set_slot_suffix("_b");
-
-    auto init = SnapshotManager::NewForFirstStageMount(rebooted);
+    auto init = SnapshotManager::NewForFirstStageMount(new TestDeviceInfo(fake_super, "_b"));
     ASSERT_NE(init, nullptr);
     ASSERT_TRUE(init->NeedSnapshotsInFirstStageMount());
     ASSERT_TRUE(init->CreateLogicalAndSnapshotPartitions("super"));
@@ -589,10 +585,7 @@
     ASSERT_TRUE(sm->FinishedSnapshotWrites());
     ASSERT_TRUE(DestroyLogicalPartition("test_partition_b-base"));
 
-    auto rebooted = new TestDeviceInfo(fake_super);
-    rebooted->set_slot_suffix("_b");
-
-    auto init = SnapshotManager::NewForFirstStageMount(rebooted);
+    auto init = SnapshotManager::NewForFirstStageMount(new TestDeviceInfo(fake_super, "_b"));
     ASSERT_NE(init, nullptr);
     ASSERT_TRUE(init->NeedSnapshotsInFirstStageMount());
     ASSERT_TRUE(init->CreateLogicalAndSnapshotPartitions("super"));
@@ -744,9 +737,6 @@
 // Also test UnmapUpdateSnapshot unmaps everything.
 // Also test first stage mount and merge after this.
 TEST_F(SnapshotUpdateTest, FullUpdateFlow) {
-    // OTA client calls BeginUpdate before doing anything.
-    ASSERT_TRUE(sm->BeginUpdate());
-
     // OTA client blindly unmaps all partitions that are possibly mapped.
     for (const auto& name : {"sys_b", "vnd_b", "prd_b"}) {
         ASSERT_TRUE(sm->UnmapUpdateSnapshot(name));
@@ -757,6 +747,8 @@
     SetSize(vnd_, 4_MiB);
     SetSize(prd_, 4_MiB);
 
+    // Execute the update.
+    ASSERT_TRUE(sm->BeginUpdate());
     ASSERT_TRUE(sm->CreateUpdateSnapshots(manifest_));
 
     // Test that partitions prioritize using space in super.
@@ -795,9 +787,7 @@
     ASSERT_TRUE(UnmapAll());
 
     // After reboot, init does first stage mount.
-    auto rebooted = new TestDeviceInfo(fake_super);
-    rebooted->set_slot_suffix("_b");
-    auto init = SnapshotManager::NewForFirstStageMount(rebooted);
+    auto init = SnapshotManager::NewForFirstStageMount(new TestDeviceInfo(fake_super, "_b"));
     ASSERT_NE(init, nullptr);
     ASSERT_TRUE(init->NeedSnapshotsInFirstStageMount());
     ASSERT_TRUE(init->CreateLogicalAndSnapshotPartitions("super"));
@@ -948,9 +938,7 @@
     ASSERT_TRUE(UnmapAll());
 
     // After reboot, init does first stage mount.
-    auto rebooted = new TestDeviceInfo(fake_super);
-    rebooted->set_slot_suffix("_b");
-    auto init = SnapshotManager::NewForFirstStageMount(rebooted);
+    auto init = SnapshotManager::NewForFirstStageMount(new TestDeviceInfo(fake_super, "_b"));
     ASSERT_NE(init, nullptr);
     ASSERT_TRUE(init->NeedSnapshotsInFirstStageMount());
     ASSERT_TRUE(init->CreateLogicalAndSnapshotPartitions("super"));
@@ -962,9 +950,7 @@
 
     // Simulate shutting down the device again.
     ASSERT_TRUE(UnmapAll());
-    rebooted = new TestDeviceInfo(fake_super);
-    rebooted->set_slot_suffix("_a");
-    init = SnapshotManager::NewForFirstStageMount(rebooted);
+    init = SnapshotManager::NewForFirstStageMount(new TestDeviceInfo(fake_super, "_a"));
     ASSERT_NE(init, nullptr);
     ASSERT_FALSE(init->NeedSnapshotsInFirstStageMount());
     ASSERT_TRUE(init->CreateLogicalAndSnapshotPartitions("super"));
@@ -982,6 +968,57 @@
     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
 
diff --git a/fs_mgr/libsnapshot/test_helpers.h b/fs_mgr/libsnapshot/test_helpers.h
index 555bcf4..bb19adc 100644
--- a/fs_mgr/libsnapshot/test_helpers.h
+++ b/fs_mgr/libsnapshot/test_helpers.h
@@ -59,6 +59,10 @@
   public:
     TestDeviceInfo() {}
     explicit TestDeviceInfo(const std::string& fake_super) { set_fake_super(fake_super); }
+    TestDeviceInfo(const std::string& fake_super, const std::string& slot_suffix)
+        : TestDeviceInfo(fake_super) {
+        set_slot_suffix(slot_suffix);
+    }
     std::string GetGsidDir() const override { return "ota/test"s; }
     std::string GetMetadataDir() const override { return "/metadata/ota/test"s; }
     std::string GetSlotSuffix() const override { return slot_suffix_; }