From 1af515b57c97ef9ce88334ffab9275df71ad87d0 Mon Sep 17 00:00:00 2001
From: Yifan Hong <elsk@google.com>
Date: Fri, 13 Dec 2019 19:28:28 -0800
Subject: [PATCH] libsnapshot: add WaitForMerge

Add an API that does not initiate the merge, but only waits for it to
finish. It is different from ProcessUpdateState API in that it also
blocks when state == UNVERIFIED and booting from the new slot.
(ProcessUpdateState immediately returns in this case).

This is useful for android.os.UpdateEngine.CleanupSuccessfulUpdate().
Bug: 138808328
Test: libsnapshot_test

Change-Id: I7cc59fcaf69616e7ec7ebe6101991b5106845b65
---
 .../include/libsnapshot/snapshot.h            |  7 ++++
 fs_mgr/libsnapshot/snapshot.cpp               | 18 ++++++++-
 fs_mgr/libsnapshot/snapshot_test.cpp          | 39 +++++++++++++++++++
 3 files changed, 63 insertions(+), 1 deletion(-)

diff --git a/fs_mgr/libsnapshot/include/libsnapshot/snapshot.h b/fs_mgr/libsnapshot/include/libsnapshot/snapshot.h
index eb48a761a..445e6dbed 100644
--- a/fs_mgr/libsnapshot/include/libsnapshot/snapshot.h
+++ b/fs_mgr/libsnapshot/include/libsnapshot/snapshot.h
@@ -198,6 +198,13 @@ class SnapshotManager final {
     //   - other states indicating an error has occurred
     UpdateState InitiateMergeAndWait();
 
+    // Wait for the merge if rebooted into the new slot. Does NOT initiate a
+    // merge. If the merge has not been initiated (but should be), wait.
+    // Returns:
+    //   - true there is no merge or merge finishes
+    //   - false indicating an error has occurred
+    bool WaitForMerge();
+
     // Find the status of the current update, if any.
     //
     // |progress| depends on the returned status:
diff --git a/fs_mgr/libsnapshot/snapshot.cpp b/fs_mgr/libsnapshot/snapshot.cpp
index ce960f90d..a0ec0685b 100644
--- a/fs_mgr/libsnapshot/snapshot.cpp
+++ b/fs_mgr/libsnapshot/snapshot.cpp
@@ -74,6 +74,7 @@ using namespace std::chrono_literals;
 using namespace std::string_literals;
 
 static constexpr char kBootIndicatorPath[] = "/metadata/ota/snapshot-boot";
+static constexpr auto kUpdateStateCheckInterval = 2s;
 
 // Note: IImageManager is an incomplete type in the header, so the default
 // destructor doesn't work.
@@ -731,7 +732,7 @@ UpdateState SnapshotManager::ProcessUpdateState(const std::function<void()>& cal
 
         // This wait is not super time sensitive, so we have a relatively
         // low polling frequency.
-        std::this_thread::sleep_for(2s);
+        std::this_thread::sleep_for(kUpdateStateCheckInterval);
     }
 }
 
@@ -2241,6 +2242,21 @@ UpdateState SnapshotManager::InitiateMergeAndWait() {
     return state;
 }
 
+bool SnapshotManager::WaitForMerge() {
+    LOG(INFO) << "Waiting for any previous merge request to complete. "
+              << "This can take up to several minutes.";
+    while (true) {
+        auto state = ProcessUpdateState();
+        if (state == UpdateState::Unverified && GetCurrentSlot() == Slot::Target) {
+            LOG(INFO) << "Wait for merge to be initiated.";
+            std::this_thread::sleep_for(kUpdateStateCheckInterval);
+            continue;
+        }
+        LOG(INFO) << "Wait for merge exits with state " << state;
+        return state == UpdateState::None || state == UpdateState::MergeCompleted;
+    }
+}
+
 bool SnapshotManager::HandleImminentDataWipe(const std::function<void()>& callback) {
     if (!device_->IsRecovery()) {
         LOG(ERROR) << "Data wipes are only allowed in recovery.";
diff --git a/fs_mgr/libsnapshot/snapshot_test.cpp b/fs_mgr/libsnapshot/snapshot_test.cpp
index df4eb0b11..0f5af141a 100644
--- a/fs_mgr/libsnapshot/snapshot_test.cpp
+++ b/fs_mgr/libsnapshot/snapshot_test.cpp
@@ -1546,6 +1546,45 @@ TEST_F(SnapshotUpdateTest, Overflow) {
             << "FinishedSnapshotWrites should detect overflow of CoW device.";
 }
 
+TEST_F(SnapshotUpdateTest, WaitForMerge) {
+    AddOperationForPartitions();
+
+    // Execute the update.
+    ASSERT_TRUE(sm->BeginUpdate());
+    ASSERT_TRUE(sm->CreateUpdateSnapshots(manifest_));
+
+    // Write some data to target partitions.
+    for (const auto& name : {"sys_b", "vnd_b", "prd_b"}) {
+        ASSERT_TRUE(WriteSnapshotAndHash(name));
+    }
+
+    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(nullptr, init);
+        ASSERT_TRUE(init->CreateLogicalAndSnapshotPartitions("super"));
+    }
+
+    auto new_sm = SnapshotManager::New(new TestDeviceInfo(fake_super, "_b"));
+    ASSERT_NE(nullptr, new_sm);
+
+    auto waiter = std::async(std::launch::async, [&new_sm] { return new_sm->WaitForMerge(); });
+    ASSERT_EQ(std::future_status::timeout, waiter.wait_for(1s))
+            << "WaitForMerge should block when not initiated";
+
+    auto merger =
+            std::async(std::launch::async, [&new_sm] { return new_sm->InitiateMergeAndWait(); });
+    // Small images, so should be merged pretty quickly.
+    ASSERT_EQ(std::future_status::ready, waiter.wait_for(3s)) << "WaitForMerge did not finish";
+    ASSERT_TRUE(waiter.get());
+    ASSERT_THAT(merger.get(), AnyOf(UpdateState::None, UpdateState::MergeCompleted));
+}
+
 class FlashAfterUpdateTest : public SnapshotUpdateTest,
                              public WithParamInterface<std::tuple<uint32_t, bool>> {
   public: