libsnapuserd: Pause/Resume Snapshot merge

Add API to support pause and resume of snapshot merge

Bug: 386142969
Test: Reboot device when snapshot merge is in progress
Change-Id: I6d1e36895a6b1b14a4e915d16afbd844a14f6175
Signed-off-by: Akilesh Kailash <akailash@google.com>
diff --git a/fs_mgr/libsnapshot/snapuserd/include/snapuserd/snapuserd_client.h b/fs_mgr/libsnapshot/snapuserd/include/snapuserd/snapuserd_client.h
index ede92dd..453627c 100644
--- a/fs_mgr/libsnapshot/snapuserd/include/snapuserd/snapuserd_client.h
+++ b/fs_mgr/libsnapshot/snapuserd/include/snapuserd/snapuserd_client.h
@@ -108,6 +108,9 @@
 
     // Notify init that snapuserd daemon is ready post selinux transition
     void NotifyTransitionDaemonIsReady();
+
+    // Pause Merge threads
+    bool PauseMerge();
 };
 
 }  // namespace snapshot
diff --git a/fs_mgr/libsnapshot/snapuserd/snapuserd_client.cpp b/fs_mgr/libsnapshot/snapuserd/snapuserd_client.cpp
index 7c820f3..0497c65 100644
--- a/fs_mgr/libsnapshot/snapuserd/snapuserd_client.cpp
+++ b/fs_mgr/libsnapshot/snapuserd/snapuserd_client.cpp
@@ -52,6 +52,7 @@
             return false;
         }
     }
+
     if (!android::base::WaitForProperty("snapuserd.ready", "true", 10s)) {
         LOG(ERROR) << "Timed out waiting for snapuserd to be ready.";
         return false;
@@ -389,5 +390,13 @@
     }
 }
 
+bool SnapuserdClient::PauseMerge() {
+    if (!Sendmsg("pause_merge")) {
+        LOG(ERROR) << "Failed to pause snapshot merge.";
+        return false;
+    }
+    return true;
+}
+
 }  // namespace snapshot
 }  // namespace android
diff --git a/fs_mgr/libsnapshot/snapuserd/user-space-merge/handler_manager.cpp b/fs_mgr/libsnapshot/snapuserd/user-space-merge/handler_manager.cpp
index fdd9cce..cf507e3 100644
--- a/fs_mgr/libsnapshot/snapuserd/user-space-merge/handler_manager.cpp
+++ b/fs_mgr/libsnapshot/snapuserd/user-space-merge/handler_manager.cpp
@@ -383,5 +383,25 @@
     return dm_users_.end();
 }
 
+void SnapshotHandlerManager::PauseMerge() {
+    std::lock_guard<std::mutex> guard(lock_);
+
+    for (auto iter = dm_users_.begin(); iter != dm_users_.end(); iter++) {
+        if (!(*iter)->ThreadTerminated()) {
+            (*iter)->snapuserd()->PauseMergeThreads();
+        }
+    }
+}
+
+void SnapshotHandlerManager::ResumeMerge() {
+    std::lock_guard<std::mutex> guard(lock_);
+
+    for (auto iter = dm_users_.begin(); iter != dm_users_.end(); iter++) {
+        if (!(*iter)->ThreadTerminated()) {
+            (*iter)->snapuserd()->ResumeMergeThreads();
+        }
+    }
+}
+
 }  // namespace snapshot
 }  // namespace android
diff --git a/fs_mgr/libsnapshot/snapuserd/user-space-merge/handler_manager.h b/fs_mgr/libsnapshot/snapuserd/user-space-merge/handler_manager.h
index c6301d4..89f3461 100644
--- a/fs_mgr/libsnapshot/snapuserd/user-space-merge/handler_manager.h
+++ b/fs_mgr/libsnapshot/snapuserd/user-space-merge/handler_manager.h
@@ -86,6 +86,12 @@
 
     // Disable partition verification
     virtual void DisableVerification() = 0;
+
+    // Pause Merge threads
+    virtual void PauseMerge() = 0;
+
+    // Resume Merge threads
+    virtual void ResumeMerge() = 0;
 };
 
 class SnapshotHandlerManager final : public ISnapshotHandlerManager {
@@ -107,6 +113,8 @@
     double GetMergePercentage() override;
     bool GetVerificationStatus() override;
     void DisableVerification() override { perform_verification_ = false; }
+    void PauseMerge() override;
+    void ResumeMerge() override;
 
   private:
     bool StartHandler(const std::shared_ptr<HandlerThread>& handler);
diff --git a/fs_mgr/libsnapshot/snapuserd/user-space-merge/merge_worker.cpp b/fs_mgr/libsnapshot/snapuserd/user-space-merge/merge_worker.cpp
index febb484..660082f 100644
--- a/fs_mgr/libsnapshot/snapuserd/user-space-merge/merge_worker.cpp
+++ b/fs_mgr/libsnapshot/snapuserd/user-space-merge/merge_worker.cpp
@@ -191,6 +191,9 @@
                                "down merge";
             return false;
         }
+
+        // Safe to check if there is a pause request.
+        snapuserd_->PauseMergeIfRequired();
     }
 
     // Any left over ops not flushed yet.
diff --git a/fs_mgr/libsnapshot/snapuserd/user-space-merge/snapuserd_core.h b/fs_mgr/libsnapshot/snapuserd/user-space-merge/snapuserd_core.h
index 2340b0b..924539f 100644
--- a/fs_mgr/libsnapshot/snapuserd/user-space-merge/snapuserd_core.h
+++ b/fs_mgr/libsnapshot/snapuserd/user-space-merge/snapuserd_core.h
@@ -175,6 +175,9 @@
     bool MergeInitiated() { return merge_initiated_; }
     bool MergeMonitored() { return merge_monitored_; }
     double GetMergePercentage() { return merge_completion_percentage_; }
+    void PauseMergeThreads();
+    void ResumeMergeThreads();
+    void PauseMergeIfRequired();
 
     // Merge Block State Transitions
     void SetMergeCompleted(size_t block_index);
@@ -255,6 +258,11 @@
     uint32_t cow_op_merge_size_ = 0;
     std::unique_ptr<UpdateVerify> update_verify_;
     std::shared_ptr<IBlockServerOpener> block_server_opener_;
+
+    // Pause merge threads
+    bool pause_merge_ = false;
+    std::mutex pause_merge_lock_;
+    std::condition_variable pause_merge_cv_;
 };
 
 std::ostream& operator<<(std::ostream& os, MERGE_IO_TRANSITION value);
diff --git a/fs_mgr/libsnapshot/snapuserd/user-space-merge/snapuserd_server.cpp b/fs_mgr/libsnapshot/snapuserd/user-space-merge/snapuserd_server.cpp
index 3bb8a30..a40617b 100644
--- a/fs_mgr/libsnapshot/snapuserd/user-space-merge/snapuserd_server.cpp
+++ b/fs_mgr/libsnapshot/snapuserd/user-space-merge/snapuserd_server.cpp
@@ -227,6 +227,9 @@
             return Sendmsg(fd, "fail");
         }
         return Sendmsg(fd, "success");
+    } else if (cmd == "pause_merge") {
+        handlers_->PauseMerge();
+        return true;
     } else {
         LOG(ERROR) << "Received unknown message type from client";
         Sendmsg(fd, "fail");
diff --git a/fs_mgr/libsnapshot/snapuserd/user-space-merge/snapuserd_test.cpp b/fs_mgr/libsnapshot/snapuserd/user-space-merge/snapuserd_test.cpp
index 469fd09..0790a19 100644
--- a/fs_mgr/libsnapshot/snapuserd/user-space-merge/snapuserd_test.cpp
+++ b/fs_mgr/libsnapshot/snapuserd/user-space-merge/snapuserd_test.cpp
@@ -934,6 +934,26 @@
     read_future.wait();
 }
 
+TEST_P(SnapuserdTest, Snapshot_MERGE_PAUSE_RESUME) {
+    if (!harness_->HasUserDevice()) {
+        GTEST_SKIP() << "Skipping snapshot read; not supported";
+    }
+    ASSERT_NO_FATAL_FAILURE(SetupDefault());
+    // Start the merge
+    ASSERT_TRUE(StartMerge());
+    std::this_thread::sleep_for(300ms);
+    // Pause merge
+    handlers_->PauseMerge();
+    // Issue I/O after pausing the merge and validate
+    auto read_future =
+            std::async(std::launch::async, &SnapuserdTest::ReadSnapshotDeviceAndValidate, this);
+    // Resume the merge
+    handlers_->ResumeMerge();
+    CheckMergeCompletion();
+    ValidateMerge();
+    read_future.wait();
+}
+
 TEST_P(SnapuserdTest, Snapshot_Merge_Resume) {
     ASSERT_NO_FATAL_FAILURE(SetupDefault());
     ASSERT_NO_FATAL_FAILURE(MergeInterrupt());
diff --git a/fs_mgr/libsnapshot/snapuserd/user-space-merge/snapuserd_transitions.cpp b/fs_mgr/libsnapshot/snapuserd/user-space-merge/snapuserd_transitions.cpp
index 714c641..90705f7 100644
--- a/fs_mgr/libsnapshot/snapuserd/user-space-merge/snapuserd_transitions.cpp
+++ b/fs_mgr/libsnapshot/snapuserd/user-space-merge/snapuserd_transitions.cpp
@@ -257,6 +257,19 @@
     return true;
 }
 
+void SnapshotHandler::PauseMergeIfRequired() {
+    {
+        std::unique_lock<std::mutex> lock(pause_merge_lock_);
+        while (pause_merge_) {
+            SNAP_LOG(INFO) << "Merge thread paused";
+            pause_merge_cv_.wait(lock);
+            if (!pause_merge_) {
+                SNAP_LOG(INFO) << "Merge thread resumed";
+            }
+        }
+    }
+}
+
 // Invoked by RA thread - Waits for merge thread to finish merging
 // RA Block N - RA thread would be ready will with Block N+1 but
 // will wait to merge thread to finish Block N. Once Block N
@@ -281,8 +294,13 @@
             }
             return false;
         }
-        return true;
     }
+
+    // This is a safe place to check if the RA thread should be
+    // paused. Since the scratch space isn't flushed yet, it is safe
+    // to wait here until resume is invoked.
+    PauseMergeIfRequired();
+    return true;
 }
 
 // Invoked by Merge thread - Notify RA thread about Merge completion
@@ -297,6 +315,11 @@
     }
 
     cv.notify_all();
+
+    // This is a safe place to check if the merge thread should be
+    // paused. The data from the scratch space is merged to disk and is safe
+    // to wait.
+    PauseMergeIfRequired();
 }
 
 // The following transitions are mostly in the failure paths
@@ -393,6 +416,20 @@
     merge_complete_ = true;
 }
 
+void SnapshotHandler::PauseMergeThreads() {
+    {
+        std::lock_guard<std::mutex> lock(pause_merge_lock_);
+        pause_merge_ = true;
+    }
+}
+
+void SnapshotHandler::ResumeMergeThreads() {
+    {
+        std::lock_guard<std::mutex> lock(pause_merge_lock_);
+        pause_merge_ = false;
+    }
+}
+
 std::string SnapshotHandler::GetMergeStatus() {
     bool merge_not_initiated = false;
     bool merge_monitored = false;