diff --git a/fastboot/device/commands.cpp b/fastboot/device/commands.cpp
index 409ef70..4c77c75 100644
--- a/fastboot/device/commands.cpp
+++ b/fastboot/device/commands.cpp
@@ -194,23 +194,6 @@
     return device->WriteStatus(FastbootResult::FAIL, "Couldn't download data");
 }
 
-bool FlashHandler(FastbootDevice* device, const std::vector<std::string>& args) {
-    if (args.size() < 2) {
-        return device->WriteStatus(FastbootResult::FAIL, "Invalid arguments");
-    }
-
-    if (GetDeviceLockStatus()) {
-        return device->WriteStatus(FastbootResult::FAIL,
-                                   "Flashing is not allowed on locked devices");
-    }
-
-    int ret = Flash(device, args[1]);
-    if (ret < 0) {
-        return device->WriteStatus(FastbootResult::FAIL, strerror(-ret));
-    }
-    return device->WriteStatus(FastbootResult::OKAY, "Flashing succeeded");
-}
-
 bool SetActiveHandler(FastbootDevice* device, const std::vector<std::string>& args) {
     if (args.size() < 2) {
         return device->WriteStatus(FastbootResult::FAIL, "Missing slot argument");
@@ -440,6 +423,11 @@
     if (!partition) {
         return device->WriteFail("Partition does not exist");
     }
+
+    // Remove the updated flag to cancel any snapshots.
+    uint32_t attrs = partition->attributes();
+    partition->set_attributes(attrs & ~LP_PARTITION_ATTR_UPDATED);
+
     if (!builder->ResizePartition(partition, partition_size)) {
         return device->WriteFail("Not enough space to resize partition");
     }
@@ -449,6 +437,42 @@
     return device->WriteOkay("Partition resized");
 }
 
+void CancelPartitionSnapshot(FastbootDevice* device, const std::string& partition_name) {
+    PartitionBuilder builder(device, partition_name);
+    if (!builder.Valid()) return;
+
+    auto partition = builder->FindPartition(partition_name);
+    if (!partition) return;
+
+    // Remove the updated flag to cancel any snapshots.
+    uint32_t attrs = partition->attributes();
+    partition->set_attributes(attrs & ~LP_PARTITION_ATTR_UPDATED);
+
+    builder.Write();
+}
+
+bool FlashHandler(FastbootDevice* device, const std::vector<std::string>& args) {
+    if (args.size() < 2) {
+        return device->WriteStatus(FastbootResult::FAIL, "Invalid arguments");
+    }
+
+    if (GetDeviceLockStatus()) {
+        return device->WriteStatus(FastbootResult::FAIL,
+                                   "Flashing is not allowed on locked devices");
+    }
+
+    const auto& partition_name = args[1];
+    if (LogicalPartitionExists(device, partition_name)) {
+        CancelPartitionSnapshot(device, partition_name);
+    }
+
+    int ret = Flash(device, partition_name);
+    if (ret < 0) {
+        return device->WriteStatus(FastbootResult::FAIL, strerror(-ret));
+    }
+    return device->WriteStatus(FastbootResult::OKAY, "Flashing succeeded");
+}
+
 bool UpdateSuperHandler(FastbootDevice* device, const std::vector<std::string>& args) {
     if (args.size() < 2) {
         return device->WriteFail("Invalid arguments");
diff --git a/fs_mgr/fs_mgr_dm_linear.cpp b/fs_mgr/fs_mgr_dm_linear.cpp
index fa756a0..1166062 100644
--- a/fs_mgr/fs_mgr_dm_linear.cpp
+++ b/fs_mgr/fs_mgr_dm_linear.cpp
@@ -79,9 +79,9 @@
     return true;
 }
 
-static bool CreateDmTable(const IPartitionOpener& opener, const LpMetadata& metadata,
-                          const LpMetadataPartition& partition, const std::string& super_device,
-                          DmTable* table) {
+bool CreateDmTable(const IPartitionOpener& opener, const LpMetadata& metadata,
+                   const LpMetadataPartition& partition, const std::string& super_device,
+                   DmTable* table) {
     uint64_t sector = 0;
     for (size_t i = 0; i < partition.num_extents; i++) {
         const auto& extent = metadata.extents[partition.first_extent_index + i];
diff --git a/fs_mgr/include/fs_mgr_dm_linear.h b/fs_mgr/include/fs_mgr_dm_linear.h
index 2054fa1..2c54755 100644
--- a/fs_mgr/include/fs_mgr_dm_linear.h
+++ b/fs_mgr/include/fs_mgr_dm_linear.h
@@ -89,6 +89,11 @@
 // is non-zero, then this will block until the device path has been unlinked.
 bool DestroyLogicalPartition(const std::string& name);
 
+// Helper for populating a DmTable for a logical partition.
+bool CreateDmTable(const IPartitionOpener& opener, const LpMetadata& metadata,
+                   const LpMetadataPartition& partition, const std::string& super_device,
+                   android::dm::DmTable* table);
+
 }  // namespace fs_mgr
 }  // namespace android
 
diff --git a/fs_mgr/liblp/include/liblp/builder.h b/fs_mgr/liblp/include/liblp/builder.h
index c3f6e91..5f854c5 100644
--- a/fs_mgr/liblp/include/liblp/builder.h
+++ b/fs_mgr/liblp/include/liblp/builder.h
@@ -123,6 +123,7 @@
     const std::string& name() const { return name_; }
     const std::string& group_name() const { return group_name_; }
     uint32_t attributes() const { return attributes_; }
+    void set_attributes(uint32_t attributes) { attributes_ = attributes; }
     const std::vector<std::unique_ptr<Extent>>& extents() const { return extents_; }
     uint64_t size() const { return size_; }
 
diff --git a/fs_mgr/libsnapshot/include/libsnapshot/snapshot.h b/fs_mgr/libsnapshot/include/libsnapshot/snapshot.h
index b982b4b..e3fc4f6 100644
--- a/fs_mgr/libsnapshot/include/libsnapshot/snapshot.h
+++ b/fs_mgr/libsnapshot/include/libsnapshot/snapshot.h
@@ -66,7 +66,11 @@
     MergeCompleted,
 
     // Merging failed due to an unrecoverable error.
-    MergeFailed
+    MergeFailed,
+
+    // The update was implicitly cancelled, either by a rollback or a flash
+    // operation via fastboot. This state can only be returned by WaitForMerge.
+    Cancelled
 };
 
 class SnapshotManager final {
@@ -82,6 +86,7 @@
         virtual std::string GetGsidDir() const = 0;
         virtual std::string GetMetadataDir() const = 0;
         virtual std::string GetSlotSuffix() const = 0;
+        virtual std::string GetSuperDevice(uint32_t slot) const = 0;
         virtual const IPartitionOpener& GetPartitionOpener() const = 0;
     };
 
@@ -117,12 +122,15 @@
     // update has been marked successful after booting.
     bool InitiateMerge();
 
-    // Wait for the current merge to finish, then perform cleanup when it
-    // completes. It is necessary to call this after InitiateMerge(), or when
-    // a merge state is detected during boot.
+    // Perform any necessary post-boot actions. This should be run soon after
+    // /data is mounted.
     //
-    // Note that after calling WaitForMerge(), GetUpdateState() may still return
-    // that a merge is in progress:
+    // 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
+    // update artifacts and return.
+    //
+    // Note that after calling this, GetUpdateState() may still return that a
+    // merge is in progress:
     //   MergeFailed indicates that a fatal error occurred. WaitForMerge() may
     //   called any number of times again to attempt to make more progress, but
     //   we do not expect it to succeed if a catastrophic error occurred.
@@ -135,7 +143,7 @@
     //
     //   MergeCompleted indicates that the update has fully completed.
     //   GetUpdateState will return None, and a new update can begin.
-    UpdateState WaitForMerge();
+    UpdateState ProcessUpdateState();
 
     // Find the status of the current update, if any.
     //
@@ -158,6 +166,7 @@
     FRIEND_TEST(SnapshotTest, CreateSnapshot);
     FRIEND_TEST(SnapshotTest, FirstStageMountAfterRollback);
     FRIEND_TEST(SnapshotTest, FirstStageMountAndMerge);
+    FRIEND_TEST(SnapshotTest, FlashSuperDuringMerge);
     FRIEND_TEST(SnapshotTest, FlashSuperDuringUpdate);
     FRIEND_TEST(SnapshotTest, MapPartialSnapshot);
     FRIEND_TEST(SnapshotTest, MapSnapshot);
@@ -245,6 +254,14 @@
     // List the known snapshot names.
     bool ListSnapshots(LockedFile* lock, std::vector<std::string>* snapshots);
 
+    // Check for a cancelled or rolled back merge, returning true if such a
+    // condition was detected and handled.
+    bool HandleCancelledUpdate(LockedFile* lock);
+
+    // Remove artifacts created by the update process, such as snapshots, and
+    // set the update state to None.
+    bool RemoveAllUpdateState(LockedFile* lock);
+
     // Interact with /metadata/ota/state.
     std::unique_ptr<LockedFile> OpenStateFile(int open_flags, int lock_flags);
     std::unique_ptr<LockedFile> LockShared();
@@ -272,6 +289,7 @@
     bool MarkSnapshotMergeCompleted(LockedFile* snapshot_lock, const std::string& snapshot_name);
     void AcknowledgeMergeSuccess(LockedFile* lock);
     void AcknowledgeMergeFailure();
+    bool IsCancelledSnapshot(const std::string& snapshot_name);
 
     // Note that these require the name of the device containing the snapshot,
     // which may be the "inner" device. Use GetsnapshotDeviecName().
diff --git a/fs_mgr/libsnapshot/snapshot.cpp b/fs_mgr/libsnapshot/snapshot.cpp
index ab1157b..71457ee 100644
--- a/fs_mgr/libsnapshot/snapshot.cpp
+++ b/fs_mgr/libsnapshot/snapshot.cpp
@@ -28,6 +28,7 @@
 #include <android-base/strings.h>
 #include <android-base/unique_fd.h>
 #include <ext4_utils/ext4_utils.h>
+#include <fs_mgr.h>
 #include <fs_mgr_dm_linear.h>
 #include <fstab/fstab.h>
 #include <libdm/dm.h>
@@ -46,6 +47,7 @@
 using android::dm::kSectorSize;
 using android::dm::SnapshotStorageMode;
 using android::fiemap::IImageManager;
+using android::fs_mgr::CreateDmTable;
 using android::fs_mgr::CreateLogicalPartition;
 using android::fs_mgr::CreateLogicalPartitionParams;
 using android::fs_mgr::GetPartitionName;
@@ -64,6 +66,9 @@
     std::string GetMetadataDir() const override { return "/metadata/ota"s; }
     std::string GetSlotSuffix() const override { return fs_mgr_get_slot_suffix(); }
     const android::fs_mgr::IPartitionOpener& GetPartitionOpener() const { return opener_; }
+    std::string GetSuperDevice(uint32_t slot) const override {
+        return fs_mgr_get_super_partition_name(slot);
+    }
 
   private:
     android::fs_mgr::PartitionOpener opener_;
@@ -123,17 +128,20 @@
         LOG(ERROR) << "Cannot cancel update after it has completed or started merging";
         return false;
     }
+    return RemoveAllUpdateState(file.get());
+}
 
-    if (!RemoveAllSnapshots(file.get())) {
+bool SnapshotManager::RemoveAllUpdateState(LockedFile* lock) {
+    if (!RemoveAllSnapshots(lock)) {
         LOG(ERROR) << "Could not remove all snapshots";
         return false;
     }
 
-    if (!WriteUpdateState(file.get(), UpdateState::None)) {
-        LOG(ERROR) << "Could not write new update state";
-        return false;
-    }
-    return true;
+    RemoveSnapshotBootIndicator();
+
+    // If this fails, we'll keep trying to remove the update state (as the
+    // device reboots or starts a new update) until it finally succeeds.
+    return WriteUpdateState(lock, UpdateState::None);
 }
 
 bool SnapshotManager::FinishedSnapshotWrites() {
@@ -362,14 +370,13 @@
     if (!EnsureImageManager()) return false;
 
     auto cow_name = GetCowName(name);
-    if (!images_->BackingImageExists(cow_name)) {
-        return true;
-    }
-    if (images_->IsImageMapped(cow_name) && !images_->UnmapImageDevice(cow_name)) {
-        return false;
-    }
-    if (!images_->DeleteBackingImage(cow_name)) {
-        return false;
+    if (images_->BackingImageExists(cow_name)) {
+        if (images_->IsImageMapped(cow_name) && !images_->UnmapImageDevice(cow_name)) {
+            return false;
+        }
+        if (!images_->DeleteBackingImage(cow_name)) {
+            return false;
+        }
     }
 
     std::string error;
@@ -575,9 +582,12 @@
 // 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::WaitForMerge() {
+UpdateState SnapshotManager::ProcessUpdateState() {
     while (true) {
         UpdateState state = CheckMergeState();
+        if (state == UpdateState::MergeFailed) {
+            AcknowledgeMergeFailure();
+        }
         if (state != UpdateState::Merging) {
             // Either there is no merge, or the merge was finished, so no need
             // to keep waiting.
@@ -593,15 +603,16 @@
 UpdateState SnapshotManager::CheckMergeState() {
     auto lock = LockExclusive();
     if (!lock) {
-        AcknowledgeMergeFailure();
         return UpdateState::MergeFailed;
     }
 
-    auto state = CheckMergeState(lock.get());
+    UpdateState state = CheckMergeState(lock.get());
     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::MergeFailed) {
-        AcknowledgeMergeFailure();
+    } else if (state == UpdateState::Cancelled) {
+        RemoveAllUpdateState(lock.get());
     }
     return state;
 }
@@ -623,10 +634,17 @@
             // run.
             break;
 
+        case UpdateState::Unverified:
+            // 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)) {
+                return UpdateState::Cancelled;
+            }
+            return state;
+
         default:
-            LOG(ERROR) << "No merge exists, cannot wait. Update state: "
-                       << static_cast<uint32_t>(state);
-            return UpdateState::None;
+            return state;
     }
 
     std::vector<std::string> snapshots;
@@ -634,6 +652,7 @@
         return UpdateState::MergeFailed;
     }
 
+    bool cancelled = false;
     bool failed = false;
     bool merging = false;
     bool needs_reboot = false;
@@ -651,6 +670,9 @@
                 break;
             case UpdateState::MergeCompleted:
                 break;
+            case UpdateState::Cancelled:
+                cancelled = true;
+                break;
             default:
                 LOG(ERROR) << "Unknown merge status: " << static_cast<uint32_t>(snapshot_state);
                 failed = true;
@@ -673,6 +695,14 @@
         WriteUpdateState(lock, UpdateState::MergeNeedsReboot);
         return UpdateState::MergeNeedsReboot;
     }
+    if (cancelled) {
+        // This is an edge case, that we handle as correctly as we sensibly can.
+        // The underlying partition has changed behind update_engine, and we've
+        // removed the snapshot as a result. The exact state of the update is
+        // undefined now, but this can only happen on an unlocked device where
+        // partitions can be flashed without wiping userdata.
+        return UpdateState::Cancelled;
+    }
     return UpdateState::MergeCompleted;
 }
 
@@ -684,17 +714,30 @@
 
     std::string dm_name = GetSnapshotDeviceName(name, snapshot_status);
 
-    // During a check, we decided the merge was complete, but we were unable to
-    // collapse the device-mapper stack and perform COW cleanup. If we haven't
-    // rebooted after this check, the device will still be a snapshot-merge
-    // target. If the have rebooted, the device will now be a linear target,
-    // and we can try cleanup again.
-    if (snapshot_status.state == SnapshotState::MergeCompleted && !IsSnapshotDevice(dm_name)) {
-        // NB: It's okay if this fails now, we gave cleanup our best effort.
-        OnSnapshotMergeComplete(lock, name, snapshot_status);
-        return UpdateState::MergeCompleted;
+    if (!IsSnapshotDevice(dm_name)) {
+        if (IsCancelledSnapshot(name)) {
+            DeleteSnapshot(lock, name);
+            return UpdateState::Cancelled;
+        }
+
+        // During a check, we decided the merge was complete, but we were unable to
+        // collapse the device-mapper stack and perform COW cleanup. If we haven't
+        // rebooted after this check, the device will still be a snapshot-merge
+        // target. If the have rebooted, the device will now be a linear target,
+        // and we can try cleanup again.
+        if (snapshot_status.state == SnapshotState::MergeCompleted) {
+            // NB: It's okay if this fails now, we gave cleanup our best effort.
+            OnSnapshotMergeComplete(lock, name, snapshot_status);
+            return UpdateState::MergeCompleted;
+        }
+
+        LOG(ERROR) << "Expected snapshot or snapshot-merge for device: " << dm_name;
+        return UpdateState::MergeFailed;
     }
 
+    // This check is expensive so it is only enabled for debugging.
+    DCHECK(!IsCancelledSnapshot(name));
+
     std::string target_type;
     DmTargetSnapshot::Status status;
     if (!QuerySnapshotStatus(dm_name, &target_type, &status)) {
@@ -750,12 +793,7 @@
 }
 
 void SnapshotManager::AcknowledgeMergeSuccess(LockedFile* lock) {
-    RemoveSnapshotBootIndicator();
-
-    if (!WriteUpdateState(lock, UpdateState::None)) {
-        // We'll try again next reboot, ad infinitum.
-        return;
-    }
+    RemoveAllUpdateState(lock);
 }
 
 void SnapshotManager::AcknowledgeMergeFailure() {
@@ -814,25 +852,16 @@
 
 bool SnapshotManager::CollapseSnapshotDevice(const std::string& name,
                                              const SnapshotStatus& status) {
-    // Ideally, we would complete the following steps to collapse the device:
-    //  (1) Rewrite the snapshot table to be identical to the base device table.
-    //  (2) Rewrite the verity table to use the "snapshot" (now linear) device.
-    //  (3) Delete the base device.
-    //
-    // This should be possible once libsnapshot understands LpMetadata. In the
-    // meantime, we implement a simpler solution: rewriting the snapshot table
-    // to be a single dm-linear segment against the base device. While not as
-    // ideal, it still lets us remove the COW device. We can remove this
-    // implementation once the new method has been tested.
     auto& dm = DeviceMapper::Instance();
     auto dm_name = GetSnapshotDeviceName(name, status);
 
+    // Verify we have a snapshot-merge device.
     DeviceMapper::TargetInfo target;
     if (!GetSingleTarget(dm_name, TableQuery::Table, &target)) {
         return false;
     }
     if (DeviceMapper::GetTargetType(target.spec) != "snapshot-merge") {
-        // This should be impossible, it was checked above.
+        // This should be impossible, it was checked earlier.
         LOG(ERROR) << "Snapshot device has invalid target type: " << dm_name;
         return false;
     }
@@ -861,7 +890,7 @@
             return false;
         }
         if (outer_table.size() != 2) {
-            LOG(ERROR) << "Expected 2 dm-linear targets for tabble " << name
+            LOG(ERROR) << "Expected 2 dm-linear targets for table " << name
                        << ", got: " << outer_table.size();
             return false;
         }
@@ -881,31 +910,91 @@
         }
     }
 
-    // Note: we are replacing the OUTER table here, so we do not use dm_name.
-    DmTargetLinear new_target(0, num_sectors, base_device, 0);
-    LOG(INFO) << "Replacing snapshot device " << name
-              << " table with: " << new_target.GetParameterString();
+    // Grab the partition metadata for the snapshot.
+    uint32_t slot = SlotNumberForSlotSuffix(device_->GetSlotSuffix());
+    auto super_device = device_->GetSuperDevice(slot);
+    const auto& opener = device_->GetPartitionOpener();
+    auto metadata = android::fs_mgr::ReadMetadata(opener, super_device, slot);
+    if (!metadata) {
+        LOG(ERROR) << "Could not read super partition metadata.";
+        return false;
+    }
+    auto partition = android::fs_mgr::FindPartition(*metadata.get(), name);
+    if (!partition) {
+        LOG(ERROR) << "Snapshot does not have a partition in super: " << name;
+        return false;
+    }
 
+    // Create a DmTable that is identical to the base device.
     DmTable table;
-    table.Emplace<DmTargetLinear>(new_target);
+    if (!CreateDmTable(opener, *metadata.get(), *partition, super_device, &table)) {
+        LOG(ERROR) << "Could not create a DmTable for partition: " << name;
+        return false;
+    }
+
+    // Note: we are replacing the *outer* table here, so we do not use dm_name.
     if (!dm.LoadTableAndActivate(name, table)) {
         return false;
     }
 
-    if (dm_name != name) {
-        // Attempt to delete the snapshot device. Nothing should be depending on
-        // the device, and device-mapper should have flushed remaining I/O. We
-        // could in theory replace with dm-zero (or re-use the table above), but
-        // for now it's better to know why this would fail.
-        if (!dm.DeleteDeviceIfExists(dm_name)) {
-            LOG(ERROR) << "Unable to delete snapshot device " << dm_name << ", COW cannot be "
-                       << "reclaimed until after reboot.";
-            return false;
-        }
+    // Attempt to delete the snapshot device if one still exists. Nothing
+    // should be depending on the device, and device-mapper should have
+    // flushed remaining I/O. We could in theory replace with dm-zero (or
+    // re-use the table above), but for now it's better to know why this
+    // would fail.
+    if (!dm.DeleteDeviceIfExists(dm_name)) {
+        LOG(ERROR) << "Unable to delete snapshot device " << dm_name << ", COW cannot be "
+                   << "reclaimed until after reboot.";
+        return false;
+    }
+
+    // Cleanup the base device as well, since it is no longer used. This does
+    // not block cleanup.
+    auto base_name = GetBaseDeviceName(name);
+    if (!dm.DeleteDeviceIfExists(base_name)) {
+        LOG(ERROR) << "Unable to delete base device for snapshot: " << base_name;
     }
     return true;
 }
 
+bool SnapshotManager::HandleCancelledUpdate(LockedFile* lock) {
+    std::string old_slot;
+    auto boot_file = GetSnapshotBootIndicatorPath();
+    if (!android::base::ReadFileToString(boot_file, &old_slot)) {
+        PLOG(ERROR) << "Unable to read the snapshot indicator file: " << boot_file;
+        return false;
+    }
+    if (device_->GetSlotSuffix() != old_slot) {
+        // We're booted into the target slot, which means we just rebooted
+        // after applying the update.
+        return false;
+    }
+
+    // The only way we can get here is if:
+    //  (1) The device rolled back to the previous slot.
+    //  (2) This function was called prematurely before rebooting the device.
+    //  (3) fastboot set_active was used.
+    //
+    // In any case, delete the snapshots. It may be worth using the boot_control
+    // HAL to differentiate case (2).
+    RemoveAllUpdateState(lock);
+    return true;
+}
+
+bool SnapshotManager::IsCancelledSnapshot(const std::string& snapshot_name) {
+    const auto& opener = device_->GetPartitionOpener();
+    uint32_t slot = SlotNumberForSlotSuffix(device_->GetSlotSuffix());
+    auto super_device = device_->GetSuperDevice(slot);
+    auto metadata = android::fs_mgr::ReadMetadata(opener, super_device, slot);
+    if (!metadata) {
+        LOG(ERROR) << "Could not read dynamic partition metadata for device: " << super_device;
+        return false;
+    }
+    auto partition = android::fs_mgr::FindPartition(*metadata.get(), snapshot_name);
+    if (!partition) return false;
+    return (partition->attributes & LP_PARTITION_ATTR_UPDATED) == 0;
+}
+
 bool SnapshotManager::RemoveAllSnapshots(LockedFile* lock) {
     std::vector<std::string> snapshots;
     if (!ListSnapshots(lock, &snapshots)) {
diff --git a/fs_mgr/libsnapshot/snapshot_test.cpp b/fs_mgr/libsnapshot/snapshot_test.cpp
index acffe8c..8487339 100644
--- a/fs_mgr/libsnapshot/snapshot_test.cpp
+++ b/fs_mgr/libsnapshot/snapshot_test.cpp
@@ -45,6 +45,7 @@
 using android::fs_mgr::BlockDeviceInfo;
 using android::fs_mgr::CreateLogicalPartitionParams;
 using android::fs_mgr::DestroyLogicalPartition;
+using android::fs_mgr::GetPartitionName;
 using android::fs_mgr::MetadataBuilder;
 using namespace ::testing;
 using namespace android::fs_mgr::testing;
@@ -198,6 +199,7 @@
                     .block_device = fake_super,
                     .metadata = metadata.get(),
                     .partition = &partition,
+                    .device_name = GetPartitionName(partition) + "-base",
                     .force_writable = true,
                     .timeout_ms = 10s,
             };
@@ -308,15 +310,20 @@
 }
 
 TEST_F(SnapshotTest, Merge) {
+    ON_CALL(*GetMockedPropertyFetcher(), GetBoolProperty("ro.virtual_ab.enabled", _))
+            .WillByDefault(Return(true));
+
     ASSERT_TRUE(AcquireLock());
 
     static const uint64_t kDeviceSize = 1024 * 1024;
-    ASSERT_TRUE(sm->CreateSnapshot(lock_.get(), "test-snapshot", kDeviceSize, kDeviceSize,
-                                   kDeviceSize));
 
     std::string base_device, snap_device;
-    ASSERT_TRUE(CreatePartition("base-device", kDeviceSize, &base_device));
-    ASSERT_TRUE(sm->MapSnapshot(lock_.get(), "test-snapshot", base_device, 10s, &snap_device));
+    ASSERT_TRUE(CreatePartition("test_partition_a", kDeviceSize));
+    ASSERT_TRUE(MapUpdatePartitions());
+    ASSERT_TRUE(dm_.GetDmDevicePathByName("test_partition_b-base", &base_device));
+    ASSERT_TRUE(sm->CreateSnapshot(lock_.get(), "test_partition_b", kDeviceSize, kDeviceSize,
+                                   kDeviceSize));
+    ASSERT_TRUE(sm->MapSnapshot(lock_.get(), "test_partition_b", base_device, 10s, &snap_device));
 
     std::string test_string = "This is a test string.";
     {
@@ -325,10 +332,10 @@
         ASSERT_TRUE(android::base::WriteFully(fd, test_string.data(), test_string.size()));
     }
 
-    // Note: we know the name of the device is test-snapshot because we didn't
-    // request a linear segment.
+    // Note: we know there is no inner/outer dm device since we didn't request
+    // a linear segment.
     DeviceMapper::TargetInfo target;
-    ASSERT_TRUE(sm->IsSnapshotDevice("test-snapshot", &target));
+    ASSERT_TRUE(sm->IsSnapshotDevice("test_partition_b", &target));
     ASSERT_EQ(DeviceMapper::GetTargetType(target.spec), "snapshot");
 
     // Release the lock.
@@ -341,20 +348,22 @@
     ASSERT_TRUE(sm->InitiateMerge());
 
     // The device should have been switched to a snapshot-merge target.
-    ASSERT_TRUE(sm->IsSnapshotDevice("test-snapshot", &target));
+    ASSERT_TRUE(sm->IsSnapshotDevice("test_partition_b", &target));
     ASSERT_EQ(DeviceMapper::GetTargetType(target.spec), "snapshot-merge");
 
     // We should not be able to cancel an update now.
     ASSERT_FALSE(sm->CancelUpdate());
 
-    ASSERT_EQ(sm->WaitForMerge(), UpdateState::MergeCompleted);
+    ASSERT_EQ(sm->ProcessUpdateState(), UpdateState::MergeCompleted);
     ASSERT_EQ(sm->GetUpdateState(), UpdateState::None);
 
     // The device should no longer be a snapshot or snapshot-merge.
-    ASSERT_FALSE(sm->IsSnapshotDevice("test-snapshot"));
+    ASSERT_FALSE(sm->IsSnapshotDevice("test_partition_b"));
 
-    // Test that we can read back the string we wrote to the snapshot.
-    unique_fd fd(open(base_device.c_str(), O_RDONLY | O_CLOEXEC));
+    // Test that we can read back the string we wrote to the snapshot. Note
+    // that the base device is gone now. |snap_device| contains the correct
+    // partition.
+    unique_fd fd(open(snap_device.c_str(), O_RDONLY | O_CLOEXEC));
     ASSERT_GE(fd, 0);
 
     std::string buffer(test_string.size(), '\0');
@@ -388,7 +397,7 @@
     ASSERT_TRUE(sm->InitiateMerge());
 
     // COW cannot be removed due to open fd, so expect a soft failure.
-    ASSERT_EQ(sm->WaitForMerge(), UpdateState::MergeNeedsReboot);
+    ASSERT_EQ(sm->ProcessUpdateState(), UpdateState::MergeNeedsReboot);
 
     // Forcefully delete the snapshot device, so it looks like we just rebooted.
     DeleteSnapshotDevice("test-snapshot");
@@ -401,7 +410,7 @@
     fd = {};
     lock_ = nullptr;
 
-    ASSERT_EQ(sm->WaitForMerge(), UpdateState::MergeCompleted);
+    ASSERT_EQ(sm->ProcessUpdateState(), UpdateState::MergeCompleted);
 }
 
 TEST_F(SnapshotTest, FirstStageMountAndMerge) {
@@ -420,7 +429,7 @@
     // Simulate a reboot into the new slot.
     lock_ = nullptr;
     ASSERT_TRUE(sm->FinishedSnapshotWrites());
-    ASSERT_TRUE(DestroyLogicalPartition("test_partition_b"));
+    ASSERT_TRUE(DestroyLogicalPartition("test_partition_b-base"));
 
     auto rebooted = new TestDeviceInfo(fake_super);
     rebooted->set_slot_suffix("_b");
@@ -459,7 +468,7 @@
     // Simulate a reboot into the new slot.
     lock_ = nullptr;
     ASSERT_TRUE(sm->FinishedSnapshotWrites());
-    ASSERT_TRUE(DestroyLogicalPartition("test_partition_b"));
+    ASSERT_TRUE(DestroyLogicalPartition("test_partition_b-base"));
 
     // Reflash the super partition.
     FormatFakeSuper();
@@ -482,6 +491,52 @@
     DeviceMapper::TargetInfo target;
     auto dm_name = init->GetSnapshotDeviceName("test_partition_b", status);
     ASSERT_FALSE(init->IsSnapshotDevice(dm_name, &target));
+
+    // We should see a cancelled update as well.
+    lock_ = nullptr;
+    ASSERT_EQ(sm->ProcessUpdateState(), UpdateState::Cancelled);
+}
+
+TEST_F(SnapshotTest, FlashSuperDuringMerge) {
+    ON_CALL(*GetMockedPropertyFetcher(), GetBoolProperty("ro.virtual_ab.enabled", _))
+            .WillByDefault(Return(true));
+
+    ASSERT_TRUE(AcquireLock());
+
+    static const uint64_t kDeviceSize = 1024 * 1024;
+
+    ASSERT_TRUE(CreatePartition("test_partition_a", kDeviceSize));
+    ASSERT_TRUE(MapUpdatePartitions());
+    ASSERT_TRUE(sm->CreateSnapshot(lock_.get(), "test_partition_b", kDeviceSize, kDeviceSize,
+                                   kDeviceSize));
+
+    // Simulate a reboot into the new slot.
+    lock_ = nullptr;
+    ASSERT_TRUE(sm->FinishedSnapshotWrites());
+    ASSERT_TRUE(DestroyLogicalPartition("test_partition_b-base"));
+
+    auto rebooted = new TestDeviceInfo(fake_super);
+    rebooted->set_slot_suffix("_b");
+
+    auto init = SnapshotManager::NewForFirstStageMount(rebooted);
+    ASSERT_NE(init, nullptr);
+    ASSERT_TRUE(init->NeedSnapshotsInFirstStageMount());
+    ASSERT_TRUE(init->CreateLogicalAndSnapshotPartitions("super"));
+    ASSERT_TRUE(init->InitiateMerge());
+
+    // Now, reflash super. Note that we haven't called ProcessUpdateState, so the
+    // status is still Merging.
+    DeleteSnapshotDevice("test_partition_b");
+    ASSERT_TRUE(init->image_manager()->UnmapImageDevice("test_partition_b-cow"));
+    FormatFakeSuper();
+    ASSERT_TRUE(CreatePartition("test_partition_b", kDeviceSize));
+    ASSERT_TRUE(init->NeedSnapshotsInFirstStageMount());
+    ASSERT_TRUE(init->CreateLogicalAndSnapshotPartitions("super"));
+
+    // Because the status is Merging, we must call ProcessUpdateState, which should
+    // detect a cancelled update.
+    ASSERT_EQ(sm->ProcessUpdateState(), UpdateState::Cancelled);
+    ASSERT_EQ(sm->GetUpdateState(), UpdateState::None);
 }
 
 }  // namespace snapshot
diff --git a/fs_mgr/libsnapshot/test_helpers.h b/fs_mgr/libsnapshot/test_helpers.h
index c87f118..9f582d9 100644
--- a/fs_mgr/libsnapshot/test_helpers.h
+++ b/fs_mgr/libsnapshot/test_helpers.h
@@ -47,6 +47,7 @@
     std::string GetGsidDir() const override { return "ota/test"s; }
     std::string GetMetadataDir() const override { return "/metadata/ota/test"s; }
     std::string GetSlotSuffix() const override { return slot_suffix_; }
+    std::string GetSuperDevice([[maybe_unused]] uint32_t slot) const override { return "super"; }
     const android::fs_mgr::IPartitionOpener& GetPartitionOpener() const override {
         return *opener_.get();
     }