libsnapshot: Cancel/merge existing update before begin

Attempt to cancel or merge an existing update before starting
a new one.

Test: libsnapshot_test

Change-Id: Ic534db846f1fe97bd7ca1f176a9494fbc6a0f8d5
diff --git a/fs_mgr/libsnapshot/include/libsnapshot/snapshot.h b/fs_mgr/libsnapshot/include/libsnapshot/snapshot.h
index 8d29e40..6130a10 100644
--- a/fs_mgr/libsnapshot/include/libsnapshot/snapshot.h
+++ b/fs_mgr/libsnapshot/include/libsnapshot/snapshot.h
@@ -127,8 +127,9 @@
     // will fail if GetUpdateState() != None.
     bool BeginUpdate();
 
-    // Cancel an update; any snapshots will be deleted. This will fail if the
-    // state != Initiated or None.
+    // Cancel an update; any snapshots will be deleted. This is allowed if the
+    // state == Initiated, None, or Unverified (before rebooting to the new
+    // slot).
     bool CancelUpdate();
 
     // Mark snapshot writes as having completed. After this, new snapshots cannot
@@ -395,6 +396,13 @@
     // The reverse of MapPartitionWithSnapshot.
     bool UnmapPartitionWithSnapshot(LockedFile* lock, const std::string& target_partition_name);
 
+    // If there isn't a previous update, return true. |needs_merge| is set to false.
+    // If there is a previous update but the device has not boot into it, tries to cancel the
+    //   update and delete any snapshots. Return true if successful. |needs_merge| is set to false.
+    // If there is a previous update and the device has boot into it, do nothing and return true.
+    //   |needs_merge| is set to true.
+    bool TryCancelUpdate(bool* needs_merge);
+
     std::string gsid_dir_;
     std::string metadata_dir_;
     std::unique_ptr<IDeviceInfo> device_;
diff --git a/fs_mgr/libsnapshot/snapshot.cpp b/fs_mgr/libsnapshot/snapshot.cpp
index 0bf0365..803bdb5 100644
--- a/fs_mgr/libsnapshot/snapshot.cpp
+++ b/fs_mgr/libsnapshot/snapshot.cpp
@@ -131,6 +131,16 @@
 }
 
 bool SnapshotManager::BeginUpdate() {
+    bool needs_merge = false;
+    if (!TryCancelUpdate(&needs_merge)) {
+        return false;
+    }
+    if (needs_merge) {
+        LOG(INFO) << "Wait for merge (if any) before beginning a new update.";
+        auto state = ProcessUpdateState();
+        LOG(INFO) << "Merged with state = " << state;
+    }
+
     auto file = LockExclusive();
     if (!file) return false;
 
@@ -143,6 +153,19 @@
 }
 
 bool SnapshotManager::CancelUpdate() {
+    bool needs_merge = false;
+    if (!TryCancelUpdate(&needs_merge)) {
+        return false;
+    }
+    if (needs_merge) {
+        LOG(ERROR) << "Cannot cancel update after it has completed or started merging";
+    }
+    return !needs_merge;
+}
+
+bool SnapshotManager::TryCancelUpdate(bool* needs_merge) {
+    *needs_merge = false;
+
     auto file = LockExclusive();
     if (!file) return false;
 
@@ -167,8 +190,8 @@
             return RemoveAllUpdateState(file.get());
         }
     }
-    LOG(ERROR) << "Cannot cancel update after it has completed or started merging";
-    return false;
+    *needs_merge = true;
+    return true;
 }
 
 bool SnapshotManager::RemoveAllUpdateState(LockedFile* lock) {
diff --git a/fs_mgr/libsnapshot/snapshot_test.cpp b/fs_mgr/libsnapshot/snapshot_test.cpp
index 36982a9..bc764fd 100644
--- a/fs_mgr/libsnapshot/snapshot_test.cpp
+++ b/fs_mgr/libsnapshot/snapshot_test.cpp
@@ -761,8 +761,7 @@
 // Also test UnmapUpdateSnapshot unmaps everything.
 // Also test first stage mount and merge after this.
 TEST_F(SnapshotUpdateTest, FullUpdateFlow) {
-    // OTA client calls CancelUpdate then BeginUpdate before doing anything.
-    ASSERT_TRUE(sm->CancelUpdate());
+    // OTA client calls BeginUpdate before doing anything.
     ASSERT_TRUE(sm->BeginUpdate());
 
     // OTA client blindly unmaps all partitions that are possibly mapped.