Merge changes from topic "snapshotctl_cleanup" into rvc-dev

* changes:
  libsnapshot: SnapshotMergeStats::Start/Finish
  libsnapshot: Expose SnapshotMergeStats
  libsnapshot: delete WaitForMerge.
  libsnapshot: handle errors in RemoveAllUpdateState appropriately.
  libsnapshot: remove snapshots properly after flashing
  libsnapshot: RemoveUpdateState on rollback.
  libsnapshot: NeedSnapshotsInFirstStageMount don't test for IsRecovery
  Allow ProcessUpdateState to be paused.
  libsnapshot: Re-expose InitiateMerge and ProcessUpdateState
  libsnapshot: Add prolog to RemoveAllUpdateStates.
  snapshotctl don't auto-merge.
  snapshotctl: init reports merge statistics
  libsnapshot/test: Re-enable the failing tests
diff --git a/CleanSpec.mk b/CleanSpec.mk
index c84bd24..0a534a2 100644
--- a/CleanSpec.mk
+++ b/CleanSpec.mk
@@ -90,3 +90,4 @@
 $(call add-clean-step, rm -rf $(PRODUCT_OUT)/debug_ramdisk/product_services)
 $(call add-clean-step, find $(PRODUCT_OUT) -type l -name "charger" -print0 | xargs -0 rm -f)
 $(call add-clean-step, rm -f $(PRODUCT_OUT)/system/bin/adbd)
+$(call add-clean-step, rm -rf $(PRODUCT_OUT)/system/etc/init/snapshotctl.rc)
diff --git a/fs_mgr/TEST_MAPPING b/fs_mgr/TEST_MAPPING
index 705d4e3..676f446 100644
--- a/fs_mgr/TEST_MAPPING
+++ b/fs_mgr/TEST_MAPPING
@@ -13,7 +13,7 @@
       "name": "fiemap_writer_test"
     },
     {
-      "name": "vts_libsnapshot_test_presubmit"
+      "name": "vts_libsnapshot_test"
     }
   ]
 }
diff --git a/fs_mgr/libsnapshot/Android.bp b/fs_mgr/libsnapshot/Android.bp
index 2d62617..0a0a21d 100644
--- a/fs_mgr/libsnapshot/Android.bp
+++ b/fs_mgr/libsnapshot/Android.bp
@@ -208,14 +208,6 @@
     defaults: ["libsnapshot_test_defaults"],
 }
 
-cc_test {
-    name: "vts_libsnapshot_test_presubmit",
-    defaults: ["libsnapshot_test_defaults"],
-    cppflags: [
-        "-DSKIP_TEST_IN_PRESUBMIT",
-    ],
-}
-
 cc_binary {
     name: "snapshotctl",
     srcs: [
@@ -243,7 +235,4 @@
         // TODO(b/148818798): remove when parent bug is fixed.
         "libutilscallstack",
     ],
-    init_rc: [
-        "snapshotctl.rc",
-    ],
 }
diff --git a/fs_mgr/libsnapshot/include/libsnapshot/return.h b/fs_mgr/libsnapshot/include/libsnapshot/return.h
index 1f132fa..dedc445 100644
--- a/fs_mgr/libsnapshot/include/libsnapshot/return.h
+++ b/fs_mgr/libsnapshot/include/libsnapshot/return.h
@@ -30,7 +30,6 @@
     enum class ErrorCode : int32_t {
         SUCCESS = static_cast<int32_t>(FiemapStatus::ErrorCode::SUCCESS),
         ERROR = static_cast<int32_t>(FiemapStatus::ErrorCode::ERROR),
-        NEEDS_REBOOT = ERROR + 1,
         NO_SPACE = static_cast<int32_t>(FiemapStatus::ErrorCode::NO_SPACE),
     };
     ErrorCode error_code() const { return error_code_; }
@@ -43,7 +42,6 @@
     static Return Ok() { return Return(ErrorCode::SUCCESS); }
     static Return Error() { return Return(ErrorCode::ERROR); }
     static Return NoSpace(uint64_t size) { return Return(ErrorCode::NO_SPACE, size); }
-    static Return NeedsReboot() { return Return(ErrorCode::NEEDS_REBOOT); }
     // Does not set required_size_ properly even when status.error_code() == NO_SPACE.
     explicit Return(const FiemapStatus& status)
         : error_code_(FromFiemapStatusErrorCode(status.error_code())), required_size_(0) {}
diff --git a/fs_mgr/libsnapshot/include/libsnapshot/snapshot.h b/fs_mgr/libsnapshot/include/libsnapshot/snapshot.h
index b440c71..32345d2 100644
--- a/fs_mgr/libsnapshot/include/libsnapshot/snapshot.h
+++ b/fs_mgr/libsnapshot/include/libsnapshot/snapshot.h
@@ -140,7 +140,6 @@
     // Before calling this function, all snapshots must be mapped.
     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();
@@ -149,7 +148,11 @@
     // /data is mounted.
     //
     // If a merge is in progress, this function will block until the merge is
-    // completed. If a merge or update was cancelled, this will clean up any
+    // completed.
+    //    - Callback is called periodically during the merge. If callback()
+    //      returns false during the merge, ProcessUpdateState() will pause
+    //      and returns Merging.
+    // If a merge or update was cancelled, this will clean up any
     // update artifacts and return.
     //
     // Note that after calling this, GetUpdateState() may still return that a
@@ -169,9 +172,9 @@
     //
     // The optional callback allows the caller to periodically check the
     // progress with GetUpdateState().
-    UpdateState ProcessUpdateState(const std::function<void()>& callback = {});
+    UpdateState ProcessUpdateState(const std::function<bool()>& callback = {},
+                                   const std::function<bool()>& before_cancel = {});
 
-  public:
     // Initiate the merge if necessary, then wait for the merge to finish.
     // See InitiateMerge() and ProcessUpdateState() for details.
     // Returns:
@@ -179,16 +182,8 @@
     //   - Unverified if called on the source slot
     //   - MergeCompleted if merge is completed
     //   - other states indicating an error has occurred
-    UpdateState InitiateMergeAndWait(SnapshotMergeReport* report = nullptr);
-
-    // 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:
-    //   - Return::Ok(): there is no merge or merge finishes
-    //   - Return::NeedsReboot(): merge finishes but need a reboot before
-    //     applying the next update.
-    //   - Return::Error(): other irrecoverable errors
-    Return WaitForMerge();
+    UpdateState InitiateMergeAndWait(SnapshotMergeReport* report = nullptr,
+                                     const std::function<bool()>& before_cancel = {});
 
     // Find the status of the current update, if any.
     //
@@ -375,14 +370,23 @@
 
     // Check for a cancelled or rolled back merge, returning true if such a
     // condition was detected and handled.
-    bool HandleCancelledUpdate(LockedFile* lock);
+    bool HandleCancelledUpdate(LockedFile* lock, const std::function<bool()>& before_cancel);
 
     // Helper for HandleCancelledUpdate. Assumes booting from new slot.
     bool AreAllSnapshotsCancelled(LockedFile* lock);
 
+    // Determine whether partition names in |snapshots| have been flashed and
+    // store result to |out|.
+    // Return true if values are successfully retrieved and false on error
+    // (e.g. super partition metadata cannot be read). When it returns true,
+    // |out| stores true for partitions that have been flashed and false for
+    // partitions that have not been flashed.
+    bool GetSnapshotFlashingStatus(LockedFile* lock, const std::vector<std::string>& snapshots,
+                                   std::map<std::string, bool>* out);
+
     // Remove artifacts created by the update process, such as snapshots, and
     // set the update state to None.
-    bool RemoveAllUpdateState(LockedFile* lock);
+    bool RemoveAllUpdateState(LockedFile* lock, const std::function<bool()>& prolog = {});
 
     // Interact with /metadata/ota.
     std::unique_ptr<LockedFile> OpenLock(int lock_flags);
@@ -437,8 +441,8 @@
     //   UpdateState::MergeCompleted
     //   UpdateState::MergeFailed
     //   UpdateState::MergeNeedsReboot
-    UpdateState CheckMergeState();
-    UpdateState CheckMergeState(LockedFile* lock);
+    UpdateState CheckMergeState(const std::function<bool()>& before_cancel);
+    UpdateState CheckMergeState(LockedFile* lock, const std::function<bool()>& before_cancel);
     UpdateState CheckTargetMergeState(LockedFile* lock, const std::string& name);
 
     // Interact with status files under /metadata/ota/snapshots.
@@ -513,6 +517,11 @@
 
     std::string ReadUpdateSourceSlotSuffix();
 
+    // Helper for RemoveAllSnapshots.
+    // Check whether |name| should be deleted as a snapshot name.
+    bool ShouldDeleteSnapshot(LockedFile* lock, const std::map<std::string, bool>& flashing_status,
+                              Slot current_slot, const std::string& name);
+
     std::string gsid_dir_;
     std::string metadata_dir_;
     std::unique_ptr<IDeviceInfo> device_;
diff --git a/fs_mgr/libsnapshot/include/libsnapshot/snapshot_stats.h b/fs_mgr/libsnapshot/include/libsnapshot/snapshot_stats.h
new file mode 100644
index 0000000..91dd34f
--- /dev/null
+++ b/fs_mgr/libsnapshot/include/libsnapshot/snapshot_stats.h
@@ -0,0 +1,59 @@
+// Copyright (C) 2020 The Android Open Source Project
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+
+#include <chrono>
+#include <memory>
+
+#include <android/snapshot/snapshot.pb.h>
+#include <libsnapshot/snapshot.h>
+
+namespace android {
+namespace snapshot {
+
+class SnapshotMergeStats {
+  public:
+    // Not thread safe.
+    static SnapshotMergeStats* GetInstance(SnapshotManager& manager);
+
+    // Called when merge starts or resumes.
+    bool Start();
+    void set_state(android::snapshot::UpdateState state);
+
+    // Called when merge ends. Properly clean up permanent storage.
+    class Result {
+      public:
+        virtual ~Result() {}
+        virtual const SnapshotMergeReport& report() const = 0;
+        // Time between successful Start() / Resume() to Finish().
+        virtual std::chrono::steady_clock::duration merge_time() const = 0;
+    };
+    std::unique_ptr<Result> Finish();
+
+  private:
+    bool ReadState();
+    bool WriteState();
+    bool DeleteState();
+    SnapshotMergeStats(const std::string& path);
+
+    std::string path_;
+    SnapshotMergeReport report_;
+    // Time of the last successful Start() / Resume() call.
+    std::chrono::time_point<std::chrono::steady_clock> start_time_;
+    bool running_{false};
+};
+
+}  // namespace snapshot
+}  // namespace android
diff --git a/fs_mgr/libsnapshot/return.cpp b/fs_mgr/libsnapshot/return.cpp
index 6559c12..cc64af5 100644
--- a/fs_mgr/libsnapshot/return.cpp
+++ b/fs_mgr/libsnapshot/return.cpp
@@ -24,8 +24,6 @@
     switch (error_code()) {
         case ErrorCode::ERROR:
             return "Error";
-        case ErrorCode::NEEDS_REBOOT:
-            return "Retry after reboot";
         case ErrorCode::SUCCESS:
             [[fallthrough]];
         case ErrorCode::NO_SPACE:
diff --git a/fs_mgr/libsnapshot/snapshot.cpp b/fs_mgr/libsnapshot/snapshot.cpp
index 2fe06fb..61fc2df 100644
--- a/fs_mgr/libsnapshot/snapshot.cpp
+++ b/fs_mgr/libsnapshot/snapshot.cpp
@@ -43,10 +43,10 @@
 #endif
 
 #include <android/snapshot/snapshot.pb.h>
+#include <libsnapshot/snapshot_stats.h>
 #include "device_info.h"
 #include "partition_cow_creator.h"
 #include "snapshot_metadata_updater.h"
-#include "snapshot_stats.h"
 #include "utility.h"
 
 namespace android {
@@ -219,7 +219,12 @@
     return true;
 }
 
-bool SnapshotManager::RemoveAllUpdateState(LockedFile* lock) {
+bool SnapshotManager::RemoveAllUpdateState(LockedFile* lock, const std::function<bool()>& prolog) {
+    if (prolog && !prolog()) {
+        LOG(WARNING) << "Can't RemoveAllUpdateState: prolog failed.";
+        return false;
+    }
+
     LOG(INFO) << "Removing all update state.";
 
 #ifdef LIBSNAPSHOT_USE_CALLSTACK
@@ -789,9 +794,10 @@
 // Note that when a merge fails, we will *always* try again to complete the
 // merge each time the device boots. There is no harm in doing so, and if
 // the problem was transient, we might manage to get a new outcome.
-UpdateState SnapshotManager::ProcessUpdateState(const std::function<void()>& callback) {
+UpdateState SnapshotManager::ProcessUpdateState(const std::function<bool()>& callback,
+                                                const std::function<bool()>& before_cancel) {
     while (true) {
-        UpdateState state = CheckMergeState();
+        UpdateState state = CheckMergeState(before_cancel);
         if (state == UpdateState::MergeFailed) {
             AcknowledgeMergeFailure();
         }
@@ -801,8 +807,8 @@
             return state;
         }
 
-        if (callback) {
-            callback();
+        if (callback && !callback()) {
+            return state;
         }
 
         // This wait is not super time sensitive, so we have a relatively
@@ -811,24 +817,27 @@
     }
 }
 
-UpdateState SnapshotManager::CheckMergeState() {
+UpdateState SnapshotManager::CheckMergeState(const std::function<bool()>& before_cancel) {
     auto lock = LockExclusive();
     if (!lock) {
         return UpdateState::MergeFailed;
     }
 
-    UpdateState state = CheckMergeState(lock.get());
+    UpdateState state = CheckMergeState(lock.get(), before_cancel);
     if (state == UpdateState::MergeCompleted) {
         // Do this inside the same lock. Failures get acknowledged without the
         // lock, because flock() might have failed.
         AcknowledgeMergeSuccess(lock.get());
     } else if (state == UpdateState::Cancelled) {
-        RemoveAllUpdateState(lock.get());
+        if (!RemoveAllUpdateState(lock.get(), before_cancel)) {
+            return ReadSnapshotUpdateStatus(lock.get()).state();
+        }
     }
     return state;
 }
 
-UpdateState SnapshotManager::CheckMergeState(LockedFile* lock) {
+UpdateState SnapshotManager::CheckMergeState(LockedFile* lock,
+                                             const std::function<bool()>& before_cancel) {
     UpdateState state = ReadUpdateState(lock);
     switch (state) {
         case UpdateState::None:
@@ -849,7 +858,7 @@
             // This is an edge case. Normally cancelled updates are detected
             // via the merge poll below, but if we never started a merge, we
             // need to also check here.
-            if (HandleCancelledUpdate(lock)) {
+            if (HandleCancelledUpdate(lock, before_cancel)) {
                 return UpdateState::Cancelled;
             }
             return state;
@@ -1169,7 +1178,8 @@
     return true;
 }
 
-bool SnapshotManager::HandleCancelledUpdate(LockedFile* lock) {
+bool SnapshotManager::HandleCancelledUpdate(LockedFile* lock,
+                                            const std::function<bool()>& before_cancel) {
     auto slot = GetCurrentSlot();
     if (slot == Slot::Unknown) {
         return false;
@@ -1177,15 +1187,30 @@
 
     // If all snapshots were reflashed, then cancel the entire update.
     if (AreAllSnapshotsCancelled(lock)) {
-        RemoveAllUpdateState(lock);
-        return true;
+        LOG(WARNING) << "Detected re-flashing, cancelling unverified update.";
+        return RemoveAllUpdateState(lock, before_cancel);
     }
 
-    // This unverified update might be rolled back, or it might not (b/147347110
-    // comment #77). Take no action, as update_engine is responsible for deciding
-    // whether to cancel.
-    LOG(ERROR) << "Update state is being processed before reboot, taking no action.";
-    return false;
+    // If update has been rolled back, then cancel the entire update.
+    // Client (update_engine) is responsible for doing additional cleanup work on its own states
+    // when ProcessUpdateState() returns UpdateState::Cancelled.
+    auto current_slot = GetCurrentSlot();
+    if (current_slot != Slot::Source) {
+        LOG(INFO) << "Update state is being processed while booting at " << current_slot
+                  << " slot, taking no action.";
+        return false;
+    }
+
+    // current_slot == Source. Attempt to detect rollbacks.
+    if (access(GetRollbackIndicatorPath().c_str(), F_OK) != 0) {
+        // This unverified update is not attempted. Take no action.
+        PLOG(INFO) << "Rollback indicator not detected. "
+                   << "Update state is being processed before reboot, taking no action.";
+        return false;
+    }
+
+    LOG(WARNING) << "Detected rollback, cancelling unverified update.";
+    return RemoveAllUpdateState(lock, before_cancel);
 }
 
 std::unique_ptr<LpMetadata> SnapshotManager::ReadCurrentMetadata() {
@@ -1219,6 +1244,28 @@
         return true;
     }
 
+    std::map<std::string, bool> flashing_status;
+
+    if (!GetSnapshotFlashingStatus(lock, snapshots, &flashing_status)) {
+        LOG(WARNING) << "Failed to determine whether partitions have been flashed. Not"
+                     << "removing update states.";
+        return false;
+    }
+
+    bool all_snapshots_cancelled = std::all_of(flashing_status.begin(), flashing_status.end(),
+                                               [](const auto& pair) { return pair.second; });
+
+    if (all_snapshots_cancelled) {
+        LOG(WARNING) << "All partitions are re-flashed after update, removing all update states.";
+    }
+    return all_snapshots_cancelled;
+}
+
+bool SnapshotManager::GetSnapshotFlashingStatus(LockedFile* lock,
+                                                const std::vector<std::string>& snapshots,
+                                                std::map<std::string, bool>* out) {
+    CHECK(lock);
+
     auto source_slot_suffix = ReadUpdateSourceSlotSuffix();
     if (source_slot_suffix.empty()) {
         return false;
@@ -1244,20 +1291,17 @@
         return false;
     }
 
-    bool all_snapshots_cancelled = true;
     for (const auto& snapshot_name : snapshots) {
         if (GetMetadataPartitionState(*metadata, snapshot_name) ==
             MetadataPartitionState::Updated) {
-            all_snapshots_cancelled = false;
-            continue;
+            out->emplace(snapshot_name, false);
+        } else {
+            // Delete snapshots for partitions that are re-flashed after the update.
+            LOG(WARNING) << "Detected re-flashing of partition " << snapshot_name << ".";
+            out->emplace(snapshot_name, true);
         }
-        // Delete snapshots for partitions that are re-flashed after the update.
-        LOG(WARNING) << "Detected re-flashing of partition " << snapshot_name << ".";
     }
-    if (all_snapshots_cancelled) {
-        LOG(WARNING) << "All partitions are re-flashed after update, removing all update states.";
-    }
-    return all_snapshots_cancelled;
+    return true;
 }
 
 bool SnapshotManager::RemoveAllSnapshots(LockedFile* lock) {
@@ -1267,10 +1311,38 @@
         return false;
     }
 
+    std::map<std::string, bool> flashing_status;
+    if (!GetSnapshotFlashingStatus(lock, snapshots, &flashing_status)) {
+        LOG(WARNING) << "Failed to get flashing status";
+    }
+
+    auto current_slot = GetCurrentSlot();
     bool ok = true;
     bool has_mapped_cow_images = false;
     for (const auto& name : snapshots) {
-        if (!UnmapPartitionWithSnapshot(lock, name) || !DeleteSnapshot(lock, name)) {
+        // If booting off source slot, it is okay to unmap and delete all the snapshots.
+        // If boot indicator is missing, update state is None or Initiated, so
+        //   it is also okay to unmap and delete all the snapshots.
+        // If booting off target slot,
+        //  - should not unmap because:
+        //    - In Android mode, snapshots are not mapped, but
+        //      filesystems are mounting off dm-linear targets directly.
+        //    - In recovery mode, assume nothing is mapped, so it is optional to unmap.
+        //  - If partition is flashed or unknown, it is okay to delete snapshots.
+        //    Otherwise (UPDATED flag), only delete snapshots if they are not mapped
+        //    as dm-snapshot (for example, after merge completes).
+        bool should_unmap = current_slot != Slot::Target;
+        bool should_delete = ShouldDeleteSnapshot(lock, flashing_status, current_slot, name);
+
+        bool partition_ok = true;
+        if (should_unmap && !UnmapPartitionWithSnapshot(lock, name)) {
+            partition_ok = false;
+        }
+        if (partition_ok && should_delete && !DeleteSnapshot(lock, name)) {
+            partition_ok = false;
+        }
+
+        if (!partition_ok) {
             // Remember whether or not we were able to unmap the cow image.
             auto cow_image_device = GetCowImageDeviceName(name);
             has_mapped_cow_images |=
@@ -1293,6 +1365,34 @@
     return ok;
 }
 
+// See comments in RemoveAllSnapshots().
+bool SnapshotManager::ShouldDeleteSnapshot(LockedFile* lock,
+                                           const std::map<std::string, bool>& flashing_status,
+                                           Slot current_slot, const std::string& name) {
+    if (current_slot != Slot::Target) {
+        return true;
+    }
+    auto it = flashing_status.find(name);
+    if (it == flashing_status.end()) {
+        LOG(WARNING) << "Can't determine flashing status for " << name;
+        return true;
+    }
+    if (it->second) {
+        // partition flashed, okay to delete obsolete snapshots
+        return true;
+    }
+    // partition updated, only delete if not dm-snapshot
+    SnapshotStatus status;
+    if (!ReadSnapshotStatus(lock, name, &status)) {
+        LOG(WARNING) << "Unable to read snapshot status for " << name
+                     << ", guessing snapshot device name";
+        auto extra_name = GetSnapshotExtraDeviceName(name);
+        return !IsSnapshotDevice(name) && !IsSnapshotDevice(extra_name);
+    }
+    auto dm_name = GetSnapshotDeviceName(name, status);
+    return !IsSnapshotDevice(dm_name);
+}
+
 UpdateState SnapshotManager::GetUpdateState(double* progress) {
     // If we've never started an update, the state file won't exist.
     auto state_file = GetStateFilePath();
@@ -1379,12 +1479,14 @@
     auto slot = GetCurrentSlot();
 
     if (slot != Slot::Target) {
-        if (slot == Slot::Source && !device_->IsRecovery()) {
+        if (slot == Slot::Source) {
             // Device is rebooting into the original slot, so mark this as a
             // rollback.
             auto path = GetRollbackIndicatorPath();
             if (!android::base::WriteStringToFile("1", path)) {
                 PLOG(ERROR) << "Unable to write rollback indicator: " << path;
+            } else {
+                LOG(INFO) << "Rollback detected, writing rollback indicator to " << path;
             }
         }
         LOG(INFO) << "Not booting from new slot. Will not mount snapshots.";
@@ -2388,7 +2490,8 @@
     return AutoUnmountDevice::New(device_->GetMetadataDir());
 }
 
-UpdateState SnapshotManager::InitiateMergeAndWait(SnapshotMergeReport* stats_report) {
+UpdateState SnapshotManager::InitiateMergeAndWait(SnapshotMergeReport* stats_report,
+                                                  const std::function<bool()>& before_cancel) {
     {
         auto lock = LockExclusive();
         // Sync update state from file with bootloader.
@@ -2398,23 +2501,24 @@
         }
     }
 
-    SnapshotMergeStats merge_stats(*this);
+    auto merge_stats = SnapshotMergeStats::GetInstance(*this);
 
     unsigned int last_progress = 0;
-    auto callback = [&]() -> void {
+    auto callback = [&]() -> bool {
         double progress;
         GetUpdateState(&progress);
         if (last_progress < static_cast<unsigned int>(progress)) {
             last_progress = progress;
             LOG(INFO) << "Waiting for merge to complete: " << last_progress << "%.";
         }
+        return true;  // continue
     };
 
     LOG(INFO) << "Waiting for any previous merge request to complete. "
               << "This can take up to several minutes.";
-    merge_stats.Resume();
-    auto state = ProcessUpdateState(callback);
-    merge_stats.set_state(state);
+    merge_stats->Start();
+    auto state = ProcessUpdateState(callback, before_cancel);
+    merge_stats->set_state(state);
     if (state == UpdateState::None) {
         LOG(INFO) << "Can't find any snapshot to merge.";
         return state;
@@ -2425,10 +2529,6 @@
             return state;
         }
 
-        // This is the first snapshot merge that is requested after OTA. We can
-        // initialize the merge duration statistics.
-        merge_stats.Start();
-
         if (!InitiateMerge()) {
             LOG(ERROR) << "Failed to initiate merge.";
             return state;
@@ -2436,43 +2536,22 @@
         // All other states can be handled by ProcessUpdateState.
         LOG(INFO) << "Waiting for merge to complete. This can take up to several minutes.";
         last_progress = 0;
-        state = ProcessUpdateState(callback);
-        merge_stats.set_state(state);
+        state = ProcessUpdateState(callback, before_cancel);
+        merge_stats->set_state(state);
     }
 
     LOG(INFO) << "Merge finished with state \"" << state << "\".";
     if (stats_report) {
-        *stats_report = merge_stats.GetReport();
+        auto result = merge_stats->Finish();
+        if (result) {
+            *stats_report = result->report();
+        } else {
+            LOG(WARNING) << "SnapshotMergeStatus::Finish failed.";
+        }
     }
     return state;
 }
 
-Return 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;
-        switch (state) {
-            case UpdateState::None:
-                [[fallthrough]];
-            case UpdateState::MergeCompleted:
-                [[fallthrough]];
-            case UpdateState::Cancelled:
-                return Return::Ok();
-            case UpdateState::MergeNeedsReboot:
-                return Return::NeedsReboot();
-            default:
-                return Return::Error();
-        }
-    }
-}
-
 bool SnapshotManager::HandleImminentDataWipe(const std::function<void()>& callback) {
     if (!device_->IsRecovery()) {
         LOG(ERROR) << "Data wipes are only allowed in recovery.";
@@ -2503,7 +2582,10 @@
         return false;
     }
 
-    UpdateState state = ProcessUpdateState(callback);
+    UpdateState state = ProcessUpdateState([&]() -> bool {
+        callback();
+        return true;
+    });
     LOG(INFO) << "Update state in recovery: " << state;
     switch (state) {
         case UpdateState::MergeFailed:
diff --git a/fs_mgr/libsnapshot/snapshot_stats.cpp b/fs_mgr/libsnapshot/snapshot_stats.cpp
index 635b47d..5da7b98 100644
--- a/fs_mgr/libsnapshot/snapshot_stats.cpp
+++ b/fs_mgr/libsnapshot/snapshot_stats.cpp
@@ -12,7 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-#include "snapshot_stats.h"
+#include <libsnapshot/snapshot_stats.h>
 
 #include <sstream>
 
@@ -23,22 +23,17 @@
 namespace android {
 namespace snapshot {
 
-SnapshotMergeStats::SnapshotMergeStats(SnapshotManager& parent) : parent_(parent) {
-    init_time_ = std::chrono::steady_clock::now();
+SnapshotMergeStats* SnapshotMergeStats::GetInstance(SnapshotManager& parent) {
+    static SnapshotMergeStats g_instance(parent.GetMergeStateFilePath());
+    CHECK(g_instance.path_ == parent.GetMergeStateFilePath());
+    return &g_instance;
 }
 
-SnapshotMergeStats::~SnapshotMergeStats() {
-    std::string error;
-    auto file_path = parent_.GetMergeStateFilePath();
-    if (!android::base::RemoveFileIfExists(file_path, &error)) {
-        LOG(ERROR) << "Failed to remove merge statistics file " << file_path << ": " << error;
-        return;
-    }
-}
+SnapshotMergeStats::SnapshotMergeStats(const std::string& path) : path_(path), running_(false) {}
 
 bool SnapshotMergeStats::ReadState() {
     std::string contents;
-    if (!android::base::ReadFileToString(parent_.GetMergeStateFilePath(), &contents)) {
+    if (!android::base::ReadFileToString(path_, &contents)) {
         PLOG(INFO) << "Read merge statistics file failed";
         return false;
     }
@@ -55,34 +50,73 @@
         LOG(ERROR) << "Unable to serialize SnapshotMergeStats.";
         return false;
     }
-    auto file_path = parent_.GetMergeStateFilePath();
-    if (!WriteStringToFileAtomic(contents, file_path)) {
+    if (!WriteStringToFileAtomic(contents, path_)) {
         PLOG(ERROR) << "Could not write to merge statistics file";
         return false;
     }
     return true;
 }
 
-void SnapshotMergeStats::Start() {
-    report_.set_resume_count(0);
-    report_.set_state(UpdateState::None);
-    WriteState();
+bool SnapshotMergeStats::DeleteState() {
+    std::string error;
+    if (!android::base::RemoveFileIfExists(path_, &error)) {
+        LOG(ERROR) << "Failed to remove merge statistics file " << path_ << ": " << error;
+        return false;
+    }
+    return true;
 }
 
-void SnapshotMergeStats::Resume() {
-    if (!ReadState()) {
-        return;
+bool SnapshotMergeStats::Start() {
+    if (running_) {
+        LOG(ERROR) << "SnapshotMergeStats running_ == " << running_;
+        return false;
     }
-    report_.set_resume_count(report_.resume_count() + 1);
-    WriteState();
+    running_ = true;
+
+    start_time_ = std::chrono::steady_clock::now();
+    if (ReadState()) {
+        report_.set_resume_count(report_.resume_count() + 1);
+    } else {
+        report_.set_resume_count(0);
+        report_.set_state(UpdateState::None);
+    }
+
+    return WriteState();
 }
 
 void SnapshotMergeStats::set_state(android::snapshot::UpdateState state) {
     report_.set_state(state);
 }
 
-SnapshotMergeReport SnapshotMergeStats::GetReport() {
-    return report_;
+class SnapshotMergeStatsResultImpl : public SnapshotMergeStats::Result {
+  public:
+    SnapshotMergeStatsResultImpl(const SnapshotMergeReport& report,
+                                 std::chrono::steady_clock::duration merge_time)
+        : report_(report), merge_time_(merge_time) {}
+    const SnapshotMergeReport& report() const override { return report_; }
+    std::chrono::steady_clock::duration merge_time() const override { return merge_time_; }
+
+  private:
+    SnapshotMergeReport report_;
+    std::chrono::steady_clock::duration merge_time_;
+};
+
+std::unique_ptr<SnapshotMergeStats::Result> SnapshotMergeStats::Finish() {
+    if (!running_) {
+        LOG(ERROR) << "SnapshotMergeStats running_ == " << running_;
+        return nullptr;
+    }
+    running_ = false;
+
+    auto result = std::make_unique<SnapshotMergeStatsResultImpl>(
+            report_, std::chrono::steady_clock::now() - start_time_);
+
+    // We still want to report result if state is not deleted. Just leave
+    // it there and move on. A side effect is that it may be reported over and
+    // over again in the future, but there is nothing we can do.
+    (void)DeleteState();
+
+    return result;
 }
 
 }  // namespace snapshot
diff --git a/fs_mgr/libsnapshot/snapshot_stats.h b/fs_mgr/libsnapshot/snapshot_stats.h
deleted file mode 100644
index 60109a4..0000000
--- a/fs_mgr/libsnapshot/snapshot_stats.h
+++ /dev/null
@@ -1,45 +0,0 @@
-// Copyright (C) 2020 The Android Open Source Project
-//
-// Licensed under the Apache License, Version 2.0 (the "License");
-// you may not use this file except in compliance with the License.
-// You may obtain a copy of the License at
-//
-//      http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing, software
-// distributed under the License is distributed on an "AS IS" BASIS,
-// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-// See the License for the specific language governing permissions and
-// limitations under the License.
-
-#pragma once
-
-#include <chrono>
-
-#include <android/snapshot/snapshot.pb.h>
-#include <libsnapshot/snapshot.h>
-
-namespace android {
-namespace snapshot {
-
-class SnapshotMergeStats {
-  public:
-    SnapshotMergeStats(SnapshotManager& parent);
-    ~SnapshotMergeStats();
-    void Start();
-    void Resume();
-    void set_state(android::snapshot::UpdateState state);
-    SnapshotMergeReport GetReport();
-
-  private:
-    bool ReadState();
-    bool WriteState();
-
-    const SnapshotManager& parent_;
-    SnapshotMergeReport report_;
-    std::chrono::time_point<std::chrono::steady_clock> init_time_;
-    std::chrono::time_point<std::chrono::steady_clock> end_time_;
-};
-
-}  // namespace snapshot
-}  // namespace android
diff --git a/fs_mgr/libsnapshot/snapshot_test.cpp b/fs_mgr/libsnapshot/snapshot_test.cpp
index 5d2840f..7d16ec2 100644
--- a/fs_mgr/libsnapshot/snapshot_test.cpp
+++ b/fs_mgr/libsnapshot/snapshot_test.cpp
@@ -506,9 +506,6 @@
 }
 
 TEST_F(SnapshotTest, FirstStageMountAndMerge) {
-#ifdef SKIP_TEST_IN_PRESUBMIT
-    GTEST_SKIP() << "WIP failure b/148889015";
-#endif
     ASSERT_TRUE(AcquireLock());
 
     static const uint64_t kDeviceSize = 1024 * 1024;
@@ -565,9 +562,6 @@
 }
 
 TEST_F(SnapshotTest, FlashSuperDuringMerge) {
-#ifdef SKIP_TEST_IN_PRESUBMIT
-    GTEST_SKIP() << "WIP failure b/148889015";
-#endif
     ASSERT_TRUE(AcquireLock());
 
     static const uint64_t kDeviceSize = 1024 * 1024;
@@ -979,9 +973,6 @@
 // Also test UnmapUpdateSnapshot unmaps everything.
 // Also test first stage mount and merge after this.
 TEST_F(SnapshotUpdateTest, FullUpdateFlow) {
-#ifdef SKIP_TEST_IN_PRESUBMIT
-    GTEST_SKIP() << "WIP failure b/148889015";
-#endif
     // 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));
@@ -1129,9 +1120,6 @@
 
 // Test that the old partitions are not modified.
 TEST_F(SnapshotUpdateTest, TestRollback) {
-#ifdef SKIP_TEST_IN_PRESUBMIT
-    GTEST_SKIP() << "WIP failure b/148889015";
-#endif
     // Execute the update.
     ASSERT_TRUE(sm->BeginUpdate());
     ASSERT_TRUE(sm->UnmapUpdateSnapshot("sys_b"));
@@ -1309,9 +1297,6 @@
 }
 
 TEST_F(SnapshotUpdateTest, MergeCannotRemoveCow) {
-#ifdef SKIP_TEST_IN_PRESUBMIT
-    GTEST_SKIP() << "WIP failure b/148889015";
-#endif
     // Make source partitions as big as possible to force COW image to be created.
     SetSize(sys_, 5_MiB);
     SetSize(vnd_, 5_MiB);
@@ -1585,48 +1570,6 @@
             << "FinishedSnapshotWrites should detect overflow of CoW device.";
 }
 
-TEST_F(SnapshotUpdateTest, WaitForMerge) {
-#ifdef SKIP_TEST_IN_PRESUBMIT
-    GTEST_SKIP() << "WIP failure b/148889015";
-#endif
-    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", snapshot_timeout_));
-    }
-
-    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));
-}
-
 TEST_F(SnapshotUpdateTest, LowSpace) {
     static constexpr auto kMaxFree = 10_MiB;
     auto userdata = std::make_unique<LowSpaceUserdata>();
diff --git a/fs_mgr/libsnapshot/snapshotctl.cpp b/fs_mgr/libsnapshot/snapshotctl.cpp
index 34d3d69..aa5e9c1 100644
--- a/fs_mgr/libsnapshot/snapshotctl.cpp
+++ b/fs_mgr/libsnapshot/snapshotctl.cpp
@@ -26,9 +26,9 @@
 #include <android-base/unique_fd.h>
 #include <android/snapshot/snapshot.pb.h>
 #include <libsnapshot/snapshot.h>
+#include <libsnapshot/snapshot_stats.h>
 #include <statslog.h>
 
-#include "snapshot_stats.h"
 #include "utility.h"
 
 using namespace std::string_literals;
@@ -178,6 +178,7 @@
     }
 
     LOG(ERROR) << "Snapshot failed to merge with state \"" << state << "\".";
+
     return false;
 }
 
diff --git a/fs_mgr/libsnapshot/snapshotctl.rc b/fs_mgr/libsnapshot/snapshotctl.rc
deleted file mode 100644
index 5dbe352..0000000
--- a/fs_mgr/libsnapshot/snapshotctl.rc
+++ /dev/null
@@ -1,2 +0,0 @@
-on property:sys.boot_completed=1
-    exec_background - root root -- /system/bin/snapshotctl merge --logcat --log-to-file