From 4913d7af12cbcc58a5c992aeef76a9a545959c2c Mon Sep 17 00:00:00 2001
From: Yifan Hong <elsk@google.com>
Date: Tue, 29 Oct 2019 16:23:10 -0700
Subject: [PATCH] [REFACTOR] libsnapshot: Add InitiateMergeAndWait

Move snapshot_ctl's merge command to SnapshotManager::
InitiateMergeAndWait function so that tests can use
it too.

Test: libsnapshot_test
Bug: 143551390

Change-Id: I0936c262afaca7ba445ee18465dca4e16b9416ad
---
 .../libsnapshot/include/libsnapshot/snapshot.h | 10 ++++++++++
 fs_mgr/libsnapshot/snapshot.cpp                | 18 ++++++++++++++++++
 fs_mgr/libsnapshot/snapshot_test.cpp           |  6 ++----
 fs_mgr/libsnapshot/snapshotctl.cpp             | 15 +--------------
 4 files changed, 31 insertions(+), 18 deletions(-)

diff --git a/fs_mgr/libsnapshot/include/libsnapshot/snapshot.h b/fs_mgr/libsnapshot/include/libsnapshot/snapshot.h
index 431fea174..260a10cd2 100644
--- a/fs_mgr/libsnapshot/include/libsnapshot/snapshot.h
+++ b/fs_mgr/libsnapshot/include/libsnapshot/snapshot.h
@@ -154,6 +154,7 @@ class SnapshotManager final {
     // rebooting or after rolling back), or merge the OTA.
     bool FinishedSnapshotWrites();
 
+  private:
     // Initiate a merge on all snapshot devices. This should only be used after an
     // update has been marked successful after booting.
     bool InitiateMerge();
@@ -181,6 +182,15 @@ class SnapshotManager final {
     //   GetUpdateState will return None, and a new update can begin.
     UpdateState ProcessUpdateState();
 
+  public:
+    // Initiate the merge if necessary, then wait for the merge to finish.
+    // See InitiateMerge() and ProcessUpdateState() for details.
+    // Returns:
+    //   - None if no merge to initiate
+    //   - MergeCompleted if merge is completed
+    //   - other states indicating an error has occurred
+    UpdateState InitiateMergeAndWait();
+
     // 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 63d97d07b..72bd30835 100644
--- a/fs_mgr/libsnapshot/snapshot.cpp
+++ b/fs_mgr/libsnapshot/snapshot.cpp
@@ -2089,5 +2089,23 @@ std::unique_ptr<AutoDevice> SnapshotManager::EnsureMetadataMounted() {
     return AutoUnmountDevice::New(device_->GetMetadataDir());
 }
 
+UpdateState SnapshotManager::InitiateMergeAndWait() {
+    auto state = GetUpdateState();
+    if (state == UpdateState::None) {
+        LOG(INFO) << "Can't find any snapshot to merge.";
+        return state;
+    }
+    if (state == UpdateState::Unverified) {
+        if (!InitiateMerge()) {
+            LOG(ERROR) << "Failed to initiate merge.";
+            return state;
+        }
+    }
+
+    // All other states can be handled by ProcessUpdateState.
+    LOG(INFO) << "Waiting for any merge to complete. This can take up to 1 minute.";
+    return ProcessUpdateState();
+}
+
 }  // namespace snapshot
 }  // namespace android
diff --git a/fs_mgr/libsnapshot/snapshot_test.cpp b/fs_mgr/libsnapshot/snapshot_test.cpp
index 57285821d..c9bf5b839 100644
--- a/fs_mgr/libsnapshot/snapshot_test.cpp
+++ b/fs_mgr/libsnapshot/snapshot_test.cpp
@@ -847,8 +847,7 @@ TEST_F(SnapshotUpdateTest, FullUpdateFlow) {
     }
 
     // Initiate the merge and wait for it to be completed.
-    ASSERT_TRUE(init->InitiateMerge());
-    ASSERT_EQ(UpdateState::MergeCompleted, init->ProcessUpdateState());
+    ASSERT_EQ(UpdateState::MergeCompleted, init->InitiateMergeAndWait());
 
     // Check that the target partitions have the same content after the merge.
     for (const auto& name : {"sys_b", "vnd_b", "prd_b"}) {
@@ -1052,8 +1051,7 @@ TEST_F(SnapshotUpdateTest, ReclaimCow) {
 
     // 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());
+    ASSERT_EQ(UpdateState::MergeCompleted, new_sm->InitiateMergeAndWait());
 
     // Execute the second update.
     ASSERT_TRUE(new_sm->BeginUpdate());
diff --git a/fs_mgr/libsnapshot/snapshotctl.cpp b/fs_mgr/libsnapshot/snapshotctl.cpp
index d65320c01..1bc03574c 100644
--- a/fs_mgr/libsnapshot/snapshotctl.cpp
+++ b/fs_mgr/libsnapshot/snapshotctl.cpp
@@ -60,24 +60,11 @@ bool MergeCmdHandler(int argc, char** argv) {
         android::base::InitLogging(argv, &android::base::StdioLogger);
     }
 
-    auto sm = SnapshotManager::New();
+    auto state = SnapshotManager::New()->InitiateMergeAndWait();
 
-    auto state = sm->GetUpdateState();
     if (state == UpdateState::None) {
-        LOG(INFO) << "Can't find any snapshot to merge.";
         return true;
     }
-    if (state == UpdateState::Unverified) {
-        if (!sm->InitiateMerge()) {
-            LOG(ERROR) << "Failed to initiate merge.";
-            return false;
-        }
-    }
-
-    // All other states can be handled by ProcessUpdateState.
-    LOG(INFO) << "Waiting for any merge to complete. This can take up to 1 minute.";
-    state = SnapshotManager::New()->ProcessUpdateState();
-
     if (state == UpdateState::MergeCompleted) {
         auto end = std::chrono::steady_clock::now();
         auto passed = std::chrono::duration_cast<std::chrono::milliseconds>(end - begin).count();