1 files changed, 11 insertions, 1 deletions
diff --git a/fs_mgr/libsnapshot/snapshot.cpp b/fs_mgr/libsnapshot/snapshot.cpp
index 4c94da28f..40cb35be1 100644
--- a/fs_mgr/libsnapshot/snapshot.cpp
+++ b/fs_mgr/libsnapshot/snapshot.cpp
@@ -87,6 +87,8 @@ static constexpr char kBootIndicatorPath[] = "/metadata/ota/snapshot-boot";
 static constexpr char kRollbackIndicatorPath[] = "/metadata/ota/rollback-indicator";
 static constexpr auto kUpdateStateCheckInterval = 2s;
 
+MergeFailureCode CheckMergeConsistency(const std::string& name, const SnapshotStatus& status);
+
 // Note: IImageManager is an incomplete type in the header, so the default
 // destructor doesn't work.
 SnapshotManager::~SnapshotManager() {}
@@ -116,6 +118,7 @@ std::unique_ptr<SnapshotManager> SnapshotManager::NewForFirstStageMount(IDeviceI
 
 SnapshotManager::SnapshotManager(IDeviceInfo* device) : device_(device) {
     metadata_dir_ = device_->GetMetadataDir();
+    merge_consistency_checker_ = android::snapshot::CheckMergeConsistency;
 }
 
 static std::string GetCowName(const std::string& snapshot_name) {
@@ -1175,6 +1178,10 @@ MergeFailureCode SnapshotManager::CheckMergeConsistency(LockedFile* lock, const
                                                         const SnapshotStatus& status) {
     CHECK(lock);
 
+    return merge_consistency_checker_(name, status);
+}
+
+MergeFailureCode CheckMergeConsistency(const std::string& name, const SnapshotStatus& status) {
     if (!status.compression_enabled()) {
         // Do not try to verify old-style COWs yet.
         return MergeFailureCode::Ok;
@@ -1252,9 +1259,11 @@ MergeFailureCode SnapshotManager::MergeSecondPhaseSnapshots(LockedFile* lock) {
     }
 
     SnapshotUpdateStatus update_status = ReadSnapshotUpdateStatus(lock);
-    CHECK(update_status.state() == UpdateState::Merging);
+    CHECK(update_status.state() == UpdateState::Merging ||
+          update_status.state() == UpdateState::MergeFailed);
     CHECK(update_status.merge_phase() == MergePhase::FIRST_PHASE);
 
+    update_status.set_state(UpdateState::Merging);
     update_status.set_merge_phase(MergePhase::SECOND_PHASE);
     if (!WriteSnapshotUpdateStatus(lock, update_status)) {
         return MergeFailureCode::WriteStatus;
@@ -2556,6 +2565,7 @@ bool SnapshotManager::WriteUpdateState(LockedFile* lock, UpdateState state,
         SnapshotUpdateStatus old_status = ReadSnapshotUpdateStatus(lock);
         status.set_compression_enabled(old_status.compression_enabled());
         status.set_source_build_fingerprint(old_status.source_build_fingerprint());
+        status.set_merge_phase(old_status.merge_phase());
     }
     return WriteSnapshotUpdateStatus(lock, status);
 }