diff --git a/fs_mgr/libsnapshot/Android.bp b/fs_mgr/libsnapshot/Android.bp
index 5ab2ce285..6b0293ada 100644
--- a/fs_mgr/libsnapshot/Android.bp
+++ b/fs_mgr/libsnapshot/Android.bp
@@ -258,11 +258,62 @@ cc_defaults {
     require_root: true,
 }
 
+cc_defaults {
+    name: "userspace_snapshot_test_defaults",
+    defaults: ["libsnapshot_defaults"],
+    srcs: [
+        "partition_cow_creator_test.cpp",
+        "snapshot_metadata_updater_test.cpp",
+        "snapshot_reader_test.cpp",
+        "userspace_snapshot_test.cpp",
+        "snapshot_writer_test.cpp",
+    ],
+    shared_libs: [
+        "libbinder",
+        "libcrypto",
+        "libhidlbase",
+        "libprotobuf-cpp-lite",
+        "libutils",
+        "libz",
+    ],
+    static_libs: [
+        "android.hardware.boot@1.0",
+        "android.hardware.boot@1.1",
+        "libbrotli",
+        "libc++fs",
+        "libfs_mgr_binder",
+        "libgsi",
+        "libgmock",
+        "liblp",
+        "libsnapshot",
+        "libsnapshot_cow",
+        "libsnapshot_test_helpers",
+        "libsparse",
+    ],
+    header_libs: [
+        "libstorage_literals_headers",
+    ],
+    test_suites: [
+        "vts",
+        "device-tests"
+    ],
+    test_options: {
+        min_shipping_api_level: 29,
+    },
+    auto_gen_config: true,
+    require_root: true,
+}
+
 cc_test {
     name: "vts_libsnapshot_test",
     defaults: ["libsnapshot_test_defaults"],
 }
 
+cc_test {
+    name: "vts_userspace_snapshot_test",
+    defaults: ["userspace_snapshot_test_defaults"],
+}
+
 cc_binary {
     name: "snapshotctl",
     srcs: [
diff --git a/fs_mgr/libsnapshot/android/snapshot/snapshot.proto b/fs_mgr/libsnapshot/android/snapshot/snapshot.proto
index e2abdbaba..532f66dda 100644
--- a/fs_mgr/libsnapshot/android/snapshot/snapshot.proto
+++ b/fs_mgr/libsnapshot/android/snapshot/snapshot.proto
@@ -194,6 +194,9 @@ message SnapshotUpdateStatus {
 
     // Source build fingerprint.
     string source_build_fingerprint = 8;
+
+    // user-space snapshots
+    bool userspace_snapshots = 9;
 }
 
 // Next: 10
diff --git a/fs_mgr/libsnapshot/include/libsnapshot/mock_snapshot.h b/fs_mgr/libsnapshot/include/libsnapshot/mock_snapshot.h
index ec58cca2e..ba62330f1 100644
--- a/fs_mgr/libsnapshot/include/libsnapshot/mock_snapshot.h
+++ b/fs_mgr/libsnapshot/include/libsnapshot/mock_snapshot.h
@@ -35,6 +35,7 @@ class MockSnapshotManager : public ISnapshotManager {
                 (override));
     MOCK_METHOD(UpdateState, GetUpdateState, (double* progress), (override));
     MOCK_METHOD(bool, UpdateUsesCompression, (), (override));
+    MOCK_METHOD(bool, UpdateUsesUserSnapshots, (), (override));
     MOCK_METHOD(Return, CreateUpdateSnapshots,
                 (const chromeos_update_engine::DeltaArchiveManifest& manifest), (override));
     MOCK_METHOD(bool, MapUpdateSnapshot,
diff --git a/fs_mgr/libsnapshot/include/libsnapshot/snapshot.h b/fs_mgr/libsnapshot/include/libsnapshot/snapshot.h
index a49b0261c..08c39205a 100644
--- a/fs_mgr/libsnapshot/include/libsnapshot/snapshot.h
+++ b/fs_mgr/libsnapshot/include/libsnapshot/snapshot.h
@@ -193,6 +193,9 @@ class ISnapshotManager {
     // UpdateState is None, or no snapshots have been created.
     virtual bool UpdateUsesCompression() = 0;
 
+    // Returns true if userspace snapshots is enabled for the current update.
+    virtual bool UpdateUsesUserSnapshots() = 0;
+
     // Create necessary COW device / files for OTA clients. New logical partitions will be added to
     // group "cow" in target_metadata. Regions of partitions of current_metadata will be
     // "write-protected" and snapshotted.
@@ -352,6 +355,7 @@ class SnapshotManager final : public ISnapshotManager {
                                    const std::function<bool()>& before_cancel = {}) override;
     UpdateState GetUpdateState(double* progress = nullptr) override;
     bool UpdateUsesCompression() override;
+    bool UpdateUsesUserSnapshots() override;
     Return CreateUpdateSnapshots(const DeltaArchiveManifest& manifest) override;
     bool MapUpdateSnapshot(const CreateLogicalPartitionParams& params,
                            std::string* snapshot_path) override;
@@ -387,6 +391,11 @@ class SnapshotManager final : public ISnapshotManager {
     // first-stage to decide whether to launch snapuserd.
     bool IsSnapuserdRequired();
 
+    enum class SnapshotDriver {
+        DM_SNAPSHOT,
+        DM_USER,
+    };
+
   private:
     FRIEND_TEST(SnapshotTest, CleanFirstStageMount);
     FRIEND_TEST(SnapshotTest, CreateSnapshot);
@@ -456,6 +465,8 @@ class SnapshotManager final : public ISnapshotManager {
     };
     static std::unique_ptr<LockedFile> OpenFile(const std::string& file, int lock_flags);
 
+    SnapshotDriver GetSnapshotDriver(LockedFile* lock);
+
     // Create a new snapshot record. This creates the backing COW store and
     // persists information needed to map the device. The device can be mapped
     // with MapSnapshot().
@@ -491,8 +502,8 @@ class SnapshotManager final : public ISnapshotManager {
 
     // Create a dm-user device for a given snapshot.
     bool MapDmUserCow(LockedFile* lock, const std::string& name, const std::string& cow_file,
-                      const std::string& base_device, const std::chrono::milliseconds& timeout_ms,
-                      std::string* path);
+                      const std::string& base_device, const std::string& base_path_merge,
+                      const std::chrono::milliseconds& timeout_ms, std::string* path);
 
     // Map the source device used for dm-user.
     bool MapSourceDevice(LockedFile* lock, const std::string& name,
@@ -591,7 +602,8 @@ class SnapshotManager final : public ISnapshotManager {
     // Internal callback for when merging is complete.
     bool OnSnapshotMergeComplete(LockedFile* lock, const std::string& name,
                                  const SnapshotStatus& status);
-    bool CollapseSnapshotDevice(const std::string& name, const SnapshotStatus& status);
+    bool CollapseSnapshotDevice(LockedFile* lock, const std::string& name,
+                                const SnapshotStatus& status);
 
     struct MergeResult {
         explicit MergeResult(UpdateState state,
@@ -689,7 +701,10 @@ class SnapshotManager final : public ISnapshotManager {
     bool UnmapPartitionWithSnapshot(LockedFile* lock, const std::string& target_partition_name);
 
     // Unmap a dm-user device through snapuserd.
-    bool UnmapDmUserDevice(const std::string& snapshot_name);
+    bool UnmapDmUserDevice(const std::string& dm_user_name);
+
+    // Unmap a dm-user device for user space snapshots
+    bool UnmapUserspaceSnapshotDevice(LockedFile* lock, const std::string& snapshot_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
@@ -778,6 +793,8 @@ class SnapshotManager final : public ISnapshotManager {
 
     // Helper of UpdateUsesCompression
     bool UpdateUsesCompression(LockedFile* lock);
+    // Helper of UpdateUsesUsersnapshots
+    bool UpdateUsesUserSnapshots(LockedFile* lock);
 
     // Wrapper around libdm, with diagnostics.
     bool DeleteDeviceIfExists(const std::string& name,
@@ -792,6 +809,7 @@ class SnapshotManager final : public ISnapshotManager {
     std::function<bool(const std::string&)> uevent_regen_callback_;
     std::unique_ptr<SnapuserdClient> snapuserd_client_;
     std::unique_ptr<LpMetadata> old_partition_metadata_;
+    std::optional<bool> is_snapshot_userspace_;
 };
 
 }  // namespace snapshot
diff --git a/fs_mgr/libsnapshot/include/libsnapshot/snapshot_stub.h b/fs_mgr/libsnapshot/include/libsnapshot/snapshot_stub.h
index 74b78c59b..318e5259d 100644
--- a/fs_mgr/libsnapshot/include/libsnapshot/snapshot_stub.h
+++ b/fs_mgr/libsnapshot/include/libsnapshot/snapshot_stub.h
@@ -35,6 +35,7 @@ class SnapshotManagerStub : public ISnapshotManager {
                                    const std::function<bool()>& before_cancel = {}) override;
     UpdateState GetUpdateState(double* progress = nullptr) override;
     bool UpdateUsesCompression() override;
+    bool UpdateUsesUserSnapshots() override;
     Return CreateUpdateSnapshots(
             const chromeos_update_engine::DeltaArchiveManifest& manifest) override;
     bool MapUpdateSnapshot(const android::fs_mgr::CreateLogicalPartitionParams& params,
diff --git a/fs_mgr/libsnapshot/snapshot.cpp b/fs_mgr/libsnapshot/snapshot.cpp
index 3d8ae29df..f4584d282 100644
--- a/fs_mgr/libsnapshot/snapshot.cpp
+++ b/fs_mgr/libsnapshot/snapshot.cpp
@@ -95,6 +95,7 @@ std::unique_ptr<SnapshotManager> SnapshotManager::New(IDeviceInfo* info) {
     if (!info) {
         info = new DeviceInfo();
     }
+
     return std::unique_ptr<SnapshotManager>(new SnapshotManager(info));
 }
 
@@ -121,8 +122,34 @@ static std::string GetCowName(const std::string& snapshot_name) {
     return snapshot_name + "-cow";
 }
 
-static std::string GetDmUserCowName(const std::string& snapshot_name) {
-    return snapshot_name + "-user-cow";
+SnapshotManager::SnapshotDriver SnapshotManager::GetSnapshotDriver(LockedFile* lock) {
+    if (UpdateUsesUserSnapshots(lock)) {
+        return SnapshotManager::SnapshotDriver::DM_USER;
+    } else {
+        return SnapshotManager::SnapshotDriver::DM_SNAPSHOT;
+    }
+}
+
+static std::string GetDmUserCowName(const std::string& snapshot_name,
+                                    SnapshotManager::SnapshotDriver driver) {
+    // dm-user block device will act as a snapshot device. We identify it with
+    // the same partition name so that when partitions can be mounted off
+    // dm-user.
+
+    switch (driver) {
+        case SnapshotManager::SnapshotDriver::DM_USER: {
+            return snapshot_name;
+        }
+
+        case SnapshotManager::SnapshotDriver::DM_SNAPSHOT: {
+            return snapshot_name + "-user-cow";
+        }
+
+        default: {
+            LOG(ERROR) << "Invalid snapshot driver";
+            return "";
+        }
+    }
 }
 
 static std::string GetCowImageDeviceName(const std::string& snapshot_name) {
@@ -398,9 +425,33 @@ Return SnapshotManager::CreateCowImage(LockedFile* lock, const std::string& name
 
 bool SnapshotManager::MapDmUserCow(LockedFile* lock, const std::string& name,
                                    const std::string& cow_file, const std::string& base_device,
+                                   const std::string& base_path_merge,
                                    const std::chrono::milliseconds& timeout_ms, std::string* path) {
     CHECK(lock);
 
+    if (UpdateUsesUserSnapshots(lock)) {
+        SnapshotStatus status;
+        if (!ReadSnapshotStatus(lock, name, &status)) {
+            LOG(ERROR) << "MapDmUserCow: ReadSnapshotStatus failed...";
+            return false;
+        }
+
+        if (status.state() == SnapshotState::NONE ||
+            status.state() == SnapshotState::MERGE_COMPLETED) {
+            LOG(ERROR) << "Should not create a snapshot device for " << name
+                       << " after merging has completed.";
+            return false;
+        }
+
+        SnapshotUpdateStatus update_status = ReadSnapshotUpdateStatus(lock);
+        if (update_status.state() == UpdateState::MergeCompleted ||
+            update_status.state() == UpdateState::MergeNeedsReboot) {
+            LOG(ERROR) << "Should not create a snapshot device for " << name
+                       << " after global merging has completed.";
+            return false;
+        }
+    }
+
     // Use an extra decoration for first-stage init, so we can transition
     // to a new table entry in second-stage.
     std::string misc_name = name;
@@ -412,18 +463,41 @@ bool SnapshotManager::MapDmUserCow(LockedFile* lock, const std::string& name,
         return false;
     }
 
-    uint64_t base_sectors = snapuserd_client_->InitDmUserCow(misc_name, cow_file, base_device);
-    if (base_sectors == 0) {
-        LOG(ERROR) << "Failed to retrieve base_sectors from Snapuserd";
-        return false;
+    uint64_t base_sectors = 0;
+    if (!UpdateUsesUserSnapshots(lock)) {
+        base_sectors = snapuserd_client_->InitDmUserCow(misc_name, cow_file, base_device);
+        if (base_sectors == 0) {
+            LOG(ERROR) << "Failed to retrieve base_sectors from Snapuserd";
+            return false;
+        }
+    } else {
+        // For userspace snapshots, the size of the base device is taken as the
+        // size of the dm-user block device. Since there is no pseudo mapping
+        // created in the daemon, we no longer need to rely on the daemon for
+        // sizing the dm-user block device.
+        unique_fd fd(TEMP_FAILURE_RETRY(open(base_path_merge.c_str(), O_RDONLY | O_CLOEXEC)));
+        if (fd < 0) {
+            LOG(ERROR) << "Cannot open block device: " << base_path_merge;
+            return false;
+        }
+
+        uint64_t dev_sz = get_block_device_size(fd.get());
+        if (!dev_sz) {
+            LOG(ERROR) << "Failed to find block device size: " << base_path_merge;
+            return false;
+        }
+
+        base_sectors = dev_sz >> 9;
     }
 
     DmTable table;
     table.Emplace<DmTargetUser>(0, base_sectors, misc_name);
     if (!dm_.CreateDevice(name, table, path, timeout_ms)) {
+        LOG(ERROR) << " dm-user: CreateDevice failed... ";
         return false;
     }
     if (!WaitForDevice(*path, timeout_ms)) {
+        LOG(ERROR) << " dm-user: timeout: Failed to create block device for: " << name;
         return false;
     }
 
@@ -432,6 +506,15 @@ bool SnapshotManager::MapDmUserCow(LockedFile* lock, const std::string& name,
         return false;
     }
 
+    if (UpdateUsesUserSnapshots(lock)) {
+        // Now that the dm-user device is created, initialize the daemon and
+        // spin up the worker threads.
+        if (!snapuserd_client_->InitDmUserCow(misc_name, cow_file, base_device, base_path_merge)) {
+            LOG(ERROR) << "InitDmUserCow failed";
+            return false;
+        }
+    }
+
     return snapuserd_client_->AttachDmUser(misc_name);
 }
 
@@ -698,13 +781,15 @@ bool SnapshotManager::InitiateMerge() {
 
     DmTargetSnapshot::Status initial_target_values = {};
     for (const auto& snapshot : snapshots) {
-        DmTargetSnapshot::Status current_status;
-        if (!QuerySnapshotStatus(snapshot, nullptr, &current_status)) {
-            return false;
+        if (!UpdateUsesUserSnapshots(lock.get())) {
+            DmTargetSnapshot::Status current_status;
+            if (!QuerySnapshotStatus(snapshot, nullptr, &current_status)) {
+                return false;
+            }
+            initial_target_values.sectors_allocated += current_status.sectors_allocated;
+            initial_target_values.total_sectors += current_status.total_sectors;
+            initial_target_values.metadata_sectors += current_status.metadata_sectors;
         }
-        initial_target_values.sectors_allocated += current_status.sectors_allocated;
-        initial_target_values.total_sectors += current_status.total_sectors;
-        initial_target_values.metadata_sectors += current_status.metadata_sectors;
 
         SnapshotStatus snapshot_status;
         if (!ReadSnapshotStatus(lock.get(), snapshot, &snapshot_status)) {
@@ -719,11 +804,14 @@ bool SnapshotManager::InitiateMerge() {
 
     SnapshotUpdateStatus initial_status = ReadSnapshotUpdateStatus(lock.get());
     initial_status.set_state(UpdateState::Merging);
-    initial_status.set_sectors_allocated(initial_target_values.sectors_allocated);
-    initial_status.set_total_sectors(initial_target_values.total_sectors);
-    initial_status.set_metadata_sectors(initial_target_values.metadata_sectors);
     initial_status.set_compression_enabled(compression_enabled);
 
+    if (!UpdateUsesUserSnapshots(lock.get())) {
+        initial_status.set_sectors_allocated(initial_target_values.sectors_allocated);
+        initial_status.set_total_sectors(initial_target_values.total_sectors);
+        initial_status.set_metadata_sectors(initial_target_values.metadata_sectors);
+    }
+
     // If any partitions shrunk, we need to merge them before we merge any other
     // partitions (see b/177935716). Otherwise, a merge from another partition
     // may overwrite the source block of a copy operation.
@@ -777,20 +865,36 @@ MergeFailureCode SnapshotManager::SwitchSnapshotToMerge(LockedFile* lock, const
                      << " has unexpected state: " << SnapshotState_Name(status.state());
     }
 
-    // After this, we return true because we technically did switch to a merge
-    // target. Everything else we do here is just informational.
-    if (auto code = RewriteSnapshotDeviceTable(name); code != MergeFailureCode::Ok) {
-        return code;
+    if (UpdateUsesUserSnapshots(lock)) {
+        if (EnsureSnapuserdConnected()) {
+            // This is the point where we inform the daemon to initiate/resume
+            // the merge
+            if (!snapuserd_client_->InitiateMerge(name)) {
+                return MergeFailureCode::UnknownTable;
+            }
+        } else {
+            LOG(ERROR) << "Failed to connect to snapuserd daemon to initiate merge";
+            return MergeFailureCode::UnknownTable;
+        }
+    } else {
+        // After this, we return true because we technically did switch to a merge
+        // target. Everything else we do here is just informational.
+        if (auto code = RewriteSnapshotDeviceTable(name); code != MergeFailureCode::Ok) {
+            return code;
+        }
     }
 
     status.set_state(SnapshotState::MERGING);
 
-    DmTargetSnapshot::Status dm_status;
-    if (!QuerySnapshotStatus(name, nullptr, &dm_status)) {
-        LOG(ERROR) << "Could not query merge status for snapshot: " << name;
+    if (!UpdateUsesUserSnapshots(lock)) {
+        DmTargetSnapshot::Status dm_status;
+        if (!QuerySnapshotStatus(name, nullptr, &dm_status)) {
+            LOG(ERROR) << "Could not query merge status for snapshot: " << name;
+        }
+        status.set_sectors_allocated(dm_status.sectors_allocated);
+        status.set_metadata_sectors(dm_status.metadata_sectors);
     }
-    status.set_sectors_allocated(dm_status.sectors_allocated);
-    status.set_metadata_sectors(dm_status.metadata_sectors);
+
     if (!WriteSnapshotStatus(lock, status)) {
         LOG(ERROR) << "Could not update status file for snapshot: " << name;
     }
@@ -856,9 +960,15 @@ bool SnapshotManager::IsSnapshotDevice(const std::string& dm_name, TargetInfo* t
         return false;
     }
     auto type = DeviceMapper::GetTargetType(snap_target.spec);
-    if (type != "snapshot" && type != "snapshot-merge") {
-        return false;
+
+    // If this is not a user-snapshot device then it should either
+    // be a dm-snapshot or dm-snapshot-merge target
+    if (type != "user") {
+        if (type != "snapshot" && type != "snapshot-merge") {
+            return false;
+        }
     }
+
     if (target) {
         *target = std::move(snap_target);
     }
@@ -1094,34 +1204,86 @@ auto SnapshotManager::CheckTargetMergeState(LockedFile* lock, const std::string&
     DCHECK((current_metadata = ReadCurrentMetadata()) &&
            GetMetadataPartitionState(*current_metadata, name) == MetadataPartitionState::Updated);
 
-    std::string target_type;
-    DmTargetSnapshot::Status status;
-    if (!QuerySnapshotStatus(name, &target_type, &status)) {
-        return MergeResult(UpdateState::MergeFailed, MergeFailureCode::QuerySnapshotStatus);
-    }
-    if (target_type == "snapshot" &&
-        DecideMergePhase(snapshot_status) == MergePhase::SECOND_PHASE &&
-        update_status.merge_phase() == MergePhase::FIRST_PHASE) {
-        // The snapshot is not being merged because it's in the wrong phase.
-        return MergeResult(UpdateState::None);
-    }
-    if (target_type != "snapshot-merge") {
-        // We can get here if we failed to rewrite the target type in
-        // InitiateMerge(). If we failed to create the target in first-stage
-        // init, boot would not succeed.
-        LOG(ERROR) << "Snapshot " << name << " has incorrect target type: " << target_type;
-        return MergeResult(UpdateState::MergeFailed, MergeFailureCode::ExpectedMergeTarget);
+    if (UpdateUsesUserSnapshots(lock)) {
+        std::string merge_status;
+        if (EnsureSnapuserdConnected()) {
+            // Query the snapshot status from the daemon
+            merge_status = snapuserd_client_->QuerySnapshotStatus(name);
+        } else {
+            MergeResult(UpdateState::MergeFailed, MergeFailureCode::QuerySnapshotStatus);
+        }
+
+        if (merge_status == "snapshot-merge-failed") {
+            return MergeResult(UpdateState::MergeFailed, MergeFailureCode::UnknownTargetType);
+        }
+
+        // This is the case when device reboots during merge. Once the device boots,
+        // snapuserd daemon will not resume merge immediately in first stage init.
+        // This is slightly different as compared to dm-snapshot-merge; In this
+        // case, metadata file will have "MERGING" state whereas the daemon will be
+        // waiting to resume the merge. Thus, we resume the merge at this point.
+        if (merge_status == "snapshot" && snapshot_status.state() == SnapshotState::MERGING) {
+            if (!snapuserd_client_->InitiateMerge(name)) {
+                return MergeResult(UpdateState::MergeFailed, MergeFailureCode::UnknownTargetType);
+            }
+            return MergeResult(UpdateState::Merging);
+        }
+
+        if (merge_status == "snapshot" &&
+            DecideMergePhase(snapshot_status) == MergePhase::SECOND_PHASE &&
+            update_status.merge_phase() == MergePhase::FIRST_PHASE) {
+            // The snapshot is not being merged because it's in the wrong phase.
+            return MergeResult(UpdateState::None);
+        }
+
+        if (merge_status == "snapshot-merge") {
+            if (snapshot_status.state() == SnapshotState::MERGE_COMPLETED) {
+                LOG(ERROR) << "Snapshot " << name
+                           << " is merging after being marked merge-complete.";
+                return MergeResult(UpdateState::MergeFailed,
+                                   MergeFailureCode::UnmergedSectorsAfterCompletion);
+            }
+            return MergeResult(UpdateState::Merging);
+        }
+
+        if (merge_status != "snapshot-merge-complete") {
+            LOG(ERROR) << "Snapshot " << name << " has incorrect status: " << merge_status;
+            return MergeResult(UpdateState::MergeFailed, MergeFailureCode::ExpectedMergeTarget);
+        }
+    } else {
+        // dm-snapshot in the kernel
+        std::string target_type;
+        DmTargetSnapshot::Status status;
+        if (!QuerySnapshotStatus(name, &target_type, &status)) {
+            return MergeResult(UpdateState::MergeFailed, MergeFailureCode::QuerySnapshotStatus);
+        }
+        if (target_type == "snapshot" &&
+            DecideMergePhase(snapshot_status) == MergePhase::SECOND_PHASE &&
+            update_status.merge_phase() == MergePhase::FIRST_PHASE) {
+            // The snapshot is not being merged because it's in the wrong phase.
+            return MergeResult(UpdateState::None);
+        }
+        if (target_type != "snapshot-merge") {
+            // We can get here if we failed to rewrite the target type in
+            // InitiateMerge(). If we failed to create the target in first-stage
+            // init, boot would not succeed.
+            LOG(ERROR) << "Snapshot " << name << " has incorrect target type: " << target_type;
+            return MergeResult(UpdateState::MergeFailed, MergeFailureCode::ExpectedMergeTarget);
+        }
+
+        // These two values are equal when merging is complete.
+        if (status.sectors_allocated != status.metadata_sectors) {
+            if (snapshot_status.state() == SnapshotState::MERGE_COMPLETED) {
+                LOG(ERROR) << "Snapshot " << name
+                           << " is merging after being marked merge-complete.";
+                return MergeResult(UpdateState::MergeFailed,
+                                   MergeFailureCode::UnmergedSectorsAfterCompletion);
+            }
+            return MergeResult(UpdateState::Merging);
+        }
     }
 
-    // These two values are equal when merging is complete.
-    if (status.sectors_allocated != status.metadata_sectors) {
-        if (snapshot_status.state() == SnapshotState::MERGE_COMPLETED) {
-            LOG(ERROR) << "Snapshot " << name << " is merging after being marked merge-complete.";
-            return MergeResult(UpdateState::MergeFailed,
-                               MergeFailureCode::UnmergedSectorsAfterCompletion);
-        }
-        return MergeResult(UpdateState::Merging);
-    }
+    // Merge is complete at this point
 
     auto code = CheckMergeConsistency(lock, name, snapshot_status);
     if (code != MergeFailureCode::Ok) {
@@ -1311,30 +1473,40 @@ void SnapshotManager::AcknowledgeMergeFailure(MergeFailureCode failure_code) {
 
 bool SnapshotManager::OnSnapshotMergeComplete(LockedFile* lock, const std::string& name,
                                               const SnapshotStatus& status) {
-    if (IsSnapshotDevice(name)) {
-        // We are extra-cautious here, to avoid deleting the wrong table.
-        std::string target_type;
-        DmTargetSnapshot::Status dm_status;
-        if (!QuerySnapshotStatus(name, &target_type, &dm_status)) {
-            return false;
+    if (!UpdateUsesUserSnapshots(lock)) {
+        if (IsSnapshotDevice(name)) {
+            // We are extra-cautious here, to avoid deleting the wrong table.
+            std::string target_type;
+            DmTargetSnapshot::Status dm_status;
+            if (!QuerySnapshotStatus(name, &target_type, &dm_status)) {
+                return false;
+            }
+            if (target_type != "snapshot-merge") {
+                LOG(ERROR) << "Unexpected target type " << target_type
+                           << " for snapshot device: " << name;
+                return false;
+            }
+            if (dm_status.sectors_allocated != dm_status.metadata_sectors) {
+                LOG(ERROR) << "Merge is unexpectedly incomplete for device " << name;
+                return false;
+            }
+            if (!CollapseSnapshotDevice(lock, name, status)) {
+                LOG(ERROR) << "Unable to collapse snapshot: " << name;
+                return false;
+            }
         }
-        if (target_type != "snapshot-merge") {
-            LOG(ERROR) << "Unexpected target type " << target_type
-                       << " for snapshot device: " << name;
-            return false;
-        }
-        if (dm_status.sectors_allocated != dm_status.metadata_sectors) {
-            LOG(ERROR) << "Merge is unexpectedly incomplete for device " << name;
-            return false;
-        }
-        if (!CollapseSnapshotDevice(name, status)) {
+    } else {
+        // Just collapse the device - no need to query again as we just did
+        // prior to calling this function
+        if (!CollapseSnapshotDevice(lock, name, status)) {
             LOG(ERROR) << "Unable to collapse snapshot: " << name;
             return false;
         }
-        // Note that collapsing is implicitly an Unmap, so we don't need to
-        // unmap the snapshot.
     }
 
+    // Note that collapsing is implicitly an Unmap, so we don't need to
+    // unmap the snapshot.
+
     if (!DeleteSnapshot(lock, name)) {
         LOG(ERROR) << "Could not delete snapshot: " << name;
         return false;
@@ -1342,23 +1514,26 @@ bool SnapshotManager::OnSnapshotMergeComplete(LockedFile* lock, const std::strin
     return true;
 }
 
-bool SnapshotManager::CollapseSnapshotDevice(const std::string& name,
+bool SnapshotManager::CollapseSnapshotDevice(LockedFile* lock, const std::string& name,
                                              const SnapshotStatus& status) {
-    // Verify we have a snapshot-merge device.
-    DeviceMapper::TargetInfo target;
-    if (!GetSingleTarget(name, TableQuery::Table, &target)) {
-        return false;
-    }
-    if (DeviceMapper::GetTargetType(target.spec) != "snapshot-merge") {
-        // This should be impossible, it was checked earlier.
-        LOG(ERROR) << "Snapshot device has invalid target type: " << name;
-        return false;
-    }
+    if (!UpdateUsesUserSnapshots(lock)) {
+        // Verify we have a snapshot-merge device.
+        DeviceMapper::TargetInfo target;
+        if (!GetSingleTarget(name, TableQuery::Table, &target)) {
+            return false;
+        }
+        if (DeviceMapper::GetTargetType(target.spec) != "snapshot-merge") {
+            // This should be impossible, it was checked earlier.
+            LOG(ERROR) << "Snapshot device has invalid target type: " << name;
+            return false;
+        }
 
-    std::string base_device, cow_device;
-    if (!DmTargetSnapshot::GetDevicesFromParams(target.data, &base_device, &cow_device)) {
-        LOG(ERROR) << "Could not parse snapshot device " << name << " parameters: " << target.data;
-        return false;
+        std::string base_device, cow_device;
+        if (!DmTargetSnapshot::GetDevicesFromParams(target.data, &base_device, &cow_device)) {
+            LOG(ERROR) << "Could not parse snapshot device " << name
+                       << " parameters: " << target.data;
+            return false;
+        }
     }
 
     uint64_t snapshot_sectors = status.snapshot_size() / kSectorSize;
@@ -1386,14 +1561,32 @@ bool SnapshotManager::CollapseSnapshotDevice(const std::string& name,
         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 (status.compression_enabled()) {
-        UnmapDmUserDevice(name);
+    if (!UpdateUsesUserSnapshots(lock)) {
+        // 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.
+        //
+        // Furthermore, we should not be trying to unmap for userspace snapshot
+        // as unmap will fail since dm-user itself was a snapshot device prior
+        // to switching of tables. Unmap will fail as the device will be mounted
+        // by system partitions
+        if (status.compression_enabled()) {
+            auto dm_user_name = GetDmUserCowName(name, GetSnapshotDriver(lock));
+            UnmapDmUserDevice(dm_user_name);
+        }
     }
+
+    // We can't delete base device immediately as daemon holds a reference.
+    // Make sure we wait for all the worker threads to terminate and release
+    // the reference
+    if (UpdateUsesUserSnapshots(lock) && EnsureSnapuserdConnected()) {
+        if (!snapuserd_client_->WaitForDeviceDelete(name)) {
+            LOG(ERROR) << "Failed to wait for " << name << " control device to delete";
+        }
+    }
+
     auto base_name = GetBaseDeviceName(name);
     if (!DeleteDeviceIfExists(base_name)) {
         LOG(ERROR) << "Unable to delete base device for snapshot: " << base_name;
@@ -1464,10 +1657,15 @@ bool SnapshotManager::PerformInitTransition(InitTransition transition,
         return false;
     }
 
+    if (UpdateUsesUserSnapshots(lock.get()) && transition == InitTransition::SELINUX_DETACH) {
+        snapuserd_argv->emplace_back("-user_snapshot");
+    }
+
     size_t num_cows = 0;
     size_t ok_cows = 0;
     for (const auto& snapshot : snapshots) {
-        std::string user_cow_name = GetDmUserCowName(snapshot);
+        std::string user_cow_name = GetDmUserCowName(snapshot, GetSnapshotDriver(lock.get()));
+
         if (dm_.GetState(user_cow_name) == DmDeviceState::INVALID) {
             continue;
         }
@@ -1513,6 +1711,12 @@ bool SnapshotManager::PerformInitTransition(InitTransition transition,
             continue;
         }
 
+        std::string base_path_merge;
+        if (!dm_.GetDmDevicePathByName(GetBaseDeviceName(snapshot), &base_path_merge)) {
+            LOG(ERROR) << "Could not get device path for " << GetSourceDeviceName(snapshot);
+            continue;
+        }
+
         std::string cow_image_name = GetMappedCowDeviceName(snapshot, snapshot_status);
 
         std::string cow_image_device;
@@ -1529,8 +1733,14 @@ bool SnapshotManager::PerformInitTransition(InitTransition transition,
         }
 
         if (transition == InitTransition::SELINUX_DETACH) {
-            auto message = misc_name + "," + cow_image_device + "," + source_device;
-            snapuserd_argv->emplace_back(std::move(message));
+            if (!UpdateUsesUserSnapshots(lock.get())) {
+                auto message = misc_name + "," + cow_image_device + "," + source_device;
+                snapuserd_argv->emplace_back(std::move(message));
+            } else {
+                auto message = misc_name + "," + cow_image_device + "," + source_device + "," +
+                               base_path_merge;
+                snapuserd_argv->emplace_back(std::move(message));
+            }
 
             // Do not attempt to connect to the new snapuserd yet, it hasn't
             // been started. We do however want to wait for the misc device
@@ -1539,8 +1749,15 @@ bool SnapshotManager::PerformInitTransition(InitTransition transition,
             continue;
         }
 
-        uint64_t base_sectors =
-                snapuserd_client_->InitDmUserCow(misc_name, cow_image_device, source_device);
+        uint64_t base_sectors;
+        if (!UpdateUsesUserSnapshots(lock.get())) {
+            base_sectors =
+                    snapuserd_client_->InitDmUserCow(misc_name, cow_image_device, source_device);
+        } else {
+            base_sectors = snapuserd_client_->InitDmUserCow(misc_name, cow_image_device,
+                                                            source_device, base_path_merge);
+        }
+
         if (base_sectors == 0) {
             // Unrecoverable as metadata reads from cow device failed
             LOG(FATAL) << "Failed to retrieve base_sectors from Snapuserd";
@@ -1775,30 +1992,36 @@ UpdateState SnapshotManager::GetUpdateState(double* progress) {
         return state;
     }
 
-    // Sum all the snapshot states as if the system consists of a single huge
-    // snapshots device, then compute the merge completion percentage of that
-    // device.
-    std::vector<std::string> snapshots;
-    if (!ListSnapshots(lock.get(), &snapshots)) {
-        LOG(ERROR) << "Could not list snapshots";
-        return state;
+    if (!UpdateUsesUserSnapshots(lock.get())) {
+        // Sum all the snapshot states as if the system consists of a single huge
+        // snapshots device, then compute the merge completion percentage of that
+        // device.
+        std::vector<std::string> snapshots;
+        if (!ListSnapshots(lock.get(), &snapshots)) {
+            LOG(ERROR) << "Could not list snapshots";
+            return state;
+        }
+
+        DmTargetSnapshot::Status fake_snapshots_status = {};
+        for (const auto& snapshot : snapshots) {
+            DmTargetSnapshot::Status current_status;
+
+            if (!IsSnapshotDevice(snapshot)) continue;
+            if (!QuerySnapshotStatus(snapshot, nullptr, &current_status)) continue;
+
+            fake_snapshots_status.sectors_allocated += current_status.sectors_allocated;
+            fake_snapshots_status.total_sectors += current_status.total_sectors;
+            fake_snapshots_status.metadata_sectors += current_status.metadata_sectors;
+        }
+
+        *progress = DmTargetSnapshot::MergePercent(fake_snapshots_status,
+                                                   update_status.sectors_allocated());
+    } else {
+        if (EnsureSnapuserdConnected()) {
+            *progress = snapuserd_client_->GetMergePercent();
+        }
     }
 
-    DmTargetSnapshot::Status fake_snapshots_status = {};
-    for (const auto& snapshot : snapshots) {
-        DmTargetSnapshot::Status current_status;
-
-        if (!IsSnapshotDevice(snapshot)) continue;
-        if (!QuerySnapshotStatus(snapshot, nullptr, &current_status)) continue;
-
-        fake_snapshots_status.sectors_allocated += current_status.sectors_allocated;
-        fake_snapshots_status.total_sectors += current_status.total_sectors;
-        fake_snapshots_status.metadata_sectors += current_status.metadata_sectors;
-    }
-
-    *progress = DmTargetSnapshot::MergePercent(fake_snapshots_status,
-                                               update_status.sectors_allocated());
-
     return state;
 }
 
@@ -1813,6 +2036,38 @@ bool SnapshotManager::UpdateUsesCompression(LockedFile* lock) {
     return update_status.compression_enabled();
 }
 
+bool SnapshotManager::UpdateUsesUserSnapshots() {
+    // This and the following function is constantly
+    // invoked during snapshot merge. We want to avoid
+    // constantly reading from disk. Hence, store this
+    // value in memory.
+    //
+    // Furthermore, this value in the disk is set
+    // only when OTA is applied and doesn't change
+    // during merge phase. Hence, once we know that
+    // the value is read from disk the very first time,
+    // it is safe to read successive checks from memory.
+    if (is_snapshot_userspace_.has_value()) {
+        return is_snapshot_userspace_.value();
+    }
+
+    auto lock = LockShared();
+    if (!lock) return false;
+
+    return UpdateUsesUserSnapshots(lock.get());
+}
+
+bool SnapshotManager::UpdateUsesUserSnapshots(LockedFile* lock) {
+    // See UpdateUsesUserSnapshots()
+    if (is_snapshot_userspace_.has_value()) {
+        return is_snapshot_userspace_.value();
+    }
+
+    SnapshotUpdateStatus update_status = ReadSnapshotUpdateStatus(lock);
+    is_snapshot_userspace_ = update_status.userspace_snapshots();
+    return is_snapshot_userspace_.value();
+}
+
 bool SnapshotManager::ListSnapshots(LockedFile* lock, std::vector<std::string>* snapshots,
                                     const std::string& suffix) {
     CHECK(lock);
@@ -2040,6 +2295,16 @@ bool SnapshotManager::MapPartitionWithSnapshot(LockedFile* lock,
         paths->target_device = base_path;
     }
 
+    auto remaining_time = GetRemainingTime(params.timeout_ms, begin);
+    if (remaining_time.count() < 0) {
+        return false;
+    }
+
+    // Wait for the base device to appear
+    if (!WaitForDevice(base_path, remaining_time)) {
+        return false;
+    }
+
     if (!live_snapshot_status.has_value()) {
         created_devices.Release();
         return true;
@@ -2053,7 +2318,7 @@ bool SnapshotManager::MapPartitionWithSnapshot(LockedFile* lock,
         return false;
     }
 
-    auto remaining_time = GetRemainingTime(params.timeout_ms, begin);
+    remaining_time = GetRemainingTime(params.timeout_ms, begin);
     if (remaining_time.count() < 0) return false;
 
     std::string cow_name;
@@ -2109,10 +2374,10 @@ bool SnapshotManager::MapPartitionWithSnapshot(LockedFile* lock,
             return false;
         }
 
-        auto name = GetDmUserCowName(params.GetPartitionName());
+        auto name = GetDmUserCowName(params.GetPartitionName(), GetSnapshotDriver(lock));
 
         std::string new_cow_device;
-        if (!MapDmUserCow(lock, name, cow_path, source_device_path, remaining_time,
+        if (!MapDmUserCow(lock, name, cow_path, source_device_path, base_path, remaining_time,
                           &new_cow_device)) {
             LOG(ERROR) << "Could not map dm-user device for partition "
                        << params.GetPartitionName();
@@ -2126,21 +2391,37 @@ bool SnapshotManager::MapPartitionWithSnapshot(LockedFile* lock,
         cow_device = new_cow_device;
     }
 
-    std::string path;
-    if (!MapSnapshot(lock, params.GetPartitionName(), base_device, cow_device, remaining_time,
-                     &path)) {
-        LOG(ERROR) << "Could not map snapshot for partition: " << params.GetPartitionName();
-        return false;
-    }
-    // No need to add params.GetPartitionName() to created_devices since it is immediately released.
+    // For userspace snapshots, dm-user block device itself will act as a
+    // snapshot device. There is one subtle difference - MapSnapshot will create
+    // either snapshot target or snapshot-merge target based on the underlying
+    // state of the snapshot device. If snapshot-merge target is created, merge
+    // will immediately start in the kernel.
+    //
+    // This is no longer true with respect to userspace snapshots. When dm-user
+    // block device is created, we just have the snapshots ready but daemon in
+    // the user-space will not start the merge. We have to explicitly inform the
+    // daemon to resume the merge. Check ProcessUpdateState() call stack.
+    if (!UpdateUsesUserSnapshots(lock)) {
+        std::string path;
+        if (!MapSnapshot(lock, params.GetPartitionName(), base_device, cow_device, remaining_time,
+                         &path)) {
+            LOG(ERROR) << "Could not map snapshot for partition: " << params.GetPartitionName();
+            return false;
+        }
+        // No need to add params.GetPartitionName() to created_devices since it is immediately
+        // released.
 
-    if (paths) {
-        paths->snapshot_device = path;
+        if (paths) {
+            paths->snapshot_device = path;
+        }
+        LOG(INFO) << "Mapped " << params.GetPartitionName() << " as snapshot device at " << path;
+    } else {
+        LOG(INFO) << "Mapped " << params.GetPartitionName() << " as snapshot device at "
+                  << cow_device;
     }
 
     created_devices.Release();
 
-    LOG(INFO) << "Mapped " << params.GetPartitionName() << " as snapshot device at " << path;
     return true;
 }
 
@@ -2148,8 +2429,10 @@ bool SnapshotManager::UnmapPartitionWithSnapshot(LockedFile* lock,
                                                  const std::string& target_partition_name) {
     CHECK(lock);
 
-    if (!UnmapSnapshot(lock, target_partition_name)) {
-        return false;
+    if (!UpdateUsesUserSnapshots(lock)) {
+        if (!UnmapSnapshot(lock, target_partition_name)) {
+            return false;
+        }
     }
 
     if (!UnmapCowDevices(lock, target_partition_name)) {
@@ -2247,8 +2530,17 @@ bool SnapshotManager::UnmapCowDevices(LockedFile* lock, const std::string& name)
     CHECK(lock);
     if (!EnsureImageManager()) return false;
 
-    if (UpdateUsesCompression(lock) && !UnmapDmUserDevice(name)) {
-        return false;
+    if (UpdateUsesCompression(lock)) {
+        if (UpdateUsesUserSnapshots(lock)) {
+            if (!UnmapUserspaceSnapshotDevice(lock, name)) {
+                return false;
+            }
+        } else {
+            auto dm_user_name = GetDmUserCowName(name, GetSnapshotDriver(lock));
+            if (!UnmapDmUserDevice(dm_user_name)) {
+                return false;
+            }
+        }
     }
 
     if (!DeleteDeviceIfExists(GetCowName(name), 4000ms)) {
@@ -2264,8 +2556,7 @@ bool SnapshotManager::UnmapCowDevices(LockedFile* lock, const std::string& name)
     return true;
 }
 
-bool SnapshotManager::UnmapDmUserDevice(const std::string& snapshot_name) {
-    auto dm_user_name = GetDmUserCowName(snapshot_name);
+bool SnapshotManager::UnmapDmUserDevice(const std::string& dm_user_name) {
     if (dm_.GetState(dm_user_name) == DmDeviceState::INVALID) {
         return true;
     }
@@ -2291,6 +2582,46 @@ bool SnapshotManager::UnmapDmUserDevice(const std::string& snapshot_name) {
     return true;
 }
 
+bool SnapshotManager::UnmapUserspaceSnapshotDevice(LockedFile* lock,
+                                                   const std::string& snapshot_name) {
+    auto dm_user_name = GetDmUserCowName(snapshot_name, GetSnapshotDriver(lock));
+    if (dm_.GetState(dm_user_name) == DmDeviceState::INVALID) {
+        return true;
+    }
+
+    CHECK(lock);
+
+    SnapshotStatus snapshot_status;
+
+    if (!ReadSnapshotStatus(lock, snapshot_name, &snapshot_status)) {
+        return false;
+    }
+    // If the merge is complete, then we switch dm tables which is equivalent
+    // to unmap; hence, we can't be deleting the device
+    // as the table would be mounted off partitions and will fail.
+    if (snapshot_status.state() != SnapshotState::MERGE_COMPLETED) {
+        if (!DeleteDeviceIfExists(dm_user_name)) {
+            LOG(ERROR) << "Cannot unmap " << dm_user_name;
+            return false;
+        }
+    }
+
+    if (EnsureSnapuserdConnected()) {
+        if (!snapuserd_client_->WaitForDeviceDelete(dm_user_name)) {
+            LOG(ERROR) << "Failed to wait for " << dm_user_name << " control device to delete";
+            return false;
+        }
+    }
+
+    // Ensure the control device is gone so we don't run into ABA problems.
+    auto control_device = "/dev/dm-user/" + dm_user_name;
+    if (!android::fs_mgr::WaitForFileDeleted(control_device, 10s)) {
+        LOG(ERROR) << "Timed out waiting for " << control_device << " to unlink";
+        return false;
+    }
+    return true;
+}
+
 bool SnapshotManager::MapAllSnapshots(const std::chrono::milliseconds& timeout_ms) {
     auto lock = LockExclusive();
     if (!lock) return false;
@@ -2527,6 +2858,7 @@ bool SnapshotManager::WriteUpdateState(LockedFile* lock, UpdateState state,
         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());
+        status.set_userspace_snapshots(old_status.userspace_snapshots());
     }
     return WriteSnapshotUpdateStatus(lock, status);
 }
@@ -2844,6 +3176,43 @@ Return SnapshotManager::CreateUpdateSnapshots(const DeltaArchiveManifest& manife
     SnapshotUpdateStatus status = ReadSnapshotUpdateStatus(lock.get());
     status.set_state(update_state);
     status.set_compression_enabled(cow_creator.compression_enabled);
+    if (cow_creator.compression_enabled) {
+        if (!device()->IsTestDevice()) {
+            // Userspace snapshots is enabled only if compression is enabled
+            status.set_userspace_snapshots(IsUserspaceSnapshotsEnabled());
+            if (IsUserspaceSnapshotsEnabled()) {
+                is_snapshot_userspace_ = true;
+                LOG(INFO) << "User-space snapshots enabled";
+            } else {
+                is_snapshot_userspace_ = false;
+                LOG(INFO) << "User-space snapshots disabled";
+            }
+
+            // Terminate stale daemon if any
+            std::unique_ptr<SnapuserdClient> snapuserd_client =
+                    SnapuserdClient::Connect(kSnapuserdSocket, 10s);
+            if (snapuserd_client) {
+                snapuserd_client->DetachSnapuserd();
+                snapuserd_client->CloseConnection();
+                snapuserd_client = nullptr;
+            }
+
+            // Clear the cached client if any
+            if (snapuserd_client_) {
+                snapuserd_client_->CloseConnection();
+                snapuserd_client_ = nullptr;
+            }
+        } else {
+            status.set_userspace_snapshots(!IsDmSnapshotTestingEnabled());
+            if (IsDmSnapshotTestingEnabled()) {
+                is_snapshot_userspace_ = false;
+                LOG(INFO) << "User-space snapshots disabled for testing";
+            } else {
+                is_snapshot_userspace_ = true;
+                LOG(INFO) << "User-space snapshots enabled for testing";
+            }
+        }
+    }
     if (!WriteSnapshotUpdateStatus(lock.get(), status)) {
         LOG(ERROR) << "Unable to write new update state";
         return Return::Error();
diff --git a/fs_mgr/libsnapshot/snapshot_stub.cpp b/fs_mgr/libsnapshot/snapshot_stub.cpp
index a8d5b8a1a..4af5367ac 100644
--- a/fs_mgr/libsnapshot/snapshot_stub.cpp
+++ b/fs_mgr/libsnapshot/snapshot_stub.cpp
@@ -121,6 +121,11 @@ bool SnapshotManagerStub::UpdateUsesCompression() {
     return false;
 }
 
+bool SnapshotManagerStub::UpdateUsesUserSnapshots() {
+    LOG(ERROR) << __FUNCTION__ << " should never be called.";
+    return false;
+}
+
 class SnapshotMergeStatsStub : public ISnapshotMergeStats {
     bool Start() override { return false; }
     void set_state(android::snapshot::UpdateState, bool) override {}
diff --git a/fs_mgr/libsnapshot/snapshot_test.cpp b/fs_mgr/libsnapshot/snapshot_test.cpp
index d78ba0a83..f1d76e74d 100644
--- a/fs_mgr/libsnapshot/snapshot_test.cpp
+++ b/fs_mgr/libsnapshot/snapshot_test.cpp
@@ -45,6 +45,8 @@
 #include "partition_cow_creator.h"
 #include "utility.h"
 
+#include <android-base/properties.h>
+
 // Mock classes are not used. Header included to ensure mocked class definition aligns with the
 // class itself.
 #include <libsnapshot/mock_device_info.h>
@@ -272,7 +274,7 @@ class SnapshotTest : public ::testing::Test {
     AssertionResult DeleteSnapshotDevice(const std::string& snapshot) {
         AssertionResult res = AssertionSuccess();
         if (!(res = DeleteDevice(snapshot))) return res;
-        if (!sm->UnmapDmUserDevice(snapshot)) {
+        if (!sm->UnmapDmUserDevice(snapshot + "-user-cow")) {
             return AssertionFailure() << "Cannot delete dm-user device for " << snapshot;
         }
         if (!(res = DeleteDevice(snapshot + "-inner"))) return res;
@@ -2559,5 +2561,15 @@ void SnapshotTestEnvironment::TearDown() {
 int main(int argc, char** argv) {
     ::testing::InitGoogleTest(&argc, argv);
     ::testing::AddGlobalTestEnvironment(new ::android::snapshot::SnapshotTestEnvironment());
-    return RUN_ALL_TESTS();
+
+    android::base::SetProperty("ctl.stop", "snapuserd");
+
+    if (!android::base::SetProperty("snapuserd.test.dm.snapshots", "1")) {
+        return testing::AssertionFailure()
+               << "Failed to disable property: virtual_ab.userspace.snapshots.enabled";
+    }
+
+    int ret = RUN_ALL_TESTS();
+    android::base::SetProperty("snapuserd.test.dm.snapshots", "0");
+    return ret;
 }
diff --git a/fs_mgr/libsnapshot/snapuserd/Android.bp b/fs_mgr/libsnapshot/snapuserd/Android.bp
index c9b05124b..93b0f7c28 100644
--- a/fs_mgr/libsnapshot/snapuserd/Android.bp
+++ b/fs_mgr/libsnapshot/snapuserd/Android.bp
@@ -142,3 +142,39 @@ cc_test {
     auto_gen_config: true,
     require_root: false,
 }
+
+cc_test {
+    name: "snapuserd_test",
+    defaults: [
+        "fs_mgr_defaults",
+    ],
+    srcs: [
+        "user-space-merge/snapuserd_test.cpp",
+    ],
+    cflags: [
+        "-Wall",
+        "-Werror",
+    ],
+    shared_libs: [
+        "libbase",
+        "liblog",
+    ],
+    static_libs: [
+        "libbrotli",
+        "libgtest",
+        "libsnapshot_cow",
+        "libsnapshot_snapuserd",
+        "libcutils_sockets",
+        "libz",
+        "libfs_mgr",
+        "libdm",
+        "libext4_utils",
+    ],
+    header_libs: [
+        "libstorage_literals_headers",
+        "libfiemap_headers",
+    ],
+    test_min_api_level: 30,
+    auto_gen_config: true,
+    require_root: false,
+}
diff --git a/fs_mgr/libsnapshot/snapuserd/include/snapuserd/snapuserd_client.h b/fs_mgr/libsnapshot/snapuserd/include/snapuserd/snapuserd_client.h
index 6ed55af2d..cebda1ccb 100644
--- a/fs_mgr/libsnapshot/snapuserd/include/snapuserd/snapuserd_client.h
+++ b/fs_mgr/libsnapshot/snapuserd/include/snapuserd/snapuserd_client.h
@@ -63,7 +63,8 @@ class SnapuserdClient {
     // The misc_name must be the "misc_name" given to dm-user in step 2.
     //
     uint64_t InitDmUserCow(const std::string& misc_name, const std::string& cow_device,
-                           const std::string& backing_device);
+                           const std::string& backing_device,
+                           const std::string& base_path_merge = "");
     bool AttachDmUser(const std::string& misc_name);
 
     // Wait for snapuserd to disassociate with a dm-user control device. This
diff --git a/fs_mgr/libsnapshot/snapuserd/snapuserd_client.cpp b/fs_mgr/libsnapshot/snapuserd/snapuserd_client.cpp
index e345269ae..7b1c7a3a6 100644
--- a/fs_mgr/libsnapshot/snapuserd/snapuserd_client.cpp
+++ b/fs_mgr/libsnapshot/snapuserd/snapuserd_client.cpp
@@ -195,8 +195,16 @@ bool SnapuserdClient::AttachDmUser(const std::string& misc_name) {
 }
 
 uint64_t SnapuserdClient::InitDmUserCow(const std::string& misc_name, const std::string& cow_device,
-                                        const std::string& backing_device) {
-    std::vector<std::string> parts = {"init", misc_name, cow_device, backing_device};
+                                        const std::string& backing_device,
+                                        const std::string& base_path_merge) {
+    std::vector<std::string> parts;
+
+    if (base_path_merge.empty()) {
+        parts = {"init", misc_name, cow_device, backing_device};
+    } else {
+        // For userspace snapshots
+        parts = {"init", misc_name, cow_device, backing_device, base_path_merge};
+    }
     std::string msg = android::base::Join(parts, ",");
     if (!Sendmsg(msg)) {
         LOG(ERROR) << "Failed to send message " << msg << " to snapuserd daemon";
diff --git a/fs_mgr/libsnapshot/snapuserd/snapuserd_daemon.cpp b/fs_mgr/libsnapshot/snapuserd/snapuserd_daemon.cpp
index 912884fd3..ddb1f7999 100644
--- a/fs_mgr/libsnapshot/snapuserd/snapuserd_daemon.cpp
+++ b/fs_mgr/libsnapshot/snapuserd/snapuserd_daemon.cpp
@@ -14,25 +14,95 @@
  * limitations under the License.
  */
 
-#include "snapuserd_daemon.h"
-
 #include <android-base/logging.h>
+#include <android-base/properties.h>
 #include <android-base/strings.h>
 #include <gflags/gflags.h>
 #include <snapuserd/snapuserd_client.h>
 
+#include "snapuserd_daemon.h"
+
 DEFINE_string(socket, android::snapshot::kSnapuserdSocket, "Named socket or socket path.");
 DEFINE_bool(no_socket, false,
             "If true, no socket is used. Each additional argument is an INIT message.");
 DEFINE_bool(socket_handoff, false,
             "If true, perform a socket hand-off with an existing snapuserd instance, then exit.");
+DEFINE_bool(user_snapshot, false, "If true, user-space snapshots are used");
 
 namespace android {
 namespace snapshot {
 
-bool Daemon::StartServer(int argc, char** argv) {
+bool Daemon::IsUserspaceSnapshotsEnabled() {
+    return android::base::GetBoolProperty("ro.virtual_ab.userspace.snapshots.enabled", false);
+}
+
+bool Daemon::IsDmSnapshotTestingEnabled() {
+    return android::base::GetBoolProperty("snapuserd.test.dm.snapshots", false);
+}
+
+bool Daemon::StartDaemon(int argc, char** argv) {
     int arg_start = gflags::ParseCommandLineFlags(&argc, &argv, true);
 
+    // Daemon launched from first stage init and during selinux transition
+    // will have the command line "-user_snapshot" flag set if the user-space
+    // snapshots are enabled.
+    //
+    // Daemon launched as a init service during "socket-handoff" and when OTA
+    // is applied will check for the property. This is ok as the system
+    // properties are valid at this point. We can't do this during first
+    // stage init and hence use the command line flags to get the information.
+    if (!IsDmSnapshotTestingEnabled() && (FLAGS_user_snapshot || IsUserspaceSnapshotsEnabled())) {
+        LOG(INFO) << "Starting daemon for user-space snapshots.....";
+        return StartServerForUserspaceSnapshots(arg_start, argc, argv);
+    } else {
+        LOG(INFO) << "Starting daemon for dm-snapshots.....";
+        return StartServerForDmSnapshot(arg_start, argc, argv);
+    }
+}
+
+bool Daemon::StartServerForUserspaceSnapshots(int arg_start, int argc, char** argv) {
+    sigfillset(&signal_mask_);
+    sigdelset(&signal_mask_, SIGINT);
+    sigdelset(&signal_mask_, SIGTERM);
+    sigdelset(&signal_mask_, SIGUSR1);
+
+    // Masking signals here ensure that after this point, we won't handle INT/TERM
+    // until after we call into ppoll()
+    signal(SIGINT, Daemon::SignalHandler);
+    signal(SIGTERM, Daemon::SignalHandler);
+    signal(SIGPIPE, Daemon::SignalHandler);
+    signal(SIGUSR1, Daemon::SignalHandler);
+
+    MaskAllSignalsExceptIntAndTerm();
+
+    if (FLAGS_socket_handoff) {
+        return user_server_.RunForSocketHandoff();
+    }
+    if (!FLAGS_no_socket) {
+        if (!user_server_.Start(FLAGS_socket)) {
+            return false;
+        }
+        return user_server_.Run();
+    }
+
+    for (int i = arg_start; i < argc; i++) {
+        auto parts = android::base::Split(argv[i], ",");
+        if (parts.size() != 4) {
+            LOG(ERROR) << "Malformed message, expected three sub-arguments.";
+            return false;
+        }
+        auto handler = user_server_.AddHandler(parts[0], parts[1], parts[2], parts[3]);
+        if (!handler || !user_server_.StartHandler(handler)) {
+            return false;
+        }
+    }
+
+    // Skip the accept() call to avoid spurious log spam. The server will still
+    // run until all handlers have completed.
+    return user_server_.WaitForSocket();
+}
+
+bool Daemon::StartServerForDmSnapshot(int arg_start, int argc, char** argv) {
     sigfillset(&signal_mask_);
     sigdelset(&signal_mask_, SIGINT);
     sigdelset(&signal_mask_, SIGTERM);
@@ -95,11 +165,19 @@ void Daemon::MaskAllSignals() {
 }
 
 void Daemon::Interrupt() {
-    server_.Interrupt();
+    if (IsUserspaceSnapshotsEnabled()) {
+        user_server_.Interrupt();
+    } else {
+        server_.Interrupt();
+    }
 }
 
 void Daemon::ReceivedSocketSignal() {
-    server_.ReceivedSocketSignal();
+    if (IsUserspaceSnapshotsEnabled()) {
+        user_server_.ReceivedSocketSignal();
+    } else {
+        server_.ReceivedSocketSignal();
+    }
 }
 
 void Daemon::SignalHandler(int signal) {
@@ -133,9 +211,10 @@ int main(int argc, char** argv) {
 
     android::snapshot::Daemon& daemon = android::snapshot::Daemon::Instance();
 
-    if (!daemon.StartServer(argc, argv)) {
-        LOG(ERROR) << "Snapuserd daemon failed to start.";
+    if (!daemon.StartDaemon(argc, argv)) {
+        LOG(ERROR) << "Snapuserd daemon failed to start";
         exit(EXIT_FAILURE);
     }
+
     return 0;
 }
diff --git a/fs_mgr/libsnapshot/snapuserd/snapuserd_daemon.h b/fs_mgr/libsnapshot/snapuserd/snapuserd_daemon.h
index fbf57d943..cf3b917e5 100644
--- a/fs_mgr/libsnapshot/snapuserd/snapuserd_daemon.h
+++ b/fs_mgr/libsnapshot/snapuserd/snapuserd_daemon.h
@@ -20,6 +20,7 @@
 #include <vector>
 
 #include "dm-snapshot-merge/snapuserd_server.h"
+#include "user-space-merge/snapuserd_server.h"
 
 namespace android {
 namespace snapshot {
@@ -35,9 +36,13 @@ class Daemon {
         return instance;
     }
 
-    bool StartServer(int argc, char** argv);
+    bool StartServerForDmSnapshot(int arg_start, int argc, char** argv);
+    bool StartServerForUserspaceSnapshots(int arg_start, int argc, char** argv);
     void Interrupt();
     void ReceivedSocketSignal();
+    bool IsUserspaceSnapshotsEnabled();
+    bool IsDmSnapshotTestingEnabled();
+    bool StartDaemon(int argc, char** argv);
 
   private:
     // Signal mask used with ppoll()
@@ -47,6 +52,7 @@ class Daemon {
     void operator=(Daemon const&) = delete;
 
     SnapuserdServer server_;
+    UserSnapshotServer user_server_;
     void MaskAllSignalsExceptIntAndTerm();
     void MaskAllSignals();
     static void SignalHandler(int signal);
diff --git a/fs_mgr/libsnapshot/snapuserd/user-space-merge/snapuserd_core.cpp b/fs_mgr/libsnapshot/snapuserd/user-space-merge/snapuserd_core.cpp
index 57e47e7ed..95d95cdb8 100644
--- a/fs_mgr/libsnapshot/snapuserd/user-space-merge/snapuserd_core.cpp
+++ b/fs_mgr/libsnapshot/snapuserd/user-space-merge/snapuserd_core.cpp
@@ -35,7 +35,7 @@ SnapshotHandler::SnapshotHandler(std::string misc_name, std::string cow_device,
 }
 
 bool SnapshotHandler::InitializeWorkers() {
-    for (int i = 0; i < NUM_THREADS_PER_PARTITION; i++) {
+    for (int i = 0; i < kNumWorkerThreads; i++) {
         std::unique_ptr<Worker> wt =
                 std::make_unique<Worker>(cow_device_, backing_store_device_, control_device_,
                                          misc_name_, base_path_merge_, GetSharedPtr());
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 13b56facb..195331691 100644
--- a/fs_mgr/libsnapshot/snapuserd/user-space-merge/snapuserd_core.h
+++ b/fs_mgr/libsnapshot/snapuserd/user-space-merge/snapuserd_core.h
@@ -48,10 +48,10 @@ namespace snapshot {
 using android::base::unique_fd;
 using namespace std::chrono_literals;
 
-static constexpr size_t PAYLOAD_SIZE = (1UL << 20);
-static_assert(PAYLOAD_SIZE >= BLOCK_SZ);
+static constexpr size_t PAYLOAD_BUFFER_SZ = (1UL << 20);
+static_assert(PAYLOAD_BUFFER_SZ >= BLOCK_SZ);
 
-static constexpr int NUM_THREADS_PER_PARTITION = 1;
+static constexpr int kNumWorkerThreads = 4;
 
 #define SNAP_LOG(level) LOG(level) << misc_name_ << ": "
 #define SNAP_PLOG(level) PLOG(level) << misc_name_ << ": "
diff --git a/fs_mgr/libsnapshot/snapuserd/user-space-merge/snapuserd_dm_user.cpp b/fs_mgr/libsnapshot/snapuserd/user-space-merge/snapuserd_dm_user.cpp
index bfbacf92e..1e300d2d4 100644
--- a/fs_mgr/libsnapshot/snapuserd/user-space-merge/snapuserd_dm_user.cpp
+++ b/fs_mgr/libsnapshot/snapuserd/user-space-merge/snapuserd_dm_user.cpp
@@ -231,8 +231,8 @@ void Worker::InitializeBufsink() {
     // Allocate the buffer which is used to communicate between
     // daemon and dm-user. The buffer comprises of header and a fixed payload.
     // If the dm-user requests a big IO, the IO will be broken into chunks
-    // of PAYLOAD_SIZE.
-    size_t buf_size = sizeof(struct dm_user_header) + PAYLOAD_SIZE;
+    // of PAYLOAD_BUFFER_SZ.
+    size_t buf_size = sizeof(struct dm_user_header) + PAYLOAD_BUFFER_SZ;
     bufsink_.Initialize(buf_size);
 }
 
@@ -326,7 +326,7 @@ bool Worker::ReadAlignedSector(sector_t sector, size_t sz, bool header_response)
 
     do {
         // Process 1MB payload at a time
-        size_t read_size = std::min(PAYLOAD_SIZE, remaining_size);
+        size_t read_size = std::min(PAYLOAD_BUFFER_SZ, remaining_size);
 
         header->type = DM_USER_RESP_SUCCESS;
         size_t total_bytes_read = 0;
diff --git a/fs_mgr/libsnapshot/snapuserd/user-space-merge/snapuserd_merge.cpp b/fs_mgr/libsnapshot/snapuserd/user-space-merge/snapuserd_merge.cpp
index 47fc7db50..fa055b730 100644
--- a/fs_mgr/libsnapshot/snapuserd/user-space-merge/snapuserd_merge.cpp
+++ b/fs_mgr/libsnapshot/snapuserd/user-space-merge/snapuserd_merge.cpp
@@ -81,11 +81,11 @@ bool Worker::MergeReplaceZeroOps(const std::unique_ptr<ICowOpIter>& cowop_iter)
     // Why 2048 ops ? We can probably increase this to bigger value but just
     // need to ensure that merge makes forward progress if there are
     // crashes repeatedly which is highly unlikely.
-    int total_ops_merged_per_commit = (PAYLOAD_SIZE / BLOCK_SZ) * 8;
+    int total_ops_merged_per_commit = (PAYLOAD_BUFFER_SZ / BLOCK_SZ) * 8;
     int num_ops_merged = 0;
 
     while (!cowop_iter->Done()) {
-        int num_ops = PAYLOAD_SIZE / BLOCK_SZ;
+        int num_ops = PAYLOAD_BUFFER_SZ / BLOCK_SZ;
         std::vector<const CowOperation*> replace_zero_vec;
         uint64_t source_offset;
 
@@ -292,6 +292,7 @@ bool Worker::RunMergeThread() {
 
     if (!Init()) {
         SNAP_LOG(ERROR) << "Merge thread initialization failed...";
+        snapuserd_->MergeFailed();
         return false;
     }
 
diff --git a/fs_mgr/libsnapshot/snapuserd/user-space-merge/snapuserd_readahead.cpp b/fs_mgr/libsnapshot/snapuserd/user-space-merge/snapuserd_readahead.cpp
index 0bcf26e78..40e7242cd 100644
--- a/fs_mgr/libsnapshot/snapuserd/user-space-merge/snapuserd_readahead.cpp
+++ b/fs_mgr/libsnapshot/snapuserd/user-space-merge/snapuserd_readahead.cpp
@@ -429,7 +429,7 @@ void ReadAhead::InitializeBuffer() {
             static_cast<void*>((char*)mapped_addr + snapuserd_->GetBufferMetadataOffset());
     read_ahead_buffer_ = static_cast<void*>((char*)mapped_addr + snapuserd_->GetBufferDataOffset());
     // For xor ops
-    bufsink_.Initialize(PAYLOAD_SIZE);
+    bufsink_.Initialize(PAYLOAD_BUFFER_SZ);
 }
 
 }  // namespace snapshot
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 a4fd5a035..a79e3e13a 100644
--- a/fs_mgr/libsnapshot/snapuserd/user-space-merge/snapuserd_server.cpp
+++ b/fs_mgr/libsnapshot/snapuserd/user-space-merge/snapuserd_server.cpp
@@ -44,7 +44,7 @@ using namespace std::string_literals;
 using android::base::borrowed_fd;
 using android::base::unique_fd;
 
-DaemonOps SnapuserServer::Resolveop(std::string& input) {
+DaemonOps UserSnapshotServer::Resolveop(std::string& input) {
     if (input == "init") return DaemonOps::INIT;
     if (input == "start") return DaemonOps::START;
     if (input == "stop") return DaemonOps::STOP;
@@ -59,14 +59,14 @@ DaemonOps SnapuserServer::Resolveop(std::string& input) {
     return DaemonOps::INVALID;
 }
 
-SnapuserServer::~SnapuserServer() {
+UserSnapshotServer::~UserSnapshotServer() {
     // Close any client sockets that were added via AcceptClient().
     for (size_t i = 1; i < watched_fds_.size(); i++) {
         close(watched_fds_[i].fd);
     }
 }
 
-std::string SnapuserServer::GetDaemonStatus() {
+std::string UserSnapshotServer::GetDaemonStatus() {
     std::string msg = "";
 
     if (IsTerminating())
@@ -77,8 +77,8 @@ std::string SnapuserServer::GetDaemonStatus() {
     return msg;
 }
 
-void SnapuserServer::Parsemsg(std::string const& msg, const char delim,
-                              std::vector<std::string>& out) {
+void UserSnapshotServer::Parsemsg(std::string const& msg, const char delim,
+                                  std::vector<std::string>& out) {
     std::stringstream ss(msg);
     std::string s;
 
@@ -87,15 +87,15 @@ void SnapuserServer::Parsemsg(std::string const& msg, const char delim,
     }
 }
 
-void SnapuserServer::ShutdownThreads() {
+void UserSnapshotServer::ShutdownThreads() {
     terminating_ = true;
     JoinAllThreads();
 }
 
-DmUserHandler::DmUserHandler(std::shared_ptr<SnapshotHandler> snapuserd)
+UserSnapshotDmUserHandler::UserSnapshotDmUserHandler(std::shared_ptr<SnapshotHandler> snapuserd)
     : snapuserd_(snapuserd), misc_name_(snapuserd_->GetMiscName()) {}
 
-bool SnapuserServer::Sendmsg(android::base::borrowed_fd fd, const std::string& msg) {
+bool UserSnapshotServer::Sendmsg(android::base::borrowed_fd fd, const std::string& msg) {
     ssize_t ret = TEMP_FAILURE_RETRY(send(fd.get(), msg.data(), msg.size(), MSG_NOSIGNAL));
     if (ret < 0) {
         PLOG(ERROR) << "Snapuserd:server: send() failed";
@@ -109,8 +109,8 @@ bool SnapuserServer::Sendmsg(android::base::borrowed_fd fd, const std::string& m
     return true;
 }
 
-bool SnapuserServer::Recv(android::base::borrowed_fd fd, std::string* data) {
-    char msg[MAX_PACKET_SIZE];
+bool UserSnapshotServer::Recv(android::base::borrowed_fd fd, std::string* data) {
+    char msg[kMaxPacketSize];
     ssize_t rv = TEMP_FAILURE_RETRY(recv(fd.get(), msg, sizeof(msg), 0));
     if (rv < 0) {
         PLOG(ERROR) << "recv failed";
@@ -120,7 +120,7 @@ bool SnapuserServer::Recv(android::base::borrowed_fd fd, std::string* data) {
     return true;
 }
 
-bool SnapuserServer::Receivemsg(android::base::borrowed_fd fd, const std::string& str) {
+bool UserSnapshotServer::Receivemsg(android::base::borrowed_fd fd, const std::string& str) {
     const char delim = ',';
 
     std::vector<std::string> out;
@@ -290,7 +290,7 @@ bool SnapuserServer::Receivemsg(android::base::borrowed_fd fd, const std::string
     }
 }
 
-void SnapuserServer::RunThread(std::shared_ptr<DmUserHandler> handler) {
+void UserSnapshotServer::RunThread(std::shared_ptr<UserSnapshotDmUserHandler> handler) {
     LOG(INFO) << "Entering thread for handler: " << handler->misc_name();
 
     handler->snapuserd()->SetSocketPresent(is_socket_present_);
@@ -337,7 +337,7 @@ void SnapuserServer::RunThread(std::shared_ptr<DmUserHandler> handler) {
     }
 }
 
-bool SnapuserServer::Start(const std::string& socketname) {
+bool UserSnapshotServer::Start(const std::string& socketname) {
     bool start_listening = true;
 
     sockfd_.reset(android_get_control_socket(socketname.c_str()));
@@ -353,7 +353,7 @@ bool SnapuserServer::Start(const std::string& socketname) {
     return StartWithSocket(start_listening);
 }
 
-bool SnapuserServer::StartWithSocket(bool start_listening) {
+bool UserSnapshotServer::StartWithSocket(bool start_listening) {
     if (start_listening && listen(sockfd_.get(), 4) < 0) {
         PLOG(ERROR) << "listen socket failed";
         return false;
@@ -374,7 +374,7 @@ bool SnapuserServer::StartWithSocket(bool start_listening) {
     return true;
 }
 
-bool SnapuserServer::Run() {
+bool UserSnapshotServer::Run() {
     LOG(INFO) << "Now listening on snapuserd socket";
 
     while (!IsTerminating()) {
@@ -406,9 +406,9 @@ bool SnapuserServer::Run() {
     return true;
 }
 
-void SnapuserServer::JoinAllThreads() {
+void UserSnapshotServer::JoinAllThreads() {
     // Acquire the thread list within the lock.
-    std::vector<std::shared_ptr<DmUserHandler>> dm_users;
+    std::vector<std::shared_ptr<UserSnapshotDmUserHandler>> dm_users;
     {
         std::lock_guard<std::mutex> guard(lock_);
         dm_users = std::move(dm_users_);
@@ -421,14 +421,14 @@ void SnapuserServer::JoinAllThreads() {
     }
 }
 
-void SnapuserServer::AddWatchedFd(android::base::borrowed_fd fd, int events) {
+void UserSnapshotServer::AddWatchedFd(android::base::borrowed_fd fd, int events) {
     struct pollfd p = {};
     p.fd = fd.get();
     p.events = events;
     watched_fds_.emplace_back(std::move(p));
 }
 
-void SnapuserServer::AcceptClient() {
+void UserSnapshotServer::AcceptClient() {
     int fd = TEMP_FAILURE_RETRY(accept4(sockfd_.get(), nullptr, nullptr, SOCK_CLOEXEC));
     if (fd < 0) {
         PLOG(ERROR) << "accept4 failed";
@@ -438,7 +438,7 @@ void SnapuserServer::AcceptClient() {
     AddWatchedFd(fd, POLLIN);
 }
 
-bool SnapuserServer::HandleClient(android::base::borrowed_fd fd, int revents) {
+bool UserSnapshotServer::HandleClient(android::base::borrowed_fd fd, int revents) {
     if (revents & POLLHUP) {
         LOG(DEBUG) << "Snapuserd client disconnected";
         return false;
@@ -455,16 +455,15 @@ bool SnapuserServer::HandleClient(android::base::borrowed_fd fd, int revents) {
     return true;
 }
 
-void SnapuserServer::Interrupt() {
+void UserSnapshotServer::Interrupt() {
     // Force close the socket so poll() fails.
     sockfd_ = {};
     SetTerminating();
 }
 
-std::shared_ptr<DmUserHandler> SnapuserServer::AddHandler(const std::string& misc_name,
-                                                          const std::string& cow_device_path,
-                                                          const std::string& backing_device,
-                                                          const std::string& base_path_merge) {
+std::shared_ptr<UserSnapshotDmUserHandler> UserSnapshotServer::AddHandler(
+        const std::string& misc_name, const std::string& cow_device_path,
+        const std::string& backing_device, const std::string& base_path_merge) {
     auto snapuserd = std::make_shared<SnapshotHandler>(misc_name, cow_device_path, backing_device,
                                                        base_path_merge);
     if (!snapuserd->InitCowDevice()) {
@@ -477,7 +476,7 @@ std::shared_ptr<DmUserHandler> SnapuserServer::AddHandler(const std::string& mis
         return nullptr;
     }
 
-    auto handler = std::make_shared<DmUserHandler>(snapuserd);
+    auto handler = std::make_shared<UserSnapshotDmUserHandler>(snapuserd);
     {
         std::lock_guard<std::mutex> lock(lock_);
         if (FindHandler(&lock, misc_name) != dm_users_.end()) {
@@ -489,7 +488,7 @@ std::shared_ptr<DmUserHandler> SnapuserServer::AddHandler(const std::string& mis
     return handler;
 }
 
-bool SnapuserServer::StartHandler(const std::shared_ptr<DmUserHandler>& handler) {
+bool UserSnapshotServer::StartHandler(const std::shared_ptr<UserSnapshotDmUserHandler>& handler) {
     if (handler->snapuserd()->IsAttached()) {
         LOG(ERROR) << "Handler already attached";
         return false;
@@ -497,11 +496,11 @@ bool SnapuserServer::StartHandler(const std::shared_ptr<DmUserHandler>& handler)
 
     handler->snapuserd()->AttachControlDevice();
 
-    handler->thread() = std::thread(std::bind(&SnapuserServer::RunThread, this, handler));
+    handler->thread() = std::thread(std::bind(&UserSnapshotServer::RunThread, this, handler));
     return true;
 }
 
-bool SnapuserServer::StartMerge(const std::shared_ptr<DmUserHandler>& handler) {
+bool UserSnapshotServer::StartMerge(const std::shared_ptr<UserSnapshotDmUserHandler>& handler) {
     if (!handler->snapuserd()->IsAttached()) {
         LOG(ERROR) << "Handler not attached to dm-user - Merge thread cannot be started";
         return false;
@@ -511,8 +510,8 @@ bool SnapuserServer::StartMerge(const std::shared_ptr<DmUserHandler>& handler) {
     return true;
 }
 
-auto SnapuserServer::FindHandler(std::lock_guard<std::mutex>* proof_of_lock,
-                                 const std::string& misc_name) -> HandlerList::iterator {
+auto UserSnapshotServer::FindHandler(std::lock_guard<std::mutex>* proof_of_lock,
+                                     const std::string& misc_name) -> HandlerList::iterator {
     CHECK(proof_of_lock);
 
     for (auto iter = dm_users_.begin(); iter != dm_users_.end(); iter++) {
@@ -523,7 +522,7 @@ auto SnapuserServer::FindHandler(std::lock_guard<std::mutex>* proof_of_lock,
     return dm_users_.end();
 }
 
-void SnapuserServer::TerminateMergeThreads(std::lock_guard<std::mutex>* proof_of_lock) {
+void UserSnapshotServer::TerminateMergeThreads(std::lock_guard<std::mutex>* proof_of_lock) {
     CHECK(proof_of_lock);
 
     for (auto iter = dm_users_.begin(); iter != dm_users_.end(); iter++) {
@@ -533,11 +532,12 @@ void SnapuserServer::TerminateMergeThreads(std::lock_guard<std::mutex>* proof_of
     }
 }
 
-std::string SnapuserServer::GetMergeStatus(const std::shared_ptr<DmUserHandler>& handler) {
+std::string UserSnapshotServer::GetMergeStatus(
+        const std::shared_ptr<UserSnapshotDmUserHandler>& handler) {
     return handler->snapuserd()->GetMergeStatus();
 }
 
-double SnapuserServer::GetMergePercentage(std::lock_guard<std::mutex>* proof_of_lock) {
+double UserSnapshotServer::GetMergePercentage(std::lock_guard<std::mutex>* proof_of_lock) {
     CHECK(proof_of_lock);
     double percentage = 0.0;
     int n = 0;
@@ -567,8 +567,8 @@ double SnapuserServer::GetMergePercentage(std::lock_guard<std::mutex>* proof_of_
     return percentage;
 }
 
-bool SnapuserServer::RemoveAndJoinHandler(const std::string& misc_name) {
-    std::shared_ptr<DmUserHandler> handler;
+bool UserSnapshotServer::RemoveAndJoinHandler(const std::string& misc_name) {
+    std::shared_ptr<UserSnapshotDmUserHandler> handler;
     {
         std::lock_guard<std::mutex> lock(lock_);
 
@@ -588,7 +588,7 @@ bool SnapuserServer::RemoveAndJoinHandler(const std::string& misc_name) {
     return true;
 }
 
-bool SnapuserServer::WaitForSocket() {
+bool UserSnapshotServer::WaitForSocket() {
     auto scope_guard = android::base::make_scope_guard([this]() -> void { JoinAllThreads(); });
 
     auto socket_path = ANDROID_SOCKET_DIR "/"s + kSnapuserdSocketProxy;
@@ -642,7 +642,7 @@ bool SnapuserServer::WaitForSocket() {
     return Run();
 }
 
-bool SnapuserServer::RunForSocketHandoff() {
+bool UserSnapshotServer::RunForSocketHandoff() {
     unique_fd proxy_fd(android_get_control_socket(kSnapuserdSocketProxy));
     if (proxy_fd < 0) {
         PLOG(FATAL) << "Proxy could not get android control socket " << kSnapuserdSocketProxy;
diff --git a/fs_mgr/libsnapshot/snapuserd/user-space-merge/snapuserd_server.h b/fs_mgr/libsnapshot/snapuserd/user-space-merge/snapuserd_server.h
index e93621ca2..c645456bc 100644
--- a/fs_mgr/libsnapshot/snapuserd/user-space-merge/snapuserd_server.h
+++ b/fs_mgr/libsnapshot/snapuserd/user-space-merge/snapuserd_server.h
@@ -33,7 +33,7 @@
 namespace android {
 namespace snapshot {
 
-static constexpr uint32_t MAX_PACKET_SIZE = 512;
+static constexpr uint32_t kMaxPacketSize = 512;
 
 enum class DaemonOps {
     INIT,
@@ -49,9 +49,9 @@ enum class DaemonOps {
     INVALID,
 };
 
-class DmUserHandler {
+class UserSnapshotDmUserHandler {
   public:
-    explicit DmUserHandler(std::shared_ptr<SnapshotHandler> snapuserd);
+    explicit UserSnapshotDmUserHandler(std::shared_ptr<SnapshotHandler> snapuserd);
 
     void FreeResources() {
         // Each worker thread holds a reference to snapuserd.
@@ -76,7 +76,7 @@ class DmUserHandler {
     bool thread_terminated_ = false;
 };
 
-class SnapuserServer {
+class UserSnapshotServer {
   private:
     android::base::unique_fd sockfd_;
     bool terminating_;
@@ -87,7 +87,7 @@ class SnapuserServer {
 
     std::mutex lock_;
 
-    using HandlerList = std::vector<std::shared_ptr<DmUserHandler>>;
+    using HandlerList = std::vector<std::shared_ptr<UserSnapshotDmUserHandler>>;
     HandlerList dm_users_;
 
     void AddWatchedFd(android::base::borrowed_fd fd, int events);
@@ -105,11 +105,11 @@ class SnapuserServer {
 
     bool IsTerminating() { return terminating_; }
 
-    void RunThread(std::shared_ptr<DmUserHandler> handler);
+    void RunThread(std::shared_ptr<UserSnapshotDmUserHandler> handler);
     void JoinAllThreads();
     bool StartWithSocket(bool start_listening);
 
-    // Find a DmUserHandler within a lock.
+    // Find a UserSnapshotDmUserHandler within a lock.
     HandlerList::iterator FindHandler(std::lock_guard<std::mutex>* proof_of_lock,
                                       const std::string& misc_name);
 
@@ -117,8 +117,8 @@ class SnapuserServer {
     void TerminateMergeThreads(std::lock_guard<std::mutex>* proof_of_lock);
 
   public:
-    SnapuserServer() { terminating_ = false; }
-    ~SnapuserServer();
+    UserSnapshotServer() { terminating_ = false; }
+    ~UserSnapshotServer();
 
     bool Start(const std::string& socketname);
     bool Run();
@@ -126,13 +126,13 @@ class SnapuserServer {
     bool RunForSocketHandoff();
     bool WaitForSocket();
 
-    std::shared_ptr<DmUserHandler> AddHandler(const std::string& misc_name,
-                                              const std::string& cow_device_path,
-                                              const std::string& backing_device,
-                                              const std::string& base_path_merge);
-    bool StartHandler(const std::shared_ptr<DmUserHandler>& handler);
-    bool StartMerge(const std::shared_ptr<DmUserHandler>& handler);
-    std::string GetMergeStatus(const std::shared_ptr<DmUserHandler>& handler);
+    std::shared_ptr<UserSnapshotDmUserHandler> AddHandler(const std::string& misc_name,
+                                                          const std::string& cow_device_path,
+                                                          const std::string& backing_device,
+                                                          const std::string& base_path_merge);
+    bool StartHandler(const std::shared_ptr<UserSnapshotDmUserHandler>& handler);
+    bool StartMerge(const std::shared_ptr<UserSnapshotDmUserHandler>& handler);
+    std::string GetMergeStatus(const std::shared_ptr<UserSnapshotDmUserHandler>& handler);
 
     void SetTerminating() { terminating_ = true; }
     void ReceivedSocketSignal() { received_socket_signal_ = true; }
diff --git a/fs_mgr/libsnapshot/snapuserd/user-space-merge/snapuserd_test.cpp b/fs_mgr/libsnapshot/snapuserd/user-space-merge/snapuserd_test.cpp
new file mode 100644
index 000000000..1c3e04b89
--- /dev/null
+++ b/fs_mgr/libsnapshot/snapuserd/user-space-merge/snapuserd_test.cpp
@@ -0,0 +1,861 @@
+// Copyright (C) 2018 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.
+
+#include <fcntl.h>
+#include <linux/fs.h>
+#include <linux/memfd.h>
+#include <sys/ioctl.h>
+#include <sys/stat.h>
+#include <sys/syscall.h>
+#include <sys/types.h>
+#include <unistd.h>
+
+#include <chrono>
+#include <iostream>
+#include <memory>
+#include <string_view>
+
+#include <android-base/file.h>
+#include <android-base/unique_fd.h>
+#include <fs_mgr/file_wait.h>
+#include <gtest/gtest.h>
+#include <libdm/dm.h>
+#include <libdm/loop_control.h>
+#include <libsnapshot/cow_writer.h>
+#include <snapuserd/snapuserd_client.h>
+#include <storage_literals/storage_literals.h>
+
+#include "snapuserd_core.h"
+
+namespace android {
+namespace snapshot {
+
+using namespace android::storage_literals;
+using android::base::unique_fd;
+using LoopDevice = android::dm::LoopDevice;
+using namespace std::chrono_literals;
+using namespace android::dm;
+using namespace std;
+
+static constexpr char kSnapuserdSocketTest[] = "snapuserdTest";
+
+class Tempdevice {
+  public:
+    Tempdevice(const std::string& name, const DmTable& table)
+        : dm_(DeviceMapper::Instance()), name_(name), valid_(false) {
+        valid_ = dm_.CreateDevice(name, table, &path_, std::chrono::seconds(5));
+    }
+    Tempdevice(Tempdevice&& other) noexcept
+        : dm_(other.dm_), name_(other.name_), path_(other.path_), valid_(other.valid_) {
+        other.valid_ = false;
+    }
+    ~Tempdevice() {
+        if (valid_) {
+            dm_.DeleteDevice(name_);
+        }
+    }
+    bool Destroy() {
+        if (!valid_) {
+            return false;
+        }
+        valid_ = false;
+        return dm_.DeleteDevice(name_);
+    }
+    const std::string& path() const { return path_; }
+    const std::string& name() const { return name_; }
+    bool valid() const { return valid_; }
+
+    Tempdevice(const Tempdevice&) = delete;
+    Tempdevice& operator=(const Tempdevice&) = delete;
+
+    Tempdevice& operator=(Tempdevice&& other) noexcept {
+        name_ = other.name_;
+        valid_ = other.valid_;
+        other.valid_ = false;
+        return *this;
+    }
+
+  private:
+    DeviceMapper& dm_;
+    std::string name_;
+    std::string path_;
+    bool valid_;
+};
+
+class SnapuserTest final {
+  public:
+    bool Setup();
+    bool SetupOrderedOps();
+    bool SetupOrderedOpsInverted();
+    bool SetupCopyOverlap_1();
+    bool SetupCopyOverlap_2();
+    bool Merge();
+    void ValidateMerge();
+    void ReadSnapshotDeviceAndValidate();
+    void Shutdown();
+    void MergeInterrupt();
+    void MergeInterruptFixed(int duration);
+    void MergeInterruptRandomly(int max_duration);
+    void StartMerge();
+    void CheckMergeCompletion();
+
+    static const uint64_t kSectorSize = 512;
+
+  private:
+    void SetupImpl();
+
+    void SimulateDaemonRestart();
+
+    void CreateCowDevice();
+    void CreateCowDeviceOrderedOps();
+    void CreateCowDeviceOrderedOpsInverted();
+    void CreateCowDeviceWithCopyOverlap_1();
+    void CreateCowDeviceWithCopyOverlap_2();
+    bool SetupDaemon();
+    void CreateBaseDevice();
+    void InitCowDevice();
+    void SetDeviceControlName();
+    void InitDaemon();
+    void CreateDmUserDevice();
+    void StartSnapuserdDaemon();
+
+    unique_ptr<LoopDevice> base_loop_;
+    unique_ptr<Tempdevice> dmuser_dev_;
+
+    std::string system_device_ctrl_name_;
+    std::string system_device_name_;
+
+    unique_fd base_fd_;
+    std::unique_ptr<TemporaryFile> cow_system_;
+    std::unique_ptr<SnapuserdClient> client_;
+    std::unique_ptr<uint8_t[]> orig_buffer_;
+    std::unique_ptr<uint8_t[]> merged_buffer_;
+    bool setup_ok_ = false;
+    bool merge_ok_ = false;
+    size_t size_ = 100_MiB;
+    int cow_num_sectors_;
+    int total_base_size_;
+};
+
+static unique_fd CreateTempFile(const std::string& name, size_t size) {
+    unique_fd fd(syscall(__NR_memfd_create, name.c_str(), MFD_ALLOW_SEALING));
+    if (fd < 0) {
+        return {};
+    }
+    if (size) {
+        if (ftruncate(fd, size) < 0) {
+            perror("ftruncate");
+            return {};
+        }
+        if (fcntl(fd, F_ADD_SEALS, F_SEAL_GROW | F_SEAL_SHRINK) < 0) {
+            perror("fcntl");
+            return {};
+        }
+    }
+    return fd;
+}
+
+void SnapuserTest::Shutdown() {
+    ASSERT_TRUE(dmuser_dev_->Destroy());
+
+    auto misc_device = "/dev/dm-user/" + system_device_ctrl_name_;
+    ASSERT_TRUE(client_->WaitForDeviceDelete(system_device_ctrl_name_));
+    ASSERT_TRUE(android::fs_mgr::WaitForFileDeleted(misc_device, 10s));
+    ASSERT_TRUE(client_->DetachSnapuserd());
+}
+
+bool SnapuserTest::Setup() {
+    SetupImpl();
+    return setup_ok_;
+}
+
+bool SnapuserTest::SetupOrderedOps() {
+    CreateBaseDevice();
+    CreateCowDeviceOrderedOps();
+    return SetupDaemon();
+}
+
+bool SnapuserTest::SetupOrderedOpsInverted() {
+    CreateBaseDevice();
+    CreateCowDeviceOrderedOpsInverted();
+    return SetupDaemon();
+}
+
+bool SnapuserTest::SetupCopyOverlap_1() {
+    CreateBaseDevice();
+    CreateCowDeviceWithCopyOverlap_1();
+    return SetupDaemon();
+}
+
+bool SnapuserTest::SetupCopyOverlap_2() {
+    CreateBaseDevice();
+    CreateCowDeviceWithCopyOverlap_2();
+    return SetupDaemon();
+}
+
+bool SnapuserTest::SetupDaemon() {
+    SetDeviceControlName();
+
+    StartSnapuserdDaemon();
+
+    CreateDmUserDevice();
+    InitCowDevice();
+    InitDaemon();
+
+    setup_ok_ = true;
+
+    return setup_ok_;
+}
+
+void SnapuserTest::StartSnapuserdDaemon() {
+    pid_t pid = fork();
+    ASSERT_GE(pid, 0);
+    if (pid == 0) {
+        std::string arg0 = "/system/bin/snapuserd";
+        std::string arg1 = "-socket="s + kSnapuserdSocketTest;
+        char* const argv[] = {arg0.data(), arg1.data(), nullptr};
+        ASSERT_GE(execv(arg0.c_str(), argv), 0);
+    } else {
+        client_ = SnapuserdClient::Connect(kSnapuserdSocketTest, 10s);
+        ASSERT_NE(client_, nullptr);
+    }
+}
+
+void SnapuserTest::CreateBaseDevice() {
+    unique_fd rnd_fd;
+
+    total_base_size_ = (size_ * 5);
+    base_fd_ = CreateTempFile("base_device", total_base_size_);
+    ASSERT_GE(base_fd_, 0);
+
+    rnd_fd.reset(open("/dev/random", O_RDONLY));
+    ASSERT_TRUE(rnd_fd > 0);
+
+    std::unique_ptr<uint8_t[]> random_buffer = std::make_unique<uint8_t[]>(1_MiB);
+
+    for (size_t j = 0; j < ((total_base_size_) / 1_MiB); j++) {
+        ASSERT_EQ(ReadFullyAtOffset(rnd_fd, (char*)random_buffer.get(), 1_MiB, 0), true);
+        ASSERT_EQ(android::base::WriteFully(base_fd_, random_buffer.get(), 1_MiB), true);
+    }
+
+    ASSERT_EQ(lseek(base_fd_, 0, SEEK_SET), 0);
+
+    base_loop_ = std::make_unique<LoopDevice>(base_fd_, 10s);
+    ASSERT_TRUE(base_loop_->valid());
+}
+
+void SnapuserTest::ReadSnapshotDeviceAndValidate() {
+    unique_fd fd(open(dmuser_dev_->path().c_str(), O_RDONLY));
+    ASSERT_GE(fd, 0);
+    std::unique_ptr<uint8_t[]> snapuserd_buffer = std::make_unique<uint8_t[]>(size_);
+
+    // COPY
+    loff_t offset = 0;
+    ASSERT_EQ(ReadFullyAtOffset(fd, snapuserd_buffer.get(), size_, offset), true);
+    ASSERT_EQ(memcmp(snapuserd_buffer.get(), orig_buffer_.get(), size_), 0);
+
+    // REPLACE
+    offset += size_;
+    ASSERT_EQ(ReadFullyAtOffset(fd, snapuserd_buffer.get(), size_, offset), true);
+    ASSERT_EQ(memcmp(snapuserd_buffer.get(), (char*)orig_buffer_.get() + size_, size_), 0);
+
+    // ZERO
+    offset += size_;
+    ASSERT_EQ(ReadFullyAtOffset(fd, snapuserd_buffer.get(), size_, offset), true);
+    ASSERT_EQ(memcmp(snapuserd_buffer.get(), (char*)orig_buffer_.get() + (size_ * 2), size_), 0);
+
+    // REPLACE
+    offset += size_;
+    ASSERT_EQ(ReadFullyAtOffset(fd, snapuserd_buffer.get(), size_, offset), true);
+    ASSERT_EQ(memcmp(snapuserd_buffer.get(), (char*)orig_buffer_.get() + (size_ * 3), size_), 0);
+
+    // XOR
+    offset += size_;
+    ASSERT_EQ(ReadFullyAtOffset(fd, snapuserd_buffer.get(), size_, offset), true);
+    ASSERT_EQ(memcmp(snapuserd_buffer.get(), (char*)orig_buffer_.get() + (size_ * 4), size_), 0);
+}
+
+void SnapuserTest::CreateCowDeviceWithCopyOverlap_2() {
+    std::string path = android::base::GetExecutableDirectory();
+    cow_system_ = std::make_unique<TemporaryFile>(path);
+
+    CowOptions options;
+    options.compression = "gz";
+    CowWriter writer(options);
+
+    ASSERT_TRUE(writer.Initialize(cow_system_->fd));
+
+    size_t num_blocks = size_ / options.block_size;
+    size_t x = num_blocks;
+    size_t blk_src_copy = 0;
+
+    // Create overlapping copy operations
+    while (1) {
+        ASSERT_TRUE(writer.AddCopy(blk_src_copy, blk_src_copy + 1));
+        x -= 1;
+        if (x == 1) {
+            break;
+        }
+        blk_src_copy += 1;
+    }
+
+    // Flush operations
+    ASSERT_TRUE(writer.Finalize());
+
+    // Construct the buffer required for validation
+    orig_buffer_ = std::make_unique<uint8_t[]>(total_base_size_);
+
+    // Read the entire base device
+    ASSERT_EQ(android::base::ReadFullyAtOffset(base_fd_, orig_buffer_.get(), total_base_size_, 0),
+              true);
+
+    // Merged operations required for validation
+    int block_size = 4096;
+    x = num_blocks;
+    loff_t src_offset = block_size;
+    loff_t dest_offset = 0;
+
+    while (1) {
+        memmove((char*)orig_buffer_.get() + dest_offset, (char*)orig_buffer_.get() + src_offset,
+                block_size);
+        x -= 1;
+        if (x == 1) {
+            break;
+        }
+        src_offset += block_size;
+        dest_offset += block_size;
+    }
+}
+
+void SnapuserTest::CreateCowDeviceWithCopyOverlap_1() {
+    std::string path = android::base::GetExecutableDirectory();
+    cow_system_ = std::make_unique<TemporaryFile>(path);
+
+    CowOptions options;
+    options.compression = "gz";
+    CowWriter writer(options);
+
+    ASSERT_TRUE(writer.Initialize(cow_system_->fd));
+
+    size_t num_blocks = size_ / options.block_size;
+    size_t x = num_blocks;
+    size_t blk_src_copy = num_blocks - 1;
+
+    // Create overlapping copy operations
+    while (1) {
+        ASSERT_TRUE(writer.AddCopy(blk_src_copy + 1, blk_src_copy));
+        x -= 1;
+        if (x == 0) {
+            ASSERT_EQ(blk_src_copy, 0);
+            break;
+        }
+        blk_src_copy -= 1;
+    }
+
+    // Flush operations
+    ASSERT_TRUE(writer.Finalize());
+
+    // Construct the buffer required for validation
+    orig_buffer_ = std::make_unique<uint8_t[]>(total_base_size_);
+
+    // Read the entire base device
+    ASSERT_EQ(android::base::ReadFullyAtOffset(base_fd_, orig_buffer_.get(), total_base_size_, 0),
+              true);
+
+    // Merged operations
+    ASSERT_EQ(android::base::ReadFullyAtOffset(base_fd_, orig_buffer_.get(), options.block_size, 0),
+              true);
+    ASSERT_EQ(android::base::ReadFullyAtOffset(
+                      base_fd_, (char*)orig_buffer_.get() + options.block_size, size_, 0),
+              true);
+}
+
+void SnapuserTest::CreateCowDeviceOrderedOpsInverted() {
+    unique_fd rnd_fd;
+    loff_t offset = 0;
+
+    std::string path = android::base::GetExecutableDirectory();
+    cow_system_ = std::make_unique<TemporaryFile>(path);
+
+    rnd_fd.reset(open("/dev/random", O_RDONLY));
+    ASSERT_TRUE(rnd_fd > 0);
+
+    std::unique_ptr<uint8_t[]> random_buffer_1_ = std::make_unique<uint8_t[]>(size_);
+
+    // Fill random data
+    for (size_t j = 0; j < (size_ / 1_MiB); j++) {
+        ASSERT_EQ(ReadFullyAtOffset(rnd_fd, (char*)random_buffer_1_.get() + offset, 1_MiB, 0),
+                  true);
+
+        offset += 1_MiB;
+    }
+
+    CowOptions options;
+    options.compression = "gz";
+    CowWriter writer(options);
+
+    ASSERT_TRUE(writer.Initialize(cow_system_->fd));
+
+    size_t num_blocks = size_ / options.block_size;
+    size_t blk_end_copy = num_blocks * 3;
+    size_t source_blk = num_blocks - 1;
+    size_t blk_src_copy = blk_end_copy - 1;
+    uint16_t xor_offset = 5;
+
+    size_t x = num_blocks;
+    while (1) {
+        ASSERT_TRUE(writer.AddCopy(source_blk, blk_src_copy));
+        x -= 1;
+        if (x == 0) {
+            break;
+        }
+        source_blk -= 1;
+        blk_src_copy -= 1;
+    }
+
+    for (size_t i = num_blocks; i > 0; i--) {
+        ASSERT_TRUE(writer.AddXorBlocks(num_blocks + i - 1,
+                                        &random_buffer_1_.get()[options.block_size * (i - 1)],
+                                        options.block_size, 2 * num_blocks + i - 1, xor_offset));
+    }
+    // Flush operations
+    ASSERT_TRUE(writer.Finalize());
+    // Construct the buffer required for validation
+    orig_buffer_ = std::make_unique<uint8_t[]>(total_base_size_);
+    // Read the entire base device
+    ASSERT_EQ(android::base::ReadFullyAtOffset(base_fd_, orig_buffer_.get(), total_base_size_, 0),
+              true);
+    // Merged Buffer
+    memmove(orig_buffer_.get(), (char*)orig_buffer_.get() + 2 * size_, size_);
+    memmove(orig_buffer_.get() + size_, (char*)orig_buffer_.get() + 2 * size_ + xor_offset, size_);
+    for (int i = 0; i < size_; i++) {
+        orig_buffer_.get()[size_ + i] ^= random_buffer_1_.get()[i];
+    }
+}
+
+void SnapuserTest::CreateCowDeviceOrderedOps() {
+    unique_fd rnd_fd;
+    loff_t offset = 0;
+
+    std::string path = android::base::GetExecutableDirectory();
+    cow_system_ = std::make_unique<TemporaryFile>(path);
+
+    rnd_fd.reset(open("/dev/random", O_RDONLY));
+    ASSERT_TRUE(rnd_fd > 0);
+
+    std::unique_ptr<uint8_t[]> random_buffer_1_ = std::make_unique<uint8_t[]>(size_);
+
+    // Fill random data
+    for (size_t j = 0; j < (size_ / 1_MiB); j++) {
+        ASSERT_EQ(ReadFullyAtOffset(rnd_fd, (char*)random_buffer_1_.get() + offset, 1_MiB, 0),
+                  true);
+
+        offset += 1_MiB;
+    }
+    memset(random_buffer_1_.get(), 0, size_);
+
+    CowOptions options;
+    options.compression = "gz";
+    CowWriter writer(options);
+
+    ASSERT_TRUE(writer.Initialize(cow_system_->fd));
+
+    size_t num_blocks = size_ / options.block_size;
+    size_t x = num_blocks;
+    size_t source_blk = 0;
+    size_t blk_src_copy = 2 * num_blocks;
+    uint16_t xor_offset = 5;
+
+    while (1) {
+        ASSERT_TRUE(writer.AddCopy(source_blk, blk_src_copy));
+
+        x -= 1;
+        if (x == 0) {
+            break;
+        }
+        source_blk += 1;
+        blk_src_copy += 1;
+    }
+
+    ASSERT_TRUE(writer.AddXorBlocks(num_blocks, random_buffer_1_.get(), size_, 2 * num_blocks,
+                                    xor_offset));
+    // Flush operations
+    ASSERT_TRUE(writer.Finalize());
+    // Construct the buffer required for validation
+    orig_buffer_ = std::make_unique<uint8_t[]>(total_base_size_);
+    // Read the entire base device
+    ASSERT_EQ(android::base::ReadFullyAtOffset(base_fd_, orig_buffer_.get(), total_base_size_, 0),
+              true);
+    // Merged Buffer
+    memmove(orig_buffer_.get(), (char*)orig_buffer_.get() + 2 * size_, size_);
+    memmove(orig_buffer_.get() + size_, (char*)orig_buffer_.get() + 2 * size_ + xor_offset, size_);
+    for (int i = 0; i < size_; i++) {
+        orig_buffer_.get()[size_ + i] ^= random_buffer_1_.get()[i];
+    }
+}
+
+void SnapuserTest::CreateCowDevice() {
+    unique_fd rnd_fd;
+    loff_t offset = 0;
+
+    std::string path = android::base::GetExecutableDirectory();
+    cow_system_ = std::make_unique<TemporaryFile>(path);
+
+    rnd_fd.reset(open("/dev/random", O_RDONLY));
+    ASSERT_TRUE(rnd_fd > 0);
+
+    std::unique_ptr<uint8_t[]> random_buffer_1_ = std::make_unique<uint8_t[]>(size_);
+
+    // Fill random data
+    for (size_t j = 0; j < (size_ / 1_MiB); j++) {
+        ASSERT_EQ(ReadFullyAtOffset(rnd_fd, (char*)random_buffer_1_.get() + offset, 1_MiB, 0),
+                  true);
+
+        offset += 1_MiB;
+    }
+
+    CowOptions options;
+    options.compression = "gz";
+    CowWriter writer(options);
+
+    ASSERT_TRUE(writer.Initialize(cow_system_->fd));
+
+    size_t num_blocks = size_ / options.block_size;
+    size_t blk_end_copy = num_blocks * 2;
+    size_t source_blk = num_blocks - 1;
+    size_t blk_src_copy = blk_end_copy - 1;
+
+    uint32_t sequence[num_blocks * 2];
+    // Sequence for Copy ops
+    for (int i = 0; i < num_blocks; i++) {
+        sequence[i] = num_blocks - 1 - i;
+    }
+    // Sequence for Xor ops
+    for (int i = 0; i < num_blocks; i++) {
+        sequence[num_blocks + i] = 5 * num_blocks - 1 - i;
+    }
+    ASSERT_TRUE(writer.AddSequenceData(2 * num_blocks, sequence));
+
+    size_t x = num_blocks;
+    while (1) {
+        ASSERT_TRUE(writer.AddCopy(source_blk, blk_src_copy));
+        x -= 1;
+        if (x == 0) {
+            break;
+        }
+        source_blk -= 1;
+        blk_src_copy -= 1;
+    }
+
+    source_blk = num_blocks;
+    blk_src_copy = blk_end_copy;
+
+    ASSERT_TRUE(writer.AddRawBlocks(source_blk, random_buffer_1_.get(), size_));
+
+    size_t blk_zero_copy_start = source_blk + num_blocks;
+    size_t blk_zero_copy_end = blk_zero_copy_start + num_blocks;
+
+    ASSERT_TRUE(writer.AddZeroBlocks(blk_zero_copy_start, num_blocks));
+
+    size_t blk_random2_replace_start = blk_zero_copy_end;
+
+    ASSERT_TRUE(writer.AddRawBlocks(blk_random2_replace_start, random_buffer_1_.get(), size_));
+
+    size_t blk_xor_start = blk_random2_replace_start + num_blocks;
+    size_t xor_offset = BLOCK_SZ / 2;
+    ASSERT_TRUE(writer.AddXorBlocks(blk_xor_start, random_buffer_1_.get(), size_, num_blocks,
+                                    xor_offset));
+
+    // Flush operations
+    ASSERT_TRUE(writer.Finalize());
+    // Construct the buffer required for validation
+    orig_buffer_ = std::make_unique<uint8_t[]>(total_base_size_);
+    std::string zero_buffer(size_, 0);
+    ASSERT_EQ(android::base::ReadFullyAtOffset(base_fd_, orig_buffer_.get(), size_, size_), true);
+    memcpy((char*)orig_buffer_.get() + size_, random_buffer_1_.get(), size_);
+    memcpy((char*)orig_buffer_.get() + (size_ * 2), (void*)zero_buffer.c_str(), size_);
+    memcpy((char*)orig_buffer_.get() + (size_ * 3), random_buffer_1_.get(), size_);
+    ASSERT_EQ(android::base::ReadFullyAtOffset(base_fd_, &orig_buffer_.get()[size_ * 4], size_,
+                                               size_ + xor_offset),
+              true);
+    for (int i = 0; i < size_; i++) {
+        orig_buffer_.get()[(size_ * 4) + i] =
+                (uint8_t)(orig_buffer_.get()[(size_ * 4) + i] ^ random_buffer_1_.get()[i]);
+    }
+}
+
+void SnapuserTest::InitCowDevice() {
+    uint64_t num_sectors = client_->InitDmUserCow(system_device_ctrl_name_, cow_system_->path,
+                                                  base_loop_->device(), base_loop_->device());
+    ASSERT_NE(num_sectors, 0);
+}
+
+void SnapuserTest::SetDeviceControlName() {
+    system_device_name_.clear();
+    system_device_ctrl_name_.clear();
+
+    std::string str(cow_system_->path);
+    std::size_t found = str.find_last_of("/\\");
+    ASSERT_NE(found, std::string::npos);
+    system_device_name_ = str.substr(found + 1);
+
+    system_device_ctrl_name_ = system_device_name_ + "-ctrl";
+}
+
+void SnapuserTest::CreateDmUserDevice() {
+    unique_fd fd(TEMP_FAILURE_RETRY(open(base_loop_->device().c_str(), O_RDONLY | O_CLOEXEC)));
+    ASSERT_TRUE(fd > 0);
+
+    uint64_t dev_sz = get_block_device_size(fd.get());
+    ASSERT_TRUE(dev_sz > 0);
+
+    cow_num_sectors_ = dev_sz >> 9;
+
+    DmTable dmuser_table;
+    ASSERT_TRUE(dmuser_table.AddTarget(
+            std::make_unique<DmTargetUser>(0, cow_num_sectors_, system_device_ctrl_name_)));
+    ASSERT_TRUE(dmuser_table.valid());
+
+    dmuser_dev_ = std::make_unique<Tempdevice>(system_device_name_, dmuser_table);
+    ASSERT_TRUE(dmuser_dev_->valid());
+    ASSERT_FALSE(dmuser_dev_->path().empty());
+
+    auto misc_device = "/dev/dm-user/" + system_device_ctrl_name_;
+    ASSERT_TRUE(android::fs_mgr::WaitForFile(misc_device, 10s));
+}
+
+void SnapuserTest::InitDaemon() {
+    bool ok = client_->AttachDmUser(system_device_ctrl_name_);
+    ASSERT_TRUE(ok);
+}
+
+void SnapuserTest::CheckMergeCompletion() {
+    while (true) {
+        double percentage = client_->GetMergePercent();
+        if ((int)percentage == 100) {
+            break;
+        }
+
+        std::this_thread::sleep_for(1s);
+    }
+}
+
+void SnapuserTest::SetupImpl() {
+    CreateBaseDevice();
+    CreateCowDevice();
+
+    SetDeviceControlName();
+
+    StartSnapuserdDaemon();
+
+    CreateDmUserDevice();
+    InitCowDevice();
+    InitDaemon();
+
+    setup_ok_ = true;
+}
+
+bool SnapuserTest::Merge() {
+    StartMerge();
+    CheckMergeCompletion();
+    merge_ok_ = true;
+    return merge_ok_;
+}
+
+void SnapuserTest::StartMerge() {
+    bool ok = client_->InitiateMerge(system_device_ctrl_name_);
+    ASSERT_TRUE(ok);
+}
+
+void SnapuserTest::ValidateMerge() {
+    merged_buffer_ = std::make_unique<uint8_t[]>(total_base_size_);
+    ASSERT_EQ(android::base::ReadFullyAtOffset(base_fd_, merged_buffer_.get(), total_base_size_, 0),
+              true);
+    ASSERT_EQ(memcmp(merged_buffer_.get(), orig_buffer_.get(), total_base_size_), 0);
+}
+
+void SnapuserTest::SimulateDaemonRestart() {
+    Shutdown();
+    std::this_thread::sleep_for(500ms);
+    SetDeviceControlName();
+    StartSnapuserdDaemon();
+    CreateDmUserDevice();
+    InitCowDevice();
+    InitDaemon();
+}
+
+void SnapuserTest::MergeInterruptRandomly(int max_duration) {
+    std::srand(std::time(nullptr));
+    StartMerge();
+
+    for (int i = 0; i < 20; i++) {
+        int duration = std::rand() % max_duration;
+        std::this_thread::sleep_for(std::chrono::milliseconds(duration));
+        SimulateDaemonRestart();
+        StartMerge();
+    }
+
+    SimulateDaemonRestart();
+    ASSERT_TRUE(Merge());
+}
+
+void SnapuserTest::MergeInterruptFixed(int duration) {
+    StartMerge();
+
+    for (int i = 0; i < 25; i++) {
+        std::this_thread::sleep_for(std::chrono::milliseconds(duration));
+        SimulateDaemonRestart();
+        StartMerge();
+    }
+
+    SimulateDaemonRestart();
+    ASSERT_TRUE(Merge());
+}
+
+void SnapuserTest::MergeInterrupt() {
+    // Interrupt merge at various intervals
+    StartMerge();
+    std::this_thread::sleep_for(250ms);
+    SimulateDaemonRestart();
+
+    StartMerge();
+    std::this_thread::sleep_for(250ms);
+    SimulateDaemonRestart();
+
+    StartMerge();
+    std::this_thread::sleep_for(150ms);
+    SimulateDaemonRestart();
+
+    StartMerge();
+    std::this_thread::sleep_for(100ms);
+    SimulateDaemonRestart();
+
+    StartMerge();
+    std::this_thread::sleep_for(800ms);
+    SimulateDaemonRestart();
+
+    StartMerge();
+    std::this_thread::sleep_for(600ms);
+    SimulateDaemonRestart();
+
+    ASSERT_TRUE(Merge());
+}
+
+TEST(Snapuserd_Test, Snapshot_IO_TEST) {
+    SnapuserTest harness;
+    ASSERT_TRUE(harness.Setup());
+    // I/O before merge
+    harness.ReadSnapshotDeviceAndValidate();
+    ASSERT_TRUE(harness.Merge());
+    harness.ValidateMerge();
+    // I/O after merge - daemon should read directly
+    // from base device
+    harness.ReadSnapshotDeviceAndValidate();
+    harness.Shutdown();
+}
+
+TEST(Snapuserd_Test, Snapshot_MERGE_IO_TEST) {
+    SnapuserTest harness;
+    ASSERT_TRUE(harness.Setup());
+    // Issue I/O before merge begins
+    std::async(std::launch::async, &SnapuserTest::ReadSnapshotDeviceAndValidate, &harness);
+    // Start the merge
+    ASSERT_TRUE(harness.Merge());
+    harness.ValidateMerge();
+    harness.Shutdown();
+}
+
+TEST(Snapuserd_Test, Snapshot_MERGE_IO_TEST_1) {
+    SnapuserTest harness;
+    ASSERT_TRUE(harness.Setup());
+    // Start the merge
+    harness.StartMerge();
+    // Issue I/O in parallel when merge is in-progress
+    std::async(std::launch::async, &SnapuserTest::ReadSnapshotDeviceAndValidate, &harness);
+    harness.CheckMergeCompletion();
+    harness.ValidateMerge();
+    harness.Shutdown();
+}
+
+TEST(Snapuserd_Test, Snapshot_Merge_Resume) {
+    SnapuserTest harness;
+    ASSERT_TRUE(harness.Setup());
+    harness.MergeInterrupt();
+    harness.ValidateMerge();
+    harness.Shutdown();
+}
+
+TEST(Snapuserd_Test, Snapshot_COPY_Overlap_TEST_1) {
+    SnapuserTest harness;
+    ASSERT_TRUE(harness.SetupCopyOverlap_1());
+    ASSERT_TRUE(harness.Merge());
+    harness.ValidateMerge();
+    harness.Shutdown();
+}
+
+TEST(Snapuserd_Test, Snapshot_COPY_Overlap_TEST_2) {
+    SnapuserTest harness;
+    ASSERT_TRUE(harness.SetupCopyOverlap_2());
+    ASSERT_TRUE(harness.Merge());
+    harness.ValidateMerge();
+    harness.Shutdown();
+}
+
+TEST(Snapuserd_Test, Snapshot_COPY_Overlap_Merge_Resume_TEST) {
+    SnapuserTest harness;
+    ASSERT_TRUE(harness.SetupCopyOverlap_1());
+    harness.MergeInterrupt();
+    harness.ValidateMerge();
+    harness.Shutdown();
+}
+
+TEST(Snapuserd_Test, Snapshot_Merge_Crash_Fixed_Ordered) {
+    SnapuserTest harness;
+    ASSERT_TRUE(harness.SetupOrderedOps());
+    harness.MergeInterruptFixed(300);
+    harness.ValidateMerge();
+    harness.Shutdown();
+}
+
+TEST(Snapuserd_Test, Snapshot_Merge_Crash_Random_Ordered) {
+    SnapuserTest harness;
+    ASSERT_TRUE(harness.SetupOrderedOps());
+    harness.MergeInterruptRandomly(500);
+    harness.ValidateMerge();
+    harness.Shutdown();
+}
+
+TEST(Snapuserd_Test, Snapshot_Merge_Crash_Fixed_Inverted) {
+    SnapuserTest harness;
+    ASSERT_TRUE(harness.SetupOrderedOpsInverted());
+    harness.MergeInterruptFixed(50);
+    harness.ValidateMerge();
+    harness.Shutdown();
+}
+
+TEST(Snapuserd_Test, Snapshot_Merge_Crash_Random_Inverted) {
+    SnapuserTest harness;
+    ASSERT_TRUE(harness.SetupOrderedOpsInverted());
+    harness.MergeInterruptRandomly(50);
+    harness.ValidateMerge();
+    harness.Shutdown();
+}
+
+}  // namespace snapshot
+}  // namespace android
+
+int main(int argc, char** argv) {
+    ::testing::InitGoogleTest(&argc, argv);
+    return RUN_ALL_TESTS();
+}
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 6c91fde6b..6dec1e2e0 100644
--- a/fs_mgr/libsnapshot/snapuserd/user-space-merge/snapuserd_transitions.cpp
+++ b/fs_mgr/libsnapshot/snapuserd/user-space-merge/snapuserd_transitions.cpp
@@ -570,7 +570,6 @@ void SnapshotHandler::NotifyIOCompletion(uint64_t new_block) {
     {
         std::unique_lock<std::mutex> lock(blk_state->m_lock);
 
-        CHECK(blk_state->merge_state_ == MERGE_GROUP_STATE::GROUP_MERGE_PENDING);
         blk_state->num_ios_in_progress -= 1;
         if (blk_state->num_ios_in_progress == 0) {
             pending_ios = false;
diff --git a/fs_mgr/libsnapshot/userspace_snapshot_test.cpp b/fs_mgr/libsnapshot/userspace_snapshot_test.cpp
new file mode 100644
index 000000000..abe67f64c
--- /dev/null
+++ b/fs_mgr/libsnapshot/userspace_snapshot_test.cpp
@@ -0,0 +1,2519 @@
+// Copyright (C) 2018 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.
+
+#include <libsnapshot/cow_format.h>
+#include <libsnapshot/snapshot.h>
+
+#include <fcntl.h>
+#include <signal.h>
+#include <sys/file.h>
+#include <sys/stat.h>
+#include <sys/types.h>
+
+#include <chrono>
+#include <deque>
+#include <future>
+#include <iostream>
+
+#include <android-base/file.h>
+#include <android-base/logging.h>
+#include <android-base/properties.h>
+#include <android-base/strings.h>
+#include <android-base/unique_fd.h>
+#include <fs_mgr/file_wait.h>
+#include <fs_mgr/roots.h>
+#include <fs_mgr_dm_linear.h>
+#include <gtest/gtest.h>
+#include <libdm/dm.h>
+#include <libfiemap/image_manager.h>
+#include <liblp/builder.h>
+#include <storage_literals/storage_literals.h>
+
+#include <android/snapshot/snapshot.pb.h>
+#include <libsnapshot/test_helpers.h>
+#include "partition_cow_creator.h"
+#include "utility.h"
+
+#include <android-base/properties.h>
+
+// Mock classes are not used. Header included to ensure mocked class definition aligns with the
+// class itself.
+#include <libsnapshot/mock_device_info.h>
+#include <libsnapshot/mock_snapshot.h>
+
+namespace android {
+namespace snapshot {
+
+using android::base::unique_fd;
+using android::dm::DeviceMapper;
+using android::dm::DmDeviceState;
+using android::dm::IDeviceMapper;
+using android::fiemap::FiemapStatus;
+using android::fiemap::IImageManager;
+using android::fs_mgr::BlockDeviceInfo;
+using android::fs_mgr::CreateLogicalPartitionParams;
+using android::fs_mgr::DestroyLogicalPartition;
+using android::fs_mgr::EnsurePathMounted;
+using android::fs_mgr::EnsurePathUnmounted;
+using android::fs_mgr::Extent;
+using android::fs_mgr::Fstab;
+using android::fs_mgr::GetPartitionGroupName;
+using android::fs_mgr::GetPartitionName;
+using android::fs_mgr::Interval;
+using android::fs_mgr::MetadataBuilder;
+using android::fs_mgr::SlotSuffixForSlotNumber;
+using chromeos_update_engine::DeltaArchiveManifest;
+using chromeos_update_engine::DynamicPartitionGroup;
+using chromeos_update_engine::PartitionUpdate;
+using namespace ::testing;
+using namespace android::storage_literals;
+using namespace std::chrono_literals;
+using namespace std::string_literals;
+
+// Global states. See test_helpers.h.
+std::unique_ptr<SnapshotManager> sm;
+TestDeviceInfo* test_device = nullptr;
+std::string fake_super;
+
+void MountMetadata();
+
+class SnapshotTest : public ::testing::Test {
+  public:
+    SnapshotTest() : dm_(DeviceMapper::Instance()) {}
+
+    // This is exposed for main.
+    void Cleanup() {
+        InitializeState();
+        CleanupTestArtifacts();
+    }
+
+  protected:
+    void SetUp() override {
+        SKIP_IF_NON_VIRTUAL_AB();
+
+        SnapshotTestPropertyFetcher::SetUp();
+        InitializeState();
+        CleanupTestArtifacts();
+        FormatFakeSuper();
+        MountMetadata();
+        ASSERT_TRUE(sm->BeginUpdate());
+    }
+
+    void TearDown() override {
+        RETURN_IF_NON_VIRTUAL_AB();
+
+        lock_ = nullptr;
+
+        CleanupTestArtifacts();
+        SnapshotTestPropertyFetcher::TearDown();
+    }
+
+    void InitializeState() {
+        ASSERT_TRUE(sm->EnsureImageManager());
+        image_manager_ = sm->image_manager();
+
+        test_device->set_slot_suffix("_a");
+
+        sm->set_use_first_stage_snapuserd(false);
+    }
+
+    void CleanupTestArtifacts() {
+        // Normally cancelling inside a merge is not allowed. Since these
+        // are tests, we don't care, destroy everything that might exist.
+        // Note we hardcode this list because of an annoying quirk: when
+        // completing a merge, the snapshot stops existing, so we can't
+        // get an accurate list to remove.
+        lock_ = nullptr;
+
+        std::vector<std::string> snapshots = {"test-snapshot", "test_partition_a",
+                                              "test_partition_b"};
+        for (const auto& snapshot : snapshots) {
+            ASSERT_TRUE(DeleteSnapshotDevice(snapshot));
+            DeleteBackingImage(image_manager_, snapshot + "-cow-img");
+
+            auto status_file = sm->GetSnapshotStatusFilePath(snapshot);
+            android::base::RemoveFileIfExists(status_file);
+        }
+
+        // Remove stale partitions in fake super.
+        std::vector<std::string> partitions = {
+                "base-device",
+                "test_partition_b",
+                "test_partition_b-base",
+                "test_partition_b-base",
+        };
+        for (const auto& partition : partitions) {
+            DeleteDevice(partition);
+        }
+
+        if (sm->GetUpdateState() != UpdateState::None) {
+            auto state_file = sm->GetStateFilePath();
+            unlink(state_file.c_str());
+        }
+    }
+
+    bool AcquireLock() {
+        lock_ = sm->LockExclusive();
+        return !!lock_;
+    }
+
+    // This is so main() can instantiate this to invoke Cleanup.
+    virtual void TestBody() override {}
+
+    void FormatFakeSuper() {
+        BlockDeviceInfo super_device("super", kSuperSize, 0, 0, 4096);
+        std::vector<BlockDeviceInfo> devices = {super_device};
+
+        auto builder = MetadataBuilder::New(devices, "super", 65536, 2);
+        ASSERT_NE(builder, nullptr);
+
+        auto metadata = builder->Export();
+        ASSERT_NE(metadata, nullptr);
+
+        TestPartitionOpener opener(fake_super);
+        ASSERT_TRUE(FlashPartitionTable(opener, fake_super, *metadata.get()));
+    }
+
+    // If |path| is non-null, the partition will be mapped after creation.
+    bool CreatePartition(const std::string& name, uint64_t size, std::string* path = nullptr,
+                         const std::optional<std::string> group = {}) {
+        TestPartitionOpener opener(fake_super);
+        auto builder = MetadataBuilder::New(opener, "super", 0);
+        if (!builder) return false;
+
+        std::string partition_group = std::string(android::fs_mgr::kDefaultGroup);
+        if (group) {
+            partition_group = *group;
+        }
+        return CreatePartition(builder.get(), name, size, path, partition_group);
+    }
+
+    bool CreatePartition(MetadataBuilder* builder, const std::string& name, uint64_t size,
+                         std::string* path, const std::string& group) {
+        auto partition = builder->AddPartition(name, group, 0);
+        if (!partition) return false;
+        if (!builder->ResizePartition(partition, size)) {
+            return false;
+        }
+
+        // Update the source slot.
+        auto metadata = builder->Export();
+        if (!metadata) return false;
+
+        TestPartitionOpener opener(fake_super);
+        if (!UpdatePartitionTable(opener, "super", *metadata.get(), 0)) {
+            return false;
+        }
+
+        if (!path) return true;
+
+        CreateLogicalPartitionParams params = {
+                .block_device = fake_super,
+                .metadata = metadata.get(),
+                .partition_name = name,
+                .force_writable = true,
+                .timeout_ms = 10s,
+        };
+        return CreateLogicalPartition(params, path);
+    }
+
+    AssertionResult MapUpdateSnapshot(const std::string& name,
+                                      std::unique_ptr<ISnapshotWriter>* writer) {
+        TestPartitionOpener opener(fake_super);
+        CreateLogicalPartitionParams params{
+                .block_device = fake_super,
+                .metadata_slot = 1,
+                .partition_name = name,
+                .timeout_ms = 10s,
+                .partition_opener = &opener,
+        };
+
+        auto old_partition = "/dev/block/mapper/" + GetOtherPartitionName(name);
+        auto result = sm->OpenSnapshotWriter(params, {old_partition});
+        if (!result) {
+            return AssertionFailure() << "Cannot open snapshot for writing: " << name;
+        }
+        if (!result->Initialize()) {
+            return AssertionFailure() << "Cannot initialize snapshot for writing: " << name;
+        }
+
+        if (writer) {
+            *writer = std::move(result);
+        }
+        return AssertionSuccess();
+    }
+
+    AssertionResult MapUpdateSnapshot(const std::string& name, std::string* path) {
+        TestPartitionOpener opener(fake_super);
+        CreateLogicalPartitionParams params{
+                .block_device = fake_super,
+                .metadata_slot = 1,
+                .partition_name = name,
+                .timeout_ms = 10s,
+                .partition_opener = &opener,
+        };
+
+        auto result = sm->MapUpdateSnapshot(params, path);
+        if (!result) {
+            return AssertionFailure() << "Cannot open snapshot for writing: " << name;
+        }
+        return AssertionSuccess();
+    }
+
+    AssertionResult DeleteSnapshotDevice(const std::string& snapshot) {
+        AssertionResult res = AssertionSuccess();
+        if (!(res = DeleteDevice(snapshot))) return res;
+        if (!sm->UnmapDmUserDevice(snapshot + "-user-cow")) {
+            return AssertionFailure() << "Cannot delete dm-user device for " << snapshot;
+        }
+        if (!(res = DeleteDevice(snapshot + "-inner"))) return res;
+        if (!(res = DeleteDevice(snapshot + "-cow"))) return res;
+        if (!image_manager_->UnmapImageIfExists(snapshot + "-cow-img")) {
+            return AssertionFailure() << "Cannot unmap image " << snapshot << "-cow-img";
+        }
+        if (!(res = DeleteDevice(snapshot + "-base"))) return res;
+        if (!(res = DeleteDevice(snapshot + "-src"))) return res;
+        return AssertionSuccess();
+    }
+
+    AssertionResult DeleteDevice(const std::string& device) {
+        if (!dm_.DeleteDeviceIfExists(device)) {
+            return AssertionFailure() << "Can't delete " << device;
+        }
+        return AssertionSuccess();
+    }
+
+    AssertionResult CreateCowImage(const std::string& name) {
+        if (!sm->CreateCowImage(lock_.get(), name)) {
+            return AssertionFailure() << "Cannot create COW image " << name;
+        }
+        std::string cow_device;
+        auto map_res = MapCowImage(name, 10s, &cow_device);
+        if (!map_res) {
+            return map_res;
+        }
+        if (!InitializeKernelCow(cow_device)) {
+            return AssertionFailure() << "Cannot zero fill " << cow_device;
+        }
+        if (!sm->UnmapCowImage(name)) {
+            return AssertionFailure() << "Cannot unmap " << name << " after zero filling it";
+        }
+        return AssertionSuccess();
+    }
+
+    AssertionResult MapCowImage(const std::string& name,
+                                const std::chrono::milliseconds& timeout_ms, std::string* path) {
+        auto cow_image_path = sm->MapCowImage(name, timeout_ms);
+        if (!cow_image_path.has_value()) {
+            return AssertionFailure() << "Cannot map cow image " << name;
+        }
+        *path = *cow_image_path;
+        return AssertionSuccess();
+    }
+
+    // Prepare A/B slot for a partition named "test_partition".
+    AssertionResult PrepareOneSnapshot(uint64_t device_size,
+                                       std::unique_ptr<ISnapshotWriter>* writer = nullptr) {
+        lock_ = nullptr;
+
+        DeltaArchiveManifest manifest;
+
+        auto dynamic_partition_metadata = manifest.mutable_dynamic_partition_metadata();
+        dynamic_partition_metadata->set_vabc_enabled(IsCompressionEnabled());
+        dynamic_partition_metadata->set_cow_version(android::snapshot::kCowVersionMajor);
+
+        auto group = dynamic_partition_metadata->add_groups();
+        group->set_name("group");
+        group->set_size(device_size * 2);
+        group->add_partition_names("test_partition");
+
+        auto pu = manifest.add_partitions();
+        pu->set_partition_name("test_partition");
+        pu->set_estimate_cow_size(device_size);
+        SetSize(pu, device_size);
+
+        auto extent = pu->add_operations()->add_dst_extents();
+        extent->set_start_block(0);
+        if (device_size) {
+            extent->set_num_blocks(device_size / manifest.block_size());
+        }
+
+        TestPartitionOpener opener(fake_super);
+        auto builder = MetadataBuilder::New(opener, "super", 0);
+        if (!builder) {
+            return AssertionFailure() << "Failed to open MetadataBuilder";
+        }
+        builder->AddGroup("group_a", 16_GiB);
+        builder->AddGroup("group_b", 16_GiB);
+        if (!CreatePartition(builder.get(), "test_partition_a", device_size, nullptr, "group_a")) {
+            return AssertionFailure() << "Failed create test_partition_a";
+        }
+
+        if (!sm->CreateUpdateSnapshots(manifest)) {
+            return AssertionFailure() << "Failed to create update snapshots";
+        }
+
+        if (writer) {
+            auto res = MapUpdateSnapshot("test_partition_b", writer);
+            if (!res) {
+                return res;
+            }
+        } else if (!IsCompressionEnabled()) {
+            std::string ignore;
+            if (!MapUpdateSnapshot("test_partition_b", &ignore)) {
+                return AssertionFailure() << "Failed to map test_partition_b";
+            }
+        }
+        if (!AcquireLock()) {
+            return AssertionFailure() << "Failed to acquire lock";
+        }
+        return AssertionSuccess();
+    }
+
+    // Simulate a reboot into the new slot.
+    AssertionResult SimulateReboot() {
+        lock_ = nullptr;
+        if (!sm->FinishedSnapshotWrites(false)) {
+            return AssertionFailure() << "Failed to finish snapshot writes";
+        }
+        if (!sm->UnmapUpdateSnapshot("test_partition_b")) {
+            return AssertionFailure() << "Failed to unmap COW for test_partition_b";
+        }
+        if (!dm_.DeleteDeviceIfExists("test_partition_b")) {
+            return AssertionFailure() << "Failed to delete test_partition_b";
+        }
+        if (!dm_.DeleteDeviceIfExists("test_partition_b-base")) {
+            return AssertionFailure() << "Failed to destroy test_partition_b-base";
+        }
+        return AssertionSuccess();
+    }
+
+    std::unique_ptr<SnapshotManager> NewManagerForFirstStageMount(
+            const std::string& slot_suffix = "_a") {
+        auto info = new TestDeviceInfo(fake_super, slot_suffix);
+        return NewManagerForFirstStageMount(info);
+    }
+
+    std::unique_ptr<SnapshotManager> NewManagerForFirstStageMount(TestDeviceInfo* info) {
+        info->set_first_stage_init(true);
+        auto init = SnapshotManager::NewForFirstStageMount(info);
+        if (!init) {
+            return nullptr;
+        }
+        init->SetUeventRegenCallback([](const std::string& device) -> bool {
+            return android::fs_mgr::WaitForFile(device, snapshot_timeout_);
+        });
+        return init;
+    }
+
+    static constexpr std::chrono::milliseconds snapshot_timeout_ = 5s;
+    DeviceMapper& dm_;
+    std::unique_ptr<SnapshotManager::LockedFile> lock_;
+    android::fiemap::IImageManager* image_manager_ = nullptr;
+    std::string fake_super_;
+};
+
+TEST_F(SnapshotTest, CreateSnapshot) {
+    ASSERT_TRUE(AcquireLock());
+
+    PartitionCowCreator cow_creator;
+    cow_creator.compression_enabled = IsCompressionEnabled();
+    if (cow_creator.compression_enabled) {
+        cow_creator.compression_algorithm = "gz";
+    } else {
+        cow_creator.compression_algorithm = "none";
+    }
+
+    static const uint64_t kDeviceSize = 1024 * 1024;
+    SnapshotStatus status;
+    status.set_name("test-snapshot");
+    status.set_device_size(kDeviceSize);
+    status.set_snapshot_size(kDeviceSize);
+    status.set_cow_file_size(kDeviceSize);
+    ASSERT_TRUE(sm->CreateSnapshot(lock_.get(), &cow_creator, &status));
+    ASSERT_TRUE(CreateCowImage("test-snapshot"));
+
+    std::vector<std::string> snapshots;
+    ASSERT_TRUE(sm->ListSnapshots(lock_.get(), &snapshots));
+    ASSERT_EQ(snapshots.size(), 1);
+    ASSERT_EQ(snapshots[0], "test-snapshot");
+
+    // Scope so delete can re-acquire the snapshot file lock.
+    {
+        SnapshotStatus status;
+        ASSERT_TRUE(sm->ReadSnapshotStatus(lock_.get(), "test-snapshot", &status));
+        ASSERT_EQ(status.state(), SnapshotState::CREATED);
+        ASSERT_EQ(status.device_size(), kDeviceSize);
+        ASSERT_EQ(status.snapshot_size(), kDeviceSize);
+        ASSERT_EQ(status.compression_enabled(), cow_creator.compression_enabled);
+        ASSERT_EQ(status.compression_algorithm(), cow_creator.compression_algorithm);
+    }
+
+    ASSERT_TRUE(sm->UnmapSnapshot(lock_.get(), "test-snapshot"));
+    ASSERT_TRUE(sm->UnmapCowImage("test-snapshot"));
+    ASSERT_TRUE(sm->DeleteSnapshot(lock_.get(), "test-snapshot"));
+}
+
+TEST_F(SnapshotTest, MapSnapshot) {
+    ASSERT_TRUE(AcquireLock());
+
+    PartitionCowCreator cow_creator;
+    cow_creator.compression_enabled = IsCompressionEnabled();
+
+    static const uint64_t kDeviceSize = 1024 * 1024;
+    SnapshotStatus status;
+    status.set_name("test-snapshot");
+    status.set_device_size(kDeviceSize);
+    status.set_snapshot_size(kDeviceSize);
+    status.set_cow_file_size(kDeviceSize);
+    ASSERT_TRUE(sm->CreateSnapshot(lock_.get(), &cow_creator, &status));
+    ASSERT_TRUE(CreateCowImage("test-snapshot"));
+
+    std::string base_device;
+    ASSERT_TRUE(CreatePartition("base-device", kDeviceSize, &base_device));
+
+    std::string cow_device;
+    ASSERT_TRUE(MapCowImage("test-snapshot", 10s, &cow_device));
+
+    std::string snap_device;
+    ASSERT_TRUE(sm->MapSnapshot(lock_.get(), "test-snapshot", base_device, cow_device, 10s,
+                                &snap_device));
+    ASSERT_TRUE(android::base::StartsWith(snap_device, "/dev/block/dm-"));
+}
+
+TEST_F(SnapshotTest, NoMergeBeforeReboot) {
+    ASSERT_TRUE(sm->FinishedSnapshotWrites(false));
+
+    // Merge should fail, since the slot hasn't changed.
+    ASSERT_FALSE(sm->InitiateMerge());
+}
+
+TEST_F(SnapshotTest, CleanFirstStageMount) {
+    // If there's no update in progress, there should be no first-stage mount
+    // needed.
+    auto sm = NewManagerForFirstStageMount();
+    ASSERT_NE(sm, nullptr);
+    ASSERT_FALSE(sm->NeedSnapshotsInFirstStageMount());
+}
+
+TEST_F(SnapshotTest, FirstStageMountAfterRollback) {
+    ASSERT_TRUE(sm->FinishedSnapshotWrites(false));
+
+    // We didn't change the slot, so we shouldn't need snapshots.
+    auto sm = NewManagerForFirstStageMount();
+    ASSERT_NE(sm, nullptr);
+    ASSERT_FALSE(sm->NeedSnapshotsInFirstStageMount());
+
+    auto indicator = sm->GetRollbackIndicatorPath();
+    ASSERT_EQ(access(indicator.c_str(), R_OK), 0);
+}
+
+TEST_F(SnapshotTest, Merge) {
+    ASSERT_TRUE(AcquireLock());
+
+    static const uint64_t kDeviceSize = 1024 * 1024;
+
+    std::unique_ptr<ISnapshotWriter> writer;
+    ASSERT_TRUE(PrepareOneSnapshot(kDeviceSize, &writer));
+
+    // Release the lock.
+    lock_ = nullptr;
+
+    std::string test_string = "This is a test string.";
+    test_string.resize(writer->options().block_size);
+    ASSERT_TRUE(writer->AddRawBlocks(0, test_string.data(), test_string.size()));
+    ASSERT_TRUE(writer->Finalize());
+    writer = nullptr;
+
+    // Done updating.
+    ASSERT_TRUE(sm->FinishedSnapshotWrites(false));
+
+    ASSERT_TRUE(sm->UnmapUpdateSnapshot("test_partition_b"));
+
+    test_device->set_slot_suffix("_b");
+    ASSERT_TRUE(sm->CreateLogicalAndSnapshotPartitions("super", snapshot_timeout_));
+    ASSERT_TRUE(sm->InitiateMerge());
+
+    // The device should have been switched to a snapshot-merge target.
+    DeviceMapper::TargetInfo target;
+    ASSERT_TRUE(sm->IsSnapshotDevice("test_partition_b", &target));
+    ASSERT_EQ(DeviceMapper::GetTargetType(target.spec), "user");
+
+    // We should not be able to cancel an update now.
+    ASSERT_FALSE(sm->CancelUpdate());
+
+    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_partition_b"));
+
+    // 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("/dev/block/mapper/test_partition_b", O_RDONLY | O_CLOEXEC));
+    ASSERT_GE(fd, 0);
+
+    std::string buffer(test_string.size(), '\0');
+    ASSERT_TRUE(android::base::ReadFully(fd, buffer.data(), buffer.size()));
+    ASSERT_EQ(test_string, buffer);
+}
+
+TEST_F(SnapshotTest, FirstStageMountAndMerge) {
+    ASSERT_TRUE(AcquireLock());
+
+    static const uint64_t kDeviceSize = 1024 * 1024;
+    ASSERT_TRUE(PrepareOneSnapshot(kDeviceSize));
+    ASSERT_TRUE(SimulateReboot());
+
+    auto init = NewManagerForFirstStageMount("_b");
+    ASSERT_NE(init, nullptr);
+    ASSERT_TRUE(init->NeedSnapshotsInFirstStageMount());
+    ASSERT_TRUE(init->CreateLogicalAndSnapshotPartitions("super", snapshot_timeout_));
+
+    ASSERT_TRUE(AcquireLock());
+
+    // Validate that we have a snapshot device.
+    SnapshotStatus status;
+    ASSERT_TRUE(init->ReadSnapshotStatus(lock_.get(), "test_partition_b", &status));
+    ASSERT_EQ(status.state(), SnapshotState::CREATED);
+    if (IsCompressionEnabled()) {
+        ASSERT_EQ(status.compression_algorithm(), "gz");
+    } else {
+        ASSERT_EQ(status.compression_algorithm(), "none");
+    }
+
+    DeviceMapper::TargetInfo target;
+    ASSERT_TRUE(init->IsSnapshotDevice("test_partition_b", &target));
+    ASSERT_EQ(DeviceMapper::GetTargetType(target.spec), "user");
+}
+
+TEST_F(SnapshotTest, FlashSuperDuringUpdate) {
+    ASSERT_TRUE(AcquireLock());
+
+    static const uint64_t kDeviceSize = 1024 * 1024;
+    ASSERT_TRUE(PrepareOneSnapshot(kDeviceSize));
+    ASSERT_TRUE(SimulateReboot());
+
+    // Reflash the super partition.
+    FormatFakeSuper();
+    ASSERT_TRUE(CreatePartition("test_partition_b", kDeviceSize));
+
+    auto init = NewManagerForFirstStageMount("_b");
+    ASSERT_NE(init, nullptr);
+    ASSERT_TRUE(init->NeedSnapshotsInFirstStageMount());
+    ASSERT_TRUE(init->CreateLogicalAndSnapshotPartitions("super", snapshot_timeout_));
+
+    ASSERT_TRUE(AcquireLock());
+
+    SnapshotStatus status;
+    ASSERT_TRUE(init->ReadSnapshotStatus(lock_.get(), "test_partition_b", &status));
+
+    // We should not get a snapshot device now.
+    DeviceMapper::TargetInfo target;
+    ASSERT_FALSE(init->IsSnapshotDevice("test_partition_b", &target));
+
+    // We should see a cancelled update as well.
+    lock_ = nullptr;
+    ASSERT_EQ(sm->ProcessUpdateState(), UpdateState::Cancelled);
+}
+
+TEST_F(SnapshotTest, FlashSuperDuringMerge) {
+    ASSERT_TRUE(AcquireLock());
+
+    static const uint64_t kDeviceSize = 1024 * 1024;
+    ASSERT_TRUE(PrepareOneSnapshot(kDeviceSize));
+    ASSERT_TRUE(SimulateReboot());
+
+    auto init = NewManagerForFirstStageMount("_b");
+    ASSERT_NE(init, nullptr);
+    ASSERT_TRUE(init->NeedSnapshotsInFirstStageMount());
+    ASSERT_TRUE(init->CreateLogicalAndSnapshotPartitions("super", snapshot_timeout_));
+    ASSERT_TRUE(init->InitiateMerge());
+
+    // Now, reflash super. Note that we haven't called ProcessUpdateState, so the
+    // status is still Merging.
+    ASSERT_TRUE(DeleteSnapshotDevice("test_partition_b"));
+    ASSERT_TRUE(init->image_manager()->UnmapImageIfExists("test_partition_b-cow-img"));
+    FormatFakeSuper();
+    ASSERT_TRUE(CreatePartition("test_partition_b", kDeviceSize));
+    ASSERT_TRUE(init->NeedSnapshotsInFirstStageMount());
+    ASSERT_TRUE(init->CreateLogicalAndSnapshotPartitions("super", snapshot_timeout_));
+
+    // Because the status is Merging, we must call ProcessUpdateState, which should
+    // detect a cancelled update.
+    ASSERT_EQ(init->ProcessUpdateState(), UpdateState::Cancelled);
+    ASSERT_EQ(init->GetUpdateState(), UpdateState::None);
+}
+
+TEST_F(SnapshotTest, UpdateBootControlHal) {
+    ASSERT_TRUE(AcquireLock());
+
+    ASSERT_TRUE(sm->WriteUpdateState(lock_.get(), UpdateState::None));
+    ASSERT_EQ(test_device->merge_status(), MergeStatus::NONE);
+
+    ASSERT_TRUE(sm->WriteUpdateState(lock_.get(), UpdateState::Initiated));
+    ASSERT_EQ(test_device->merge_status(), MergeStatus::NONE);
+
+    ASSERT_TRUE(sm->WriteUpdateState(lock_.get(), UpdateState::Unverified));
+    ASSERT_EQ(test_device->merge_status(), MergeStatus::SNAPSHOTTED);
+
+    ASSERT_TRUE(sm->WriteUpdateState(lock_.get(), UpdateState::Merging));
+    ASSERT_EQ(test_device->merge_status(), MergeStatus::MERGING);
+
+    ASSERT_TRUE(sm->WriteUpdateState(lock_.get(), UpdateState::MergeNeedsReboot));
+    ASSERT_EQ(test_device->merge_status(), MergeStatus::NONE);
+
+    ASSERT_TRUE(sm->WriteUpdateState(lock_.get(), UpdateState::MergeCompleted));
+    ASSERT_EQ(test_device->merge_status(), MergeStatus::NONE);
+
+    ASSERT_TRUE(sm->WriteUpdateState(lock_.get(), UpdateState::MergeFailed));
+    ASSERT_EQ(test_device->merge_status(), MergeStatus::MERGING);
+}
+
+TEST_F(SnapshotTest, MergeFailureCode) {
+    ASSERT_TRUE(AcquireLock());
+
+    ASSERT_TRUE(sm->WriteUpdateState(lock_.get(), UpdateState::MergeFailed,
+                                     MergeFailureCode::ListSnapshots));
+    ASSERT_EQ(test_device->merge_status(), MergeStatus::MERGING);
+
+    SnapshotUpdateStatus status = sm->ReadSnapshotUpdateStatus(lock_.get());
+    ASSERT_EQ(status.state(), UpdateState::MergeFailed);
+    ASSERT_EQ(status.merge_failure_code(), MergeFailureCode::ListSnapshots);
+}
+
+enum class Request { UNKNOWN, LOCK_SHARED, LOCK_EXCLUSIVE, UNLOCK, EXIT };
+std::ostream& operator<<(std::ostream& os, Request request) {
+    switch (request) {
+        case Request::LOCK_SHARED:
+            return os << "Shared";
+        case Request::LOCK_EXCLUSIVE:
+            return os << "Exclusive";
+        case Request::UNLOCK:
+            return os << "Unlock";
+        case Request::EXIT:
+            return os << "Exit";
+        case Request::UNKNOWN:
+            [[fallthrough]];
+        default:
+            return os << "Unknown";
+    }
+}
+
+class LockTestConsumer {
+  public:
+    AssertionResult MakeRequest(Request new_request) {
+        {
+            std::unique_lock<std::mutex> ulock(mutex_);
+            requests_.push_back(new_request);
+        }
+        cv_.notify_all();
+        return AssertionSuccess() << "Request " << new_request << " successful";
+    }
+
+    template <typename R, typename P>
+    AssertionResult WaitFulfill(std::chrono::duration<R, P> timeout) {
+        std::unique_lock<std::mutex> ulock(mutex_);
+        if (cv_.wait_for(ulock, timeout, [this] { return requests_.empty(); })) {
+            return AssertionSuccess() << "All requests_ fulfilled.";
+        }
+        return AssertionFailure() << "Timeout waiting for fulfilling " << requests_.size()
+                                  << " request(s), first one is "
+                                  << (requests_.empty() ? Request::UNKNOWN : requests_.front());
+    }
+
+    void StartHandleRequestsInBackground() {
+        future_ = std::async(std::launch::async, &LockTestConsumer::HandleRequests, this);
+    }
+
+  private:
+    void HandleRequests() {
+        static constexpr auto consumer_timeout = 3s;
+
+        auto next_request = Request::UNKNOWN;
+        do {
+            // Peek next request.
+            {
+                std::unique_lock<std::mutex> ulock(mutex_);
+                if (cv_.wait_for(ulock, consumer_timeout, [this] { return !requests_.empty(); })) {
+                    next_request = requests_.front();
+                } else {
+                    next_request = Request::EXIT;
+                }
+            }
+
+            // Handle next request.
+            switch (next_request) {
+                case Request::LOCK_SHARED: {
+                    lock_ = sm->LockShared();
+                } break;
+                case Request::LOCK_EXCLUSIVE: {
+                    lock_ = sm->LockExclusive();
+                } break;
+                case Request::EXIT:
+                    [[fallthrough]];
+                case Request::UNLOCK: {
+                    lock_.reset();
+                } break;
+                case Request::UNKNOWN:
+                    [[fallthrough]];
+                default:
+                    break;
+            }
+
+            // Pop next request. This thread is the only thread that
+            // pops from the front of the requests_ deque.
+            {
+                std::unique_lock<std::mutex> ulock(mutex_);
+                if (next_request == Request::EXIT) {
+                    requests_.clear();
+                } else {
+                    requests_.pop_front();
+                }
+            }
+            cv_.notify_all();
+        } while (next_request != Request::EXIT);
+    }
+
+    std::mutex mutex_;
+    std::condition_variable cv_;
+    std::deque<Request> requests_;
+    std::unique_ptr<SnapshotManager::LockedFile> lock_;
+    std::future<void> future_;
+};
+
+class LockTest : public ::testing::Test {
+  public:
+    void SetUp() {
+        SKIP_IF_NON_VIRTUAL_AB();
+        first_consumer.StartHandleRequestsInBackground();
+        second_consumer.StartHandleRequestsInBackground();
+    }
+
+    void TearDown() {
+        RETURN_IF_NON_VIRTUAL_AB();
+        EXPECT_TRUE(first_consumer.MakeRequest(Request::EXIT));
+        EXPECT_TRUE(second_consumer.MakeRequest(Request::EXIT));
+    }
+
+    static constexpr auto request_timeout = 500ms;
+    LockTestConsumer first_consumer;
+    LockTestConsumer second_consumer;
+};
+
+TEST_F(LockTest, SharedShared) {
+    ASSERT_TRUE(first_consumer.MakeRequest(Request::LOCK_SHARED));
+    ASSERT_TRUE(first_consumer.WaitFulfill(request_timeout));
+    ASSERT_TRUE(second_consumer.MakeRequest(Request::LOCK_SHARED));
+    ASSERT_TRUE(second_consumer.WaitFulfill(request_timeout));
+}
+
+using LockTestParam = std::pair<Request, Request>;
+class LockTestP : public LockTest, public ::testing::WithParamInterface<LockTestParam> {};
+TEST_P(LockTestP, Test) {
+    ASSERT_TRUE(first_consumer.MakeRequest(GetParam().first));
+    ASSERT_TRUE(first_consumer.WaitFulfill(request_timeout));
+    ASSERT_TRUE(second_consumer.MakeRequest(GetParam().second));
+    ASSERT_FALSE(second_consumer.WaitFulfill(request_timeout))
+            << "Should not be able to " << GetParam().second << " while separate thread "
+            << GetParam().first;
+    ASSERT_TRUE(first_consumer.MakeRequest(Request::UNLOCK));
+    ASSERT_TRUE(second_consumer.WaitFulfill(request_timeout))
+            << "Should be able to hold lock that is released by separate thread";
+}
+INSTANTIATE_TEST_SUITE_P(
+        LockTest, LockTestP,
+        testing::Values(LockTestParam{Request::LOCK_EXCLUSIVE, Request::LOCK_EXCLUSIVE},
+                        LockTestParam{Request::LOCK_EXCLUSIVE, Request::LOCK_SHARED},
+                        LockTestParam{Request::LOCK_SHARED, Request::LOCK_EXCLUSIVE}),
+        [](const testing::TestParamInfo<LockTestP::ParamType>& info) {
+            std::stringstream ss;
+            ss << info.param.first << info.param.second;
+            return ss.str();
+        });
+
+class SnapshotUpdateTest : public SnapshotTest {
+  public:
+    void SetUp() override {
+        SKIP_IF_NON_VIRTUAL_AB();
+
+        SnapshotTest::SetUp();
+        Cleanup();
+
+        // Cleanup() changes slot suffix, so initialize it again.
+        test_device->set_slot_suffix("_a");
+
+        opener_ = std::make_unique<TestPartitionOpener>(fake_super);
+
+        auto dynamic_partition_metadata = manifest_.mutable_dynamic_partition_metadata();
+        dynamic_partition_metadata->set_vabc_enabled(IsCompressionEnabled());
+        dynamic_partition_metadata->set_cow_version(android::snapshot::kCowVersionMajor);
+
+        // Create a fake update package metadata.
+        // Not using full name "system", "vendor", "product" because these names collide with the
+        // mapped partitions on the running device.
+        // Each test modifies manifest_ slightly to indicate changes to the partition layout.
+        group_ = dynamic_partition_metadata->add_groups();
+        group_->set_name("group");
+        group_->set_size(kGroupSize);
+        group_->add_partition_names("sys");
+        group_->add_partition_names("vnd");
+        group_->add_partition_names("prd");
+        sys_ = manifest_.add_partitions();
+        sys_->set_partition_name("sys");
+        sys_->set_estimate_cow_size(2_MiB);
+        SetSize(sys_, 3_MiB);
+        vnd_ = manifest_.add_partitions();
+        vnd_->set_partition_name("vnd");
+        vnd_->set_estimate_cow_size(2_MiB);
+        SetSize(vnd_, 3_MiB);
+        prd_ = manifest_.add_partitions();
+        prd_->set_partition_name("prd");
+        prd_->set_estimate_cow_size(2_MiB);
+        SetSize(prd_, 3_MiB);
+
+        // Initialize source partition metadata using |manifest_|.
+        src_ = MetadataBuilder::New(*opener_, "super", 0);
+        ASSERT_NE(src_, nullptr);
+        ASSERT_TRUE(FillFakeMetadata(src_.get(), manifest_, "_a"));
+        // Add sys_b which is like system_other.
+        ASSERT_TRUE(src_->AddGroup("group_b", kGroupSize));
+        auto partition = src_->AddPartition("sys_b", "group_b", 0);
+        ASSERT_NE(nullptr, partition);
+        ASSERT_TRUE(src_->ResizePartition(partition, 1_MiB));
+        auto metadata = src_->Export();
+        ASSERT_NE(nullptr, metadata);
+        ASSERT_TRUE(UpdatePartitionTable(*opener_, "super", *metadata.get(), 0));
+
+        // Map source partitions. Additionally, map sys_b to simulate system_other after flashing.
+        std::string path;
+        for (const auto& name : {"sys_a", "vnd_a", "prd_a", "sys_b"}) {
+            ASSERT_TRUE(CreateLogicalPartition(
+                    CreateLogicalPartitionParams{
+                            .block_device = fake_super,
+                            .metadata_slot = 0,
+                            .partition_name = name,
+                            .timeout_ms = 1s,
+                            .partition_opener = opener_.get(),
+                    },
+                    &path));
+            ASSERT_TRUE(WriteRandomData(path));
+            auto hash = GetHash(path);
+            ASSERT_TRUE(hash.has_value());
+            hashes_[name] = *hash;
+        }
+
+        // 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));
+        }
+    }
+    void TearDown() override {
+        RETURN_IF_NON_VIRTUAL_AB();
+
+        Cleanup();
+        SnapshotTest::TearDown();
+    }
+    void Cleanup() {
+        if (!image_manager_) {
+            InitializeState();
+        }
+        MountMetadata();
+        for (const auto& suffix : {"_a", "_b"}) {
+            test_device->set_slot_suffix(suffix);
+
+            // Cheat our way out of merge failed states.
+            if (sm->ProcessUpdateState() == UpdateState::MergeFailed) {
+                ASSERT_TRUE(AcquireLock());
+                ASSERT_TRUE(sm->WriteUpdateState(lock_.get(), UpdateState::None));
+                lock_ = {};
+            }
+
+            EXPECT_TRUE(sm->CancelUpdate()) << suffix;
+        }
+        EXPECT_TRUE(UnmapAll());
+    }
+
+    AssertionResult IsPartitionUnchanged(const std::string& name) {
+        std::string path;
+        if (!dm_.GetDmDevicePathByName(name, &path)) {
+            return AssertionFailure() << "Path of " << name << " cannot be determined";
+        }
+        auto hash = GetHash(path);
+        if (!hash.has_value()) {
+            return AssertionFailure() << "Cannot read partition " << name << ": " << path;
+        }
+        auto it = hashes_.find(name);
+        if (it == hashes_.end()) {
+            return AssertionFailure() << "No existing hash for " << name << ". Bad test code?";
+        }
+        if (it->second != *hash) {
+            return AssertionFailure() << "Content of " << name << " has changed";
+        }
+        return AssertionSuccess();
+    }
+
+    std::optional<uint64_t> GetSnapshotSize(const std::string& name) {
+        if (!AcquireLock()) {
+            return std::nullopt;
+        }
+        auto local_lock = std::move(lock_);
+
+        SnapshotStatus status;
+        if (!sm->ReadSnapshotStatus(local_lock.get(), name, &status)) {
+            return std::nullopt;
+        }
+        return status.snapshot_size();
+    }
+
+    AssertionResult UnmapAll() {
+        for (const auto& name : {"sys", "vnd", "prd", "dlkm"}) {
+            if (!dm_.DeleteDeviceIfExists(name + "_a"s)) {
+                return AssertionFailure() << "Cannot unmap " << name << "_a";
+            }
+            if (!DeleteSnapshotDevice(name + "_b"s)) {
+                return AssertionFailure() << "Cannot delete snapshot " << name << "_b";
+            }
+        }
+        return AssertionSuccess();
+    }
+
+    AssertionResult MapOneUpdateSnapshot(const std::string& name) {
+        if (IsCompressionEnabled()) {
+            std::unique_ptr<ISnapshotWriter> writer;
+            return MapUpdateSnapshot(name, &writer);
+        } else {
+            std::string path;
+            return MapUpdateSnapshot(name, &path);
+        }
+    }
+
+    AssertionResult WriteSnapshotAndHash(const std::string& name) {
+        if (IsCompressionEnabled()) {
+            std::unique_ptr<ISnapshotWriter> writer;
+            auto res = MapUpdateSnapshot(name, &writer);
+            if (!res) {
+                return res;
+            }
+            if (!WriteRandomData(writer.get(), &hashes_[name])) {
+                return AssertionFailure() << "Unable to write random data to snapshot " << name;
+            }
+            if (!writer->Finalize()) {
+                return AssertionFailure() << "Unable to finalize COW for " << name;
+            }
+        } else {
+            std::string path;
+            auto res = MapUpdateSnapshot(name, &path);
+            if (!res) {
+                return res;
+            }
+            if (!WriteRandomData(path, std::nullopt, &hashes_[name])) {
+                return AssertionFailure() << "Unable to write random data to snapshot " << name;
+            }
+        }
+
+        // Make sure updates to one device are seen by all devices.
+        sync();
+
+        return AssertionSuccess() << "Written random data to snapshot " << name
+                                  << ", hash: " << hashes_[name];
+    }
+
+    // Generate a snapshot that moves all the upper blocks down to the start.
+    // It doesn't really matter the order, we just want copies that reference
+    // blocks that won't exist if the partition shrinks.
+    AssertionResult ShiftAllSnapshotBlocks(const std::string& name, uint64_t old_size) {
+        std::unique_ptr<ISnapshotWriter> writer;
+        if (auto res = MapUpdateSnapshot(name, &writer); !res) {
+            return res;
+        }
+        if (!writer->options().max_blocks || !*writer->options().max_blocks) {
+            return AssertionFailure() << "No max blocks set for " << name << " writer";
+        }
+
+        uint64_t src_block = (old_size / writer->options().block_size) - 1;
+        uint64_t dst_block = 0;
+        uint64_t max_blocks = *writer->options().max_blocks;
+        while (dst_block < max_blocks && dst_block < src_block) {
+            if (!writer->AddCopy(dst_block, src_block)) {
+                return AssertionFailure() << "Unable to add copy for " << name << " for blocks "
+                                          << src_block << ", " << dst_block;
+            }
+            dst_block++;
+            src_block--;
+        }
+        if (!writer->Finalize()) {
+            return AssertionFailure() << "Unable to finalize writer for " << name;
+        }
+
+        auto hash = HashSnapshot(writer.get());
+        if (hash.empty()) {
+            return AssertionFailure() << "Unable to hash snapshot writer for " << name;
+        }
+        hashes_[name] = hash;
+
+        return AssertionSuccess();
+    }
+
+    AssertionResult MapUpdateSnapshots(const std::vector<std::string>& names = {"sys_b", "vnd_b",
+                                                                                "prd_b"}) {
+        for (const auto& name : names) {
+            auto res = MapOneUpdateSnapshot(name);
+            if (!res) {
+                return res;
+            }
+        }
+        return AssertionSuccess();
+    }
+
+    // Create fake install operations to grow the COW device size.
+    void AddOperation(PartitionUpdate* partition_update, uint64_t size_bytes = 0) {
+        auto e = partition_update->add_operations()->add_dst_extents();
+        e->set_start_block(0);
+        if (size_bytes == 0) {
+            size_bytes = GetSize(partition_update);
+        }
+        e->set_num_blocks(size_bytes / manifest_.block_size());
+    }
+
+    void AddOperationForPartitions(std::vector<PartitionUpdate*> partitions = {}) {
+        if (partitions.empty()) {
+            partitions = {sys_, vnd_, prd_};
+        }
+        for (auto* partition : partitions) {
+            AddOperation(partition);
+        }
+    }
+
+    std::unique_ptr<TestPartitionOpener> opener_;
+    DeltaArchiveManifest manifest_;
+    std::unique_ptr<MetadataBuilder> src_;
+    std::map<std::string, std::string> hashes_;
+
+    PartitionUpdate* sys_ = nullptr;
+    PartitionUpdate* vnd_ = nullptr;
+    PartitionUpdate* prd_ = nullptr;
+    DynamicPartitionGroup* group_ = nullptr;
+};
+
+// Test full update flow executed by update_engine. Some partitions uses super empty space,
+// some uses images, and some uses both.
+// Also test UnmapUpdateSnapshot unmaps everything.
+// Also test first stage mount and merge after this.
+TEST_F(SnapshotUpdateTest, FullUpdateFlow) {
+    // Grow all partitions. Set |prd| large enough that |sys| and |vnd|'s COWs
+    // fit in super, but not |prd|.
+    constexpr uint64_t partition_size = 3788_KiB;
+    SetSize(sys_, partition_size);
+    SetSize(vnd_, partition_size);
+    SetSize(prd_, 18_MiB);
+
+    // Make sure |prd| does not fit in super at all. On VABC, this means we
+    // fake an extra large COW for |vnd| to fill up super.
+    vnd_->set_estimate_cow_size(30_MiB);
+    prd_->set_estimate_cow_size(30_MiB);
+
+    AddOperationForPartitions();
+
+    // Execute the update.
+    ASSERT_TRUE(sm->BeginUpdate());
+    ASSERT_TRUE(sm->CreateUpdateSnapshots(manifest_));
+
+    // Test that partitions prioritize using space in super.
+    auto tgt = MetadataBuilder::New(*opener_, "super", 1);
+    ASSERT_NE(tgt, nullptr);
+    ASSERT_NE(nullptr, tgt->FindPartition("sys_b-cow"));
+    ASSERT_NE(nullptr, tgt->FindPartition("vnd_b-cow"));
+    ASSERT_EQ(nullptr, tgt->FindPartition("prd_b-cow"));
+
+    // Write some data to target partitions.
+    for (const auto& name : {"sys_b", "vnd_b", "prd_b"}) {
+        ASSERT_TRUE(WriteSnapshotAndHash(name));
+    }
+
+    // Assert that source partitions aren't affected.
+    for (const auto& name : {"sys_a", "vnd_a", "prd_a"}) {
+        ASSERT_TRUE(IsPartitionUnchanged(name));
+    }
+
+    ASSERT_TRUE(sm->FinishedSnapshotWrites(false));
+
+    // Simulate shutting down the device.
+    ASSERT_TRUE(UnmapAll());
+
+    // After reboot, init does first stage mount.
+    auto init = NewManagerForFirstStageMount("_b");
+    ASSERT_NE(init, nullptr);
+    ASSERT_TRUE(init->NeedSnapshotsInFirstStageMount());
+    ASSERT_TRUE(init->CreateLogicalAndSnapshotPartitions("super", snapshot_timeout_));
+
+    auto indicator = sm->GetRollbackIndicatorPath();
+    ASSERT_NE(access(indicator.c_str(), R_OK), 0);
+
+    // Check that the target partitions have the same content.
+    for (const auto& name : {"sys_b", "vnd_b", "prd_b"}) {
+        ASSERT_TRUE(IsPartitionUnchanged(name));
+    }
+
+    // Initiate the merge and wait for it to be completed.
+    ASSERT_TRUE(init->InitiateMerge());
+    ASSERT_EQ(init->IsSnapuserdRequired(), IsCompressionEnabled());
+    {
+        // We should have started in SECOND_PHASE since nothing shrinks.
+        ASSERT_TRUE(AcquireLock());
+        auto local_lock = std::move(lock_);
+        auto status = init->ReadSnapshotUpdateStatus(local_lock.get());
+        ASSERT_EQ(status.merge_phase(), MergePhase::SECOND_PHASE);
+    }
+    ASSERT_EQ(UpdateState::MergeCompleted, init->ProcessUpdateState());
+
+    // Make sure the second phase ran and deleted snapshots.
+    {
+        ASSERT_TRUE(AcquireLock());
+        auto local_lock = std::move(lock_);
+        std::vector<std::string> snapshots;
+        ASSERT_TRUE(init->ListSnapshots(local_lock.get(), &snapshots));
+        ASSERT_TRUE(snapshots.empty());
+    }
+
+    // Check that the target partitions have the same content after the merge.
+    for (const auto& name : {"sys_b", "vnd_b", "prd_b"}) {
+        ASSERT_TRUE(IsPartitionUnchanged(name))
+                << "Content of " << name << " changes after the merge";
+    }
+}
+
+TEST_F(SnapshotUpdateTest, DuplicateOps) {
+    if (!IsCompressionEnabled()) {
+        GTEST_SKIP() << "Compression-only test";
+    }
+
+    // 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));
+    }
+
+    std::vector<PartitionUpdate*> partitions = {sys_, vnd_, prd_};
+    for (auto* partition : partitions) {
+        AddOperation(partition);
+
+        std::unique_ptr<ISnapshotWriter> writer;
+        auto res = MapUpdateSnapshot(partition->partition_name() + "_b", &writer);
+        ASSERT_TRUE(res);
+        ASSERT_TRUE(writer->AddZeroBlocks(0, 1));
+        ASSERT_TRUE(writer->AddZeroBlocks(0, 1));
+        ASSERT_TRUE(writer->Finalize());
+    }
+
+    ASSERT_TRUE(sm->FinishedSnapshotWrites(false));
+
+    // Simulate shutting down the device.
+    ASSERT_TRUE(UnmapAll());
+
+    // After reboot, init does first stage mount.
+    auto init = NewManagerForFirstStageMount("_b");
+    ASSERT_NE(init, nullptr);
+    ASSERT_TRUE(init->NeedSnapshotsInFirstStageMount());
+    ASSERT_TRUE(init->CreateLogicalAndSnapshotPartitions("super", snapshot_timeout_));
+
+    // Initiate the merge and wait for it to be completed.
+    ASSERT_TRUE(init->InitiateMerge());
+    ASSERT_EQ(UpdateState::MergeCompleted, init->ProcessUpdateState());
+}
+
+// Test that shrinking and growing partitions at the same time is handled
+// correctly in VABC.
+TEST_F(SnapshotUpdateTest, SpaceSwapUpdate) {
+    if (!IsCompressionEnabled()) {
+        // b/179111359
+        GTEST_SKIP() << "Skipping Virtual A/B Compression test";
+    }
+
+    auto old_sys_size = GetSize(sys_);
+    auto old_prd_size = GetSize(prd_);
+
+    // Grow |sys| but shrink |prd|.
+    SetSize(sys_, old_sys_size * 2);
+    sys_->set_estimate_cow_size(8_MiB);
+    SetSize(prd_, old_prd_size / 2);
+    prd_->set_estimate_cow_size(1_MiB);
+
+    AddOperationForPartitions();
+
+    ASSERT_TRUE(sm->BeginUpdate());
+    ASSERT_TRUE(sm->CreateUpdateSnapshots(manifest_));
+
+    // Check that the old partition sizes were saved correctly.
+    {
+        ASSERT_TRUE(AcquireLock());
+        auto local_lock = std::move(lock_);
+
+        SnapshotStatus status;
+        ASSERT_TRUE(sm->ReadSnapshotStatus(local_lock.get(), "prd_b", &status));
+        ASSERT_EQ(status.old_partition_size(), 3145728);
+        ASSERT_TRUE(sm->ReadSnapshotStatus(local_lock.get(), "sys_b", &status));
+        ASSERT_EQ(status.old_partition_size(), 3145728);
+    }
+
+    ASSERT_TRUE(WriteSnapshotAndHash("sys_b"));
+    ASSERT_TRUE(WriteSnapshotAndHash("vnd_b"));
+    ASSERT_TRUE(ShiftAllSnapshotBlocks("prd_b", old_prd_size));
+
+    sync();
+
+    // Assert that source partitions aren't affected.
+    for (const auto& name : {"sys_a", "vnd_a", "prd_a"}) {
+        ASSERT_TRUE(IsPartitionUnchanged(name));
+    }
+
+    ASSERT_TRUE(sm->FinishedSnapshotWrites(false));
+
+    // Simulate shutting down the device.
+    ASSERT_TRUE(UnmapAll());
+
+    // After reboot, init does first stage mount.
+    auto init = NewManagerForFirstStageMount("_b");
+    ASSERT_NE(init, nullptr);
+    ASSERT_TRUE(init->NeedSnapshotsInFirstStageMount());
+    ASSERT_TRUE(init->CreateLogicalAndSnapshotPartitions("super", snapshot_timeout_));
+
+    auto indicator = sm->GetRollbackIndicatorPath();
+    ASSERT_NE(access(indicator.c_str(), R_OK), 0);
+
+    // Check that the target partitions have the same content.
+    for (const auto& name : {"sys_b", "vnd_b", "prd_b"}) {
+        ASSERT_TRUE(IsPartitionUnchanged(name));
+    }
+
+    // Initiate the merge and wait for it to be completed.
+    ASSERT_TRUE(init->InitiateMerge());
+    ASSERT_EQ(init->IsSnapuserdRequired(), IsCompressionEnabled());
+    {
+        // Check that the merge phase is FIRST_PHASE until at least one call
+        // to ProcessUpdateState() occurs.
+        ASSERT_TRUE(AcquireLock());
+        auto local_lock = std::move(lock_);
+        auto status = init->ReadSnapshotUpdateStatus(local_lock.get());
+        ASSERT_EQ(status.merge_phase(), MergePhase::FIRST_PHASE);
+    }
+
+    // Simulate shutting down the device and creating partitions again.
+    ASSERT_TRUE(UnmapAll());
+    ASSERT_TRUE(init->CreateLogicalAndSnapshotPartitions("super", snapshot_timeout_));
+
+    // Check that we used the correct types after rebooting mid-merge.
+    DeviceMapper::TargetInfo target;
+    ASSERT_TRUE(init->IsSnapshotDevice("prd_b", &target));
+    ASSERT_EQ(DeviceMapper::GetTargetType(target.spec), "user");
+    ASSERT_TRUE(init->IsSnapshotDevice("sys_b", &target));
+    ASSERT_EQ(DeviceMapper::GetTargetType(target.spec), "user");
+    ASSERT_TRUE(init->IsSnapshotDevice("vnd_b", &target));
+    ASSERT_EQ(DeviceMapper::GetTargetType(target.spec), "user");
+
+    // Complete the merge.
+    ASSERT_EQ(UpdateState::MergeCompleted, init->ProcessUpdateState());
+
+    // Make sure the second phase ran and deleted snapshots.
+    {
+        ASSERT_TRUE(AcquireLock());
+        auto local_lock = std::move(lock_);
+        std::vector<std::string> snapshots;
+        ASSERT_TRUE(init->ListSnapshots(local_lock.get(), &snapshots));
+        ASSERT_TRUE(snapshots.empty());
+    }
+
+    // Check that the target partitions have the same content after the merge.
+    for (const auto& name : {"sys_b", "vnd_b", "prd_b"}) {
+        ASSERT_TRUE(IsPartitionUnchanged(name))
+                << "Content of " << name << " changes after the merge";
+    }
+}
+
+// Test that if new system partitions uses empty space in super, that region is not snapshotted.
+TEST_F(SnapshotUpdateTest, DirectWriteEmptySpace) {
+    GTEST_SKIP() << "b/141889746";
+    SetSize(sys_, 4_MiB);
+    // vnd_b and prd_b are unchanged.
+    ASSERT_TRUE(sm->CreateUpdateSnapshots(manifest_));
+    ASSERT_EQ(3_MiB, GetSnapshotSize("sys_b").value_or(0));
+}
+
+// Test that if new system partitions uses space of old vendor partition, that region is
+// snapshotted.
+TEST_F(SnapshotUpdateTest, SnapshotOldPartitions) {
+    SetSize(sys_, 4_MiB);  // grows
+    SetSize(vnd_, 2_MiB);  // shrinks
+    // prd_b is unchanged
+    ASSERT_TRUE(sm->BeginUpdate());
+    ASSERT_TRUE(sm->CreateUpdateSnapshots(manifest_));
+    ASSERT_EQ(4_MiB, GetSnapshotSize("sys_b").value_or(0));
+}
+
+// Test that even if there seem to be empty space in target metadata, COW partition won't take
+// it because they are used by old partitions.
+TEST_F(SnapshotUpdateTest, CowPartitionDoNotTakeOldPartitions) {
+    SetSize(sys_, 2_MiB);  // shrinks
+    // vnd_b and prd_b are unchanged.
+    ASSERT_TRUE(sm->BeginUpdate());
+    ASSERT_TRUE(sm->CreateUpdateSnapshots(manifest_));
+
+    auto tgt = MetadataBuilder::New(*opener_, "super", 1);
+    ASSERT_NE(nullptr, tgt);
+    auto metadata = tgt->Export();
+    ASSERT_NE(nullptr, metadata);
+    std::vector<std::string> written;
+    // Write random data to all COW partitions in super
+    for (auto p : metadata->partitions) {
+        if (GetPartitionGroupName(metadata->groups[p.group_index]) != kCowGroupName) {
+            continue;
+        }
+        std::string path;
+        ASSERT_TRUE(CreateLogicalPartition(
+                CreateLogicalPartitionParams{
+                        .block_device = fake_super,
+                        .metadata = metadata.get(),
+                        .partition = &p,
+                        .timeout_ms = 1s,
+                        .partition_opener = opener_.get(),
+                },
+                &path));
+        ASSERT_TRUE(WriteRandomData(path));
+        written.push_back(GetPartitionName(p));
+    }
+    ASSERT_FALSE(written.empty())
+            << "No COW partitions are created even if there are empty space in super partition";
+
+    // Make sure source partitions aren't affected.
+    for (const auto& name : {"sys_a", "vnd_a", "prd_a"}) {
+        ASSERT_TRUE(IsPartitionUnchanged(name));
+    }
+}
+
+// Test that it crashes after creating snapshot status file but before creating COW image, then
+// calling CreateUpdateSnapshots again works.
+TEST_F(SnapshotUpdateTest, SnapshotStatusFileWithoutCow) {
+    // Write some trash snapshot files to simulate leftovers from previous runs.
+    {
+        ASSERT_TRUE(AcquireLock());
+        auto local_lock = std::move(lock_);
+        SnapshotStatus status;
+        status.set_name("sys_b");
+        ASSERT_TRUE(sm->WriteSnapshotStatus(local_lock.get(), status));
+        ASSERT_TRUE(image_manager_->CreateBackingImage("sys_b-cow-img", 1_MiB,
+                                                       IImageManager::CREATE_IMAGE_DEFAULT));
+    }
+
+    // Redo the update.
+    ASSERT_TRUE(sm->BeginUpdate());
+    ASSERT_TRUE(sm->UnmapUpdateSnapshot("sys_b"));
+
+    ASSERT_TRUE(sm->CreateUpdateSnapshots(manifest_));
+
+    // Check that target partitions can be mapped.
+    EXPECT_TRUE(MapUpdateSnapshots());
+}
+
+// Test that the old partitions are not modified.
+TEST_F(SnapshotUpdateTest, TestRollback) {
+    // Execute the update.
+    ASSERT_TRUE(sm->BeginUpdate());
+    ASSERT_TRUE(sm->UnmapUpdateSnapshot("sys_b"));
+
+    AddOperationForPartitions();
+
+    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 that source partitions aren't affected.
+    for (const auto& name : {"sys_a", "vnd_a", "prd_a"}) {
+        ASSERT_TRUE(IsPartitionUnchanged(name));
+    }
+
+    ASSERT_TRUE(sm->FinishedSnapshotWrites(false));
+
+    // Simulate shutting down the device.
+    ASSERT_TRUE(UnmapAll());
+
+    // After reboot, init does first stage mount.
+    auto init = NewManagerForFirstStageMount("_b");
+    ASSERT_NE(init, nullptr);
+    ASSERT_TRUE(init->NeedSnapshotsInFirstStageMount());
+    ASSERT_TRUE(init->CreateLogicalAndSnapshotPartitions("super", snapshot_timeout_));
+
+    // Check that the target partitions have the same content.
+    for (const auto& name : {"sys_b", "vnd_b", "prd_b"}) {
+        ASSERT_TRUE(IsPartitionUnchanged(name));
+    }
+
+    // Simulate shutting down the device again.
+    ASSERT_TRUE(UnmapAll());
+    init = NewManagerForFirstStageMount("_a");
+    ASSERT_NE(init, nullptr);
+    ASSERT_FALSE(init->NeedSnapshotsInFirstStageMount());
+    ASSERT_TRUE(init->CreateLogicalAndSnapshotPartitions("super", snapshot_timeout_));
+
+    // Assert that the source partitions aren't affected.
+    for (const auto& name : {"sys_a", "vnd_a", "prd_a"}) {
+        ASSERT_TRUE(IsPartitionUnchanged(name));
+    }
+}
+
+// Test that if an update is applied but not booted into, it can be canceled.
+TEST_F(SnapshotUpdateTest, CancelAfterApply) {
+    ASSERT_TRUE(sm->BeginUpdate());
+    ASSERT_TRUE(sm->FinishedSnapshotWrites(false));
+    ASSERT_TRUE(sm->CancelUpdate());
+}
+
+static std::vector<Interval> ToIntervals(const std::vector<std::unique_ptr<Extent>>& extents) {
+    std::vector<Interval> ret;
+    std::transform(extents.begin(), extents.end(), std::back_inserter(ret),
+                   [](const auto& extent) { return extent->AsLinearExtent()->AsInterval(); });
+    return ret;
+}
+
+// Test that at the second update, old COW partition spaces are reclaimed.
+TEST_F(SnapshotUpdateTest, ReclaimCow) {
+    // Make sure VABC cows are small enough that they fit in fake_super.
+    sys_->set_estimate_cow_size(64_KiB);
+    vnd_->set_estimate_cow_size(64_KiB);
+    prd_->set_estimate_cow_size(64_KiB);
+
+    // Execute the first update.
+    ASSERT_TRUE(sm->BeginUpdate());
+    ASSERT_TRUE(sm->CreateUpdateSnapshots(manifest_));
+    ASSERT_TRUE(MapUpdateSnapshots());
+    ASSERT_TRUE(sm->FinishedSnapshotWrites(false));
+
+    // Simulate shutting down the device.
+    ASSERT_TRUE(UnmapAll());
+
+    // After reboot, init does first stage mount.
+    auto init = NewManagerForFirstStageMount("_b");
+    ASSERT_NE(init, nullptr);
+    ASSERT_TRUE(init->NeedSnapshotsInFirstStageMount());
+    ASSERT_TRUE(init->CreateLogicalAndSnapshotPartitions("super", snapshot_timeout_));
+    init = nullptr;
+
+    // Initiate the merge and wait for it to be completed.
+    auto new_sm = SnapshotManager::New(new TestDeviceInfo(fake_super, "_b"));
+    ASSERT_TRUE(new_sm->InitiateMerge());
+    ASSERT_EQ(UpdateState::MergeCompleted, new_sm->ProcessUpdateState());
+
+    // Execute the second update.
+    ASSERT_TRUE(new_sm->BeginUpdate());
+    ASSERT_TRUE(new_sm->CreateUpdateSnapshots(manifest_));
+
+    // Check that the old COW space is reclaimed and does not occupy space of mapped partitions.
+    auto src = MetadataBuilder::New(*opener_, "super", 1);
+    ASSERT_NE(src, nullptr);
+    auto tgt = MetadataBuilder::New(*opener_, "super", 0);
+    ASSERT_NE(tgt, nullptr);
+    for (const auto& cow_part_name : {"sys_a-cow", "vnd_a-cow", "prd_a-cow"}) {
+        auto* cow_part = tgt->FindPartition(cow_part_name);
+        ASSERT_NE(nullptr, cow_part) << cow_part_name << " does not exist in target metadata";
+        auto cow_intervals = ToIntervals(cow_part->extents());
+        for (const auto& old_part_name : {"sys_b", "vnd_b", "prd_b"}) {
+            auto* old_part = src->FindPartition(old_part_name);
+            ASSERT_NE(nullptr, old_part) << old_part_name << " does not exist in source metadata";
+            auto old_intervals = ToIntervals(old_part->extents());
+
+            auto intersect = Interval::Intersect(cow_intervals, old_intervals);
+            ASSERT_TRUE(intersect.empty()) << "COW uses space of source partitions";
+        }
+    }
+}
+
+TEST_F(SnapshotUpdateTest, RetrofitAfterRegularAb) {
+    constexpr auto kRetrofitGroupSize = kGroupSize / 2;
+
+    // Initialize device-mapper / disk
+    ASSERT_TRUE(UnmapAll());
+    FormatFakeSuper();
+
+    // Setup source partition metadata to have both _a and _b partitions.
+    src_ = MetadataBuilder::New(*opener_, "super", 0);
+    ASSERT_NE(nullptr, src_);
+    for (const auto& suffix : {"_a"s, "_b"s}) {
+        ASSERT_TRUE(src_->AddGroup(group_->name() + suffix, kRetrofitGroupSize));
+        for (const auto& name : {"sys"s, "vnd"s, "prd"s}) {
+            auto partition = src_->AddPartition(name + suffix, group_->name() + suffix, 0);
+            ASSERT_NE(nullptr, partition);
+            ASSERT_TRUE(src_->ResizePartition(partition, 2_MiB));
+        }
+    }
+    auto metadata = src_->Export();
+    ASSERT_NE(nullptr, metadata);
+    ASSERT_TRUE(UpdatePartitionTable(*opener_, "super", *metadata.get(), 0));
+
+    // Flash source partitions
+    std::string path;
+    for (const auto& name : {"sys_a", "vnd_a", "prd_a"}) {
+        ASSERT_TRUE(CreateLogicalPartition(
+                CreateLogicalPartitionParams{
+                        .block_device = fake_super,
+                        .metadata_slot = 0,
+                        .partition_name = name,
+                        .timeout_ms = 1s,
+                        .partition_opener = opener_.get(),
+                },
+                &path));
+        ASSERT_TRUE(WriteRandomData(path));
+        auto hash = GetHash(path);
+        ASSERT_TRUE(hash.has_value());
+        hashes_[name] = *hash;
+    }
+
+    // Setup manifest.
+    group_->set_size(kRetrofitGroupSize);
+    for (auto* partition : {sys_, vnd_, prd_}) {
+        SetSize(partition, 2_MiB);
+    }
+    AddOperationForPartitions();
+
+    ASSERT_TRUE(sm->BeginUpdate());
+    ASSERT_TRUE(sm->CreateUpdateSnapshots(manifest_));
+
+    // Test that COW image should not be created for retrofit devices; super
+    // should be big enough.
+    ASSERT_FALSE(image_manager_->BackingImageExists("sys_b-cow-img"));
+    ASSERT_FALSE(image_manager_->BackingImageExists("vnd_b-cow-img"));
+    ASSERT_FALSE(image_manager_->BackingImageExists("prd_b-cow-img"));
+
+    // Write some data to target partitions.
+    for (const auto& name : {"sys_b", "vnd_b", "prd_b"}) {
+        ASSERT_TRUE(WriteSnapshotAndHash(name));
+    }
+
+    // Assert that source partitions aren't affected.
+    for (const auto& name : {"sys_a", "vnd_a", "prd_a"}) {
+        ASSERT_TRUE(IsPartitionUnchanged(name));
+    }
+
+    ASSERT_TRUE(sm->FinishedSnapshotWrites(false));
+}
+
+TEST_F(SnapshotUpdateTest, MergeCannotRemoveCow) {
+    // Make source partitions as big as possible to force COW image to be created.
+    SetSize(sys_, 10_MiB);
+    SetSize(vnd_, 10_MiB);
+    SetSize(prd_, 10_MiB);
+    sys_->set_estimate_cow_size(12_MiB);
+    vnd_->set_estimate_cow_size(12_MiB);
+    prd_->set_estimate_cow_size(12_MiB);
+
+    src_ = MetadataBuilder::New(*opener_, "super", 0);
+    ASSERT_NE(src_, nullptr);
+    src_->RemoveGroupAndPartitions(group_->name() + "_a");
+    src_->RemoveGroupAndPartitions(group_->name() + "_b");
+    ASSERT_TRUE(FillFakeMetadata(src_.get(), manifest_, "_a"));
+    auto metadata = src_->Export();
+    ASSERT_NE(nullptr, metadata);
+    ASSERT_TRUE(UpdatePartitionTable(*opener_, "super", *metadata.get(), 0));
+
+    // Add operations for sys. The whole device is written.
+    AddOperation(sys_);
+
+    // Execute the update.
+    ASSERT_TRUE(sm->BeginUpdate());
+    ASSERT_TRUE(sm->CreateUpdateSnapshots(manifest_));
+    ASSERT_TRUE(MapUpdateSnapshots());
+    ASSERT_TRUE(sm->FinishedSnapshotWrites(false));
+
+    // Simulate shutting down the device.
+    ASSERT_TRUE(UnmapAll());
+
+    // After reboot, init does first stage mount.
+    // Normally we should use NewManagerForFirstStageMount, but if so,
+    // "gsid.mapped_image.sys_b-cow-img" won't be set.
+    auto init = SnapshotManager::New(new TestDeviceInfo(fake_super, "_b"));
+    ASSERT_NE(init, nullptr);
+    ASSERT_TRUE(init->CreateLogicalAndSnapshotPartitions("super", snapshot_timeout_));
+
+    // Keep an open handle to the cow device. This should cause the merge to
+    // be incomplete.
+    auto cow_path = android::base::GetProperty("gsid.mapped_image.sys_b-cow-img", "");
+    unique_fd fd(open(cow_path.c_str(), O_RDONLY | O_CLOEXEC));
+    ASSERT_GE(fd, 0);
+
+    // COW cannot be removed due to open fd, so expect a soft failure.
+    ASSERT_TRUE(init->InitiateMerge());
+    ASSERT_EQ(UpdateState::MergeNeedsReboot, init->ProcessUpdateState());
+
+    // Simulate shutting down the device.
+    fd.reset();
+    ASSERT_TRUE(UnmapAll());
+
+    // init does first stage mount again.
+    ASSERT_TRUE(init->CreateLogicalAndSnapshotPartitions("super", snapshot_timeout_));
+
+    // sys_b should be mapped as a dm-linear device directly.
+    ASSERT_FALSE(sm->IsSnapshotDevice("sys_b", nullptr));
+
+    // Merge should be able to complete now.
+    ASSERT_EQ(UpdateState::MergeCompleted, init->ProcessUpdateState());
+}
+
+class MetadataMountedTest : public ::testing::Test {
+  public:
+    // This is so main() can instantiate this to invoke Cleanup.
+    virtual void TestBody() override {}
+    void SetUp() override {
+        SKIP_IF_NON_VIRTUAL_AB();
+        metadata_dir_ = test_device->GetMetadataDir();
+        ASSERT_TRUE(ReadDefaultFstab(&fstab_));
+    }
+    void TearDown() override {
+        RETURN_IF_NON_VIRTUAL_AB();
+        SetUp();
+        // Remount /metadata
+        test_device->set_recovery(false);
+        EXPECT_TRUE(android::fs_mgr::EnsurePathMounted(&fstab_, metadata_dir_));
+    }
+    AssertionResult IsMetadataMounted() {
+        Fstab mounted_fstab;
+        if (!ReadFstabFromFile("/proc/mounts", &mounted_fstab)) {
+            ADD_FAILURE() << "Failed to scan mounted volumes";
+            return AssertionFailure() << "Failed to scan mounted volumes";
+        }
+
+        auto entry = GetEntryForPath(&fstab_, metadata_dir_);
+        if (entry == nullptr) {
+            return AssertionFailure() << "No mount point found in fstab for path " << metadata_dir_;
+        }
+
+        auto mv = GetEntryForMountPoint(&mounted_fstab, entry->mount_point);
+        if (mv == nullptr) {
+            return AssertionFailure() << metadata_dir_ << " is not mounted";
+        }
+        return AssertionSuccess() << metadata_dir_ << " is mounted";
+    }
+    std::string metadata_dir_;
+    Fstab fstab_;
+};
+
+void MountMetadata() {
+    MetadataMountedTest().TearDown();
+}
+
+TEST_F(MetadataMountedTest, Android) {
+    auto device = sm->EnsureMetadataMounted();
+    EXPECT_NE(nullptr, device);
+    device.reset();
+
+    EXPECT_TRUE(IsMetadataMounted());
+    EXPECT_TRUE(sm->CancelUpdate()) << "Metadata dir should never be unmounted in Android mode";
+}
+
+TEST_F(MetadataMountedTest, Recovery) {
+    test_device->set_recovery(true);
+    metadata_dir_ = test_device->GetMetadataDir();
+
+    EXPECT_TRUE(android::fs_mgr::EnsurePathUnmounted(&fstab_, metadata_dir_));
+    EXPECT_FALSE(IsMetadataMounted());
+
+    auto device = sm->EnsureMetadataMounted();
+    EXPECT_NE(nullptr, device);
+    EXPECT_TRUE(IsMetadataMounted());
+
+    device.reset();
+    EXPECT_FALSE(IsMetadataMounted());
+}
+
+// Test that during a merge, we can wipe data in recovery.
+TEST_F(SnapshotUpdateTest, MergeInRecovery) {
+    // Execute the first update.
+    ASSERT_TRUE(sm->BeginUpdate());
+    ASSERT_TRUE(sm->CreateUpdateSnapshots(manifest_));
+    ASSERT_TRUE(MapUpdateSnapshots());
+    ASSERT_TRUE(sm->FinishedSnapshotWrites(false));
+
+    // Simulate shutting down the device.
+    ASSERT_TRUE(UnmapAll());
+
+    // After reboot, init does first stage mount.
+    auto init = NewManagerForFirstStageMount("_b");
+    ASSERT_NE(init, nullptr);
+    ASSERT_TRUE(init->NeedSnapshotsInFirstStageMount());
+    ASSERT_TRUE(init->CreateLogicalAndSnapshotPartitions("super", snapshot_timeout_));
+    init = nullptr;
+
+    // Initiate the merge and then immediately stop it to simulate a reboot.
+    auto new_sm = SnapshotManager::New(new TestDeviceInfo(fake_super, "_b"));
+    ASSERT_TRUE(new_sm->InitiateMerge());
+    ASSERT_TRUE(UnmapAll());
+
+    // Simulate a reboot into recovery.
+    auto test_device = std::make_unique<TestDeviceInfo>(fake_super, "_b");
+    test_device->set_recovery(true);
+    new_sm = NewManagerForFirstStageMount(test_device.release());
+
+    ASSERT_TRUE(new_sm->HandleImminentDataWipe());
+    ASSERT_EQ(new_sm->GetUpdateState(), UpdateState::None);
+}
+
+// Test that a merge does not clear the snapshot state in fastboot.
+TEST_F(SnapshotUpdateTest, MergeInFastboot) {
+    // Execute the first update.
+    ASSERT_TRUE(sm->BeginUpdate());
+    ASSERT_TRUE(sm->CreateUpdateSnapshots(manifest_));
+    ASSERT_TRUE(MapUpdateSnapshots());
+    ASSERT_TRUE(sm->FinishedSnapshotWrites(false));
+
+    // Simulate shutting down the device.
+    ASSERT_TRUE(UnmapAll());
+
+    // After reboot, init does first stage mount.
+    auto init = NewManagerForFirstStageMount("_b");
+    ASSERT_NE(init, nullptr);
+    ASSERT_TRUE(init->NeedSnapshotsInFirstStageMount());
+    ASSERT_TRUE(init->CreateLogicalAndSnapshotPartitions("super", snapshot_timeout_));
+    init = nullptr;
+
+    // Initiate the merge and then immediately stop it to simulate a reboot.
+    auto new_sm = SnapshotManager::New(new TestDeviceInfo(fake_super, "_b"));
+    ASSERT_TRUE(new_sm->InitiateMerge());
+    ASSERT_TRUE(UnmapAll());
+
+    // Simulate a reboot into recovery.
+    auto test_device = std::make_unique<TestDeviceInfo>(fake_super, "_b");
+    test_device->set_recovery(true);
+    new_sm = NewManagerForFirstStageMount(test_device.release());
+
+    ASSERT_TRUE(new_sm->FinishMergeInRecovery());
+
+    ASSERT_TRUE(UnmapAll());
+
+    auto mount = new_sm->EnsureMetadataMounted();
+    ASSERT_TRUE(mount && mount->HasDevice());
+    ASSERT_EQ(new_sm->ProcessUpdateState(), UpdateState::MergeCompleted);
+
+    // Finish the merge in a normal boot.
+    test_device = std::make_unique<TestDeviceInfo>(fake_super, "_b");
+    init = NewManagerForFirstStageMount(test_device.release());
+    ASSERT_TRUE(init->CreateLogicalAndSnapshotPartitions("super", snapshot_timeout_));
+    init = nullptr;
+
+    test_device = std::make_unique<TestDeviceInfo>(fake_super, "_b");
+    new_sm = NewManagerForFirstStageMount(test_device.release());
+    ASSERT_EQ(new_sm->ProcessUpdateState(), UpdateState::MergeCompleted);
+    ASSERT_EQ(new_sm->ProcessUpdateState(), UpdateState::None);
+}
+
+// Test that after an OTA, before a merge, we can wipe data in recovery.
+TEST_F(SnapshotUpdateTest, DataWipeRollbackInRecovery) {
+    // Execute the first update.
+    ASSERT_TRUE(sm->BeginUpdate());
+    ASSERT_TRUE(sm->CreateUpdateSnapshots(manifest_));
+    ASSERT_TRUE(MapUpdateSnapshots());
+    ASSERT_TRUE(sm->FinishedSnapshotWrites(false));
+
+    // Simulate shutting down the device.
+    ASSERT_TRUE(UnmapAll());
+
+    // Simulate a reboot into recovery.
+    auto test_device = new TestDeviceInfo(fake_super, "_b");
+    test_device->set_recovery(true);
+    auto new_sm = NewManagerForFirstStageMount(test_device);
+
+    ASSERT_TRUE(new_sm->HandleImminentDataWipe());
+    // Manually mount metadata so that we can call GetUpdateState() below.
+    MountMetadata();
+    EXPECT_EQ(new_sm->GetUpdateState(), UpdateState::None);
+    EXPECT_TRUE(test_device->IsSlotUnbootable(1));
+    EXPECT_FALSE(test_device->IsSlotUnbootable(0));
+}
+
+// Test that after an OTA and a bootloader rollback with no merge, we can wipe
+// data in recovery.
+TEST_F(SnapshotUpdateTest, DataWipeAfterRollback) {
+    // Execute the first update.
+    ASSERT_TRUE(sm->BeginUpdate());
+    ASSERT_TRUE(sm->CreateUpdateSnapshots(manifest_));
+    ASSERT_TRUE(MapUpdateSnapshots());
+    ASSERT_TRUE(sm->FinishedSnapshotWrites(false));
+
+    // Simulate shutting down the device.
+    ASSERT_TRUE(UnmapAll());
+
+    // Simulate a rollback, with reboot into recovery.
+    auto test_device = new TestDeviceInfo(fake_super, "_a");
+    test_device->set_recovery(true);
+    auto new_sm = NewManagerForFirstStageMount(test_device);
+
+    ASSERT_TRUE(new_sm->HandleImminentDataWipe());
+    EXPECT_EQ(new_sm->GetUpdateState(), UpdateState::None);
+    EXPECT_FALSE(test_device->IsSlotUnbootable(0));
+    EXPECT_FALSE(test_device->IsSlotUnbootable(1));
+}
+
+// Test update package that requests data wipe.
+TEST_F(SnapshotUpdateTest, DataWipeRequiredInPackage) {
+    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)) << name;
+    }
+
+    ASSERT_TRUE(sm->FinishedSnapshotWrites(true /* wipe */));
+
+    // Simulate shutting down the device.
+    ASSERT_TRUE(UnmapAll());
+
+    // Simulate a reboot into recovery.
+    auto test_device = new TestDeviceInfo(fake_super, "_b");
+    test_device->set_recovery(true);
+    auto new_sm = NewManagerForFirstStageMount(test_device);
+
+    ASSERT_TRUE(new_sm->HandleImminentDataWipe());
+    // Manually mount metadata so that we can call GetUpdateState() below.
+    MountMetadata();
+    EXPECT_EQ(new_sm->GetUpdateState(), UpdateState::None);
+    ASSERT_FALSE(test_device->IsSlotUnbootable(1));
+    ASSERT_FALSE(test_device->IsSlotUnbootable(0));
+
+    ASSERT_TRUE(UnmapAll());
+
+    // Now reboot into new slot.
+    test_device = new TestDeviceInfo(fake_super, "_b");
+    auto init = NewManagerForFirstStageMount(test_device);
+    ASSERT_TRUE(init->CreateLogicalAndSnapshotPartitions("super", snapshot_timeout_));
+    // Verify that we are on the downgraded build.
+    for (const auto& name : {"sys_b", "vnd_b", "prd_b"}) {
+        ASSERT_TRUE(IsPartitionUnchanged(name)) << name;
+    }
+}
+
+// Test update package that requests data wipe.
+TEST_F(SnapshotUpdateTest, DataWipeWithStaleSnapshots) {
+    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)) << name;
+    }
+
+    // Create a stale snapshot that should not exist.
+    {
+        ASSERT_TRUE(AcquireLock());
+
+        PartitionCowCreator cow_creator = {
+                .compression_enabled = IsCompressionEnabled(),
+                .compression_algorithm = IsCompressionEnabled() ? "gz" : "none",
+        };
+        SnapshotStatus status;
+        status.set_name("sys_a");
+        status.set_device_size(1_MiB);
+        status.set_snapshot_size(2_MiB);
+        status.set_cow_partition_size(2_MiB);
+
+        ASSERT_TRUE(sm->CreateSnapshot(lock_.get(), &cow_creator, &status));
+        lock_ = nullptr;
+
+        ASSERT_TRUE(sm->EnsureImageManager());
+        ASSERT_TRUE(sm->image_manager()->CreateBackingImage("sys_a", 1_MiB, 0));
+    }
+
+    ASSERT_TRUE(sm->FinishedSnapshotWrites(true /* wipe */));
+
+    // Simulate shutting down the device.
+    ASSERT_TRUE(UnmapAll());
+
+    // Simulate a reboot into recovery.
+    auto test_device = new TestDeviceInfo(fake_super, "_b");
+    test_device->set_recovery(true);
+    auto new_sm = NewManagerForFirstStageMount(test_device);
+
+    ASSERT_TRUE(new_sm->HandleImminentDataWipe());
+    // Manually mount metadata so that we can call GetUpdateState() below.
+    MountMetadata();
+    EXPECT_EQ(new_sm->GetUpdateState(), UpdateState::None);
+    ASSERT_FALSE(test_device->IsSlotUnbootable(1));
+    ASSERT_FALSE(test_device->IsSlotUnbootable(0));
+
+    ASSERT_TRUE(UnmapAll());
+
+    // Now reboot into new slot.
+    test_device = new TestDeviceInfo(fake_super, "_b");
+    auto init = NewManagerForFirstStageMount(test_device);
+    ASSERT_TRUE(init->CreateLogicalAndSnapshotPartitions("super", snapshot_timeout_));
+    // Verify that we are on the downgraded build.
+    for (const auto& name : {"sys_b", "vnd_b", "prd_b"}) {
+        ASSERT_TRUE(IsPartitionUnchanged(name)) << name;
+    }
+}
+
+TEST_F(SnapshotUpdateTest, Hashtree) {
+    constexpr auto partition_size = 4_MiB;
+    constexpr auto data_size = 3_MiB;
+    constexpr auto hashtree_size = 512_KiB;
+    constexpr auto fec_size = partition_size - data_size - hashtree_size;
+
+    const auto block_size = manifest_.block_size();
+    SetSize(sys_, partition_size);
+    AddOperation(sys_, data_size);
+
+    sys_->set_estimate_cow_size(partition_size + data_size);
+
+    // Set hastree extents.
+    sys_->mutable_hash_tree_data_extent()->set_start_block(0);
+    sys_->mutable_hash_tree_data_extent()->set_num_blocks(data_size / block_size);
+
+    sys_->mutable_hash_tree_extent()->set_start_block(data_size / block_size);
+    sys_->mutable_hash_tree_extent()->set_num_blocks(hashtree_size / block_size);
+
+    // Set FEC extents.
+    sys_->mutable_fec_data_extent()->set_start_block(0);
+    sys_->mutable_fec_data_extent()->set_num_blocks((data_size + hashtree_size) / block_size);
+
+    sys_->mutable_fec_extent()->set_start_block((data_size + hashtree_size) / block_size);
+    sys_->mutable_fec_extent()->set_num_blocks(fec_size / block_size);
+
+    ASSERT_TRUE(sm->BeginUpdate());
+    ASSERT_TRUE(sm->CreateUpdateSnapshots(manifest_));
+
+    // Map and write some data to target partition.
+    ASSERT_TRUE(MapUpdateSnapshots({"vnd_b", "prd_b"}));
+    ASSERT_TRUE(WriteSnapshotAndHash("sys_b"));
+
+    // Finish update.
+    ASSERT_TRUE(sm->FinishedSnapshotWrites(false));
+
+    // Simulate shutting down the device.
+    ASSERT_TRUE(UnmapAll());
+
+    // After reboot, init does first stage mount.
+    auto init = NewManagerForFirstStageMount("_b");
+    ASSERT_NE(init, nullptr);
+    ASSERT_TRUE(init->NeedSnapshotsInFirstStageMount());
+    ASSERT_TRUE(init->CreateLogicalAndSnapshotPartitions("super", snapshot_timeout_));
+
+    // Check that the target partition have the same content. Hashtree and FEC extents
+    // should be accounted for.
+    ASSERT_TRUE(IsPartitionUnchanged("sys_b"));
+}
+
+// Test for overflow bit after update
+TEST_F(SnapshotUpdateTest, Overflow) {
+    if (IsCompressionEnabled()) {
+        GTEST_SKIP() << "No overflow bit set for userspace COWs";
+    }
+
+    const auto actual_write_size = GetSize(sys_);
+    const auto declared_write_size = actual_write_size - 1_MiB;
+
+    AddOperation(sys_, declared_write_size);
+
+    // Execute the update.
+    ASSERT_TRUE(sm->BeginUpdate());
+    ASSERT_TRUE(sm->CreateUpdateSnapshots(manifest_));
+
+    // Map and write some data to target partitions.
+    ASSERT_TRUE(MapUpdateSnapshots({"vnd_b", "prd_b"}));
+    ASSERT_TRUE(WriteSnapshotAndHash("sys_b"));
+
+    std::vector<android::dm::DeviceMapper::TargetInfo> table;
+    ASSERT_TRUE(DeviceMapper::Instance().GetTableStatus("sys_b", &table));
+    ASSERT_EQ(1u, table.size());
+    EXPECT_TRUE(table[0].IsOverflowSnapshot());
+
+    ASSERT_FALSE(sm->FinishedSnapshotWrites(false))
+            << "FinishedSnapshotWrites should detect overflow of CoW device.";
+}
+
+TEST_F(SnapshotUpdateTest, LowSpace) {
+    static constexpr auto kMaxFree = 10_MiB;
+    auto userdata = std::make_unique<LowSpaceUserdata>();
+    ASSERT_TRUE(userdata->Init(kMaxFree));
+
+    // Grow all partitions to 10_MiB, total 30_MiB. This requires 30 MiB of CoW space. After
+    // using the empty space in super (< 1 MiB), it uses 30 MiB of /userdata space.
+    constexpr uint64_t partition_size = 10_MiB;
+    SetSize(sys_, partition_size);
+    SetSize(vnd_, partition_size);
+    SetSize(prd_, partition_size);
+    sys_->set_estimate_cow_size(partition_size);
+    vnd_->set_estimate_cow_size(partition_size);
+    prd_->set_estimate_cow_size(partition_size);
+
+    AddOperationForPartitions();
+
+    // Execute the update.
+    ASSERT_TRUE(sm->BeginUpdate());
+    auto res = sm->CreateUpdateSnapshots(manifest_);
+    ASSERT_FALSE(res);
+    ASSERT_EQ(Return::ErrorCode::NO_SPACE, res.error_code());
+    ASSERT_GE(res.required_size(), 14_MiB);
+    ASSERT_LT(res.required_size(), 40_MiB);
+}
+
+TEST_F(SnapshotUpdateTest, AddPartition) {
+    group_->add_partition_names("dlkm");
+
+    auto dlkm = manifest_.add_partitions();
+    dlkm->set_partition_name("dlkm");
+    dlkm->set_estimate_cow_size(2_MiB);
+    SetSize(dlkm, 3_MiB);
+
+    // Grow all partitions. Set |prd| large enough that |sys| and |vnd|'s COWs
+    // fit in super, but not |prd|.
+    constexpr uint64_t partition_size = 3788_KiB;
+    SetSize(sys_, partition_size);
+    SetSize(vnd_, partition_size);
+    SetSize(prd_, partition_size);
+    SetSize(dlkm, partition_size);
+
+    AddOperationForPartitions({sys_, vnd_, prd_, dlkm});
+
+    // 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", "dlkm_b"}) {
+        ASSERT_TRUE(WriteSnapshotAndHash(name));
+    }
+
+    // Assert that source partitions aren't affected.
+    for (const auto& name : {"sys_a", "vnd_a", "prd_a"}) {
+        ASSERT_TRUE(IsPartitionUnchanged(name));
+    }
+
+    ASSERT_TRUE(sm->FinishedSnapshotWrites(false));
+
+    // Simulate shutting down the device.
+    ASSERT_TRUE(UnmapAll());
+
+    // After reboot, init does first stage mount.
+    auto init = NewManagerForFirstStageMount("_b");
+    ASSERT_NE(init, nullptr);
+
+    ASSERT_TRUE(init->EnsureSnapuserdConnected());
+    init->set_use_first_stage_snapuserd(true);
+
+    ASSERT_TRUE(init->NeedSnapshotsInFirstStageMount());
+    ASSERT_TRUE(init->CreateLogicalAndSnapshotPartitions("super", snapshot_timeout_));
+
+    // Check that the target partitions have the same content.
+    std::vector<std::string> partitions = {"sys_b", "vnd_b", "prd_b", "dlkm_b"};
+    for (const auto& name : partitions) {
+        ASSERT_TRUE(IsPartitionUnchanged(name));
+    }
+
+    ASSERT_TRUE(init->PerformInitTransition(SnapshotManager::InitTransition::SECOND_STAGE));
+    for (const auto& name : partitions) {
+        ASSERT_TRUE(init->snapuserd_client()->WaitForDeviceDelete(name + "-user-cow-init"));
+    }
+
+    // Initiate the merge and wait for it to be completed.
+    ASSERT_TRUE(init->InitiateMerge());
+    ASSERT_EQ(UpdateState::MergeCompleted, init->ProcessUpdateState());
+
+    // Check that the target partitions have the same content after the merge.
+    for (const auto& name : {"sys_b", "vnd_b", "prd_b", "dlkm_b"}) {
+        ASSERT_TRUE(IsPartitionUnchanged(name))
+                << "Content of " << name << " changes after the merge";
+    }
+}
+
+class AutoKill final {
+  public:
+    explicit AutoKill(pid_t pid) : pid_(pid) {}
+    ~AutoKill() {
+        if (pid_ > 0) kill(pid_, SIGKILL);
+    }
+
+    bool valid() const { return pid_ > 0; }
+
+  private:
+    pid_t pid_;
+};
+
+TEST_F(SnapshotUpdateTest, DaemonTransition) {
+    if (!IsCompressionEnabled()) {
+        GTEST_SKIP() << "Skipping Virtual A/B Compression test";
+    }
+
+    // Ensure a connection to the second-stage daemon, but use the first-stage
+    // code paths thereafter.
+    ASSERT_TRUE(sm->EnsureSnapuserdConnected());
+    sm->set_use_first_stage_snapuserd(true);
+
+    AddOperationForPartitions();
+    // Execute the update.
+    ASSERT_TRUE(sm->BeginUpdate());
+    ASSERT_TRUE(sm->CreateUpdateSnapshots(manifest_));
+    ASSERT_TRUE(MapUpdateSnapshots());
+    ASSERT_TRUE(sm->FinishedSnapshotWrites(false));
+    ASSERT_TRUE(UnmapAll());
+
+    auto init = NewManagerForFirstStageMount("_b");
+    ASSERT_NE(init, nullptr);
+
+    ASSERT_TRUE(init->EnsureSnapuserdConnected());
+    init->set_use_first_stage_snapuserd(true);
+
+    ASSERT_TRUE(init->NeedSnapshotsInFirstStageMount());
+    ASSERT_TRUE(init->CreateLogicalAndSnapshotPartitions("super", snapshot_timeout_));
+
+    ASSERT_EQ(access("/dev/dm-user/sys_b-init", F_OK), 0);
+    ASSERT_EQ(access("/dev/dm-user/sys_b", F_OK), -1);
+
+    ASSERT_TRUE(init->PerformInitTransition(SnapshotManager::InitTransition::SECOND_STAGE));
+
+    // :TODO: this is a workaround to ensure the handler list stays empty. We
+    // should make this test more like actual init, and spawn two copies of
+    // snapuserd, given how many other tests we now have for normal snapuserd.
+    ASSERT_TRUE(init->snapuserd_client()->WaitForDeviceDelete("sys_b-init"));
+    ASSERT_TRUE(init->snapuserd_client()->WaitForDeviceDelete("vnd_b-init"));
+    ASSERT_TRUE(init->snapuserd_client()->WaitForDeviceDelete("prd_b-init"));
+
+    // The control device should have been renamed.
+    ASSERT_TRUE(android::fs_mgr::WaitForFileDeleted("/dev/dm-user/sys_b-init", 10s));
+    ASSERT_EQ(access("/dev/dm-user/sys_b", F_OK), 0);
+}
+
+TEST_F(SnapshotUpdateTest, MapAllSnapshots) {
+    AddOperationForPartitions();
+    // Execute the update.
+    ASSERT_TRUE(sm->BeginUpdate());
+    ASSERT_TRUE(sm->CreateUpdateSnapshots(manifest_));
+    for (const auto& name : {"sys_b", "vnd_b", "prd_b"}) {
+        ASSERT_TRUE(WriteSnapshotAndHash(name));
+    }
+    ASSERT_TRUE(sm->FinishedSnapshotWrites(false));
+    ASSERT_TRUE(sm->MapAllSnapshots(10s));
+
+    // Read bytes back and verify they match the cache.
+    ASSERT_TRUE(IsPartitionUnchanged("sys_b"));
+
+    ASSERT_TRUE(sm->UnmapAllSnapshots());
+}
+
+TEST_F(SnapshotUpdateTest, CancelOnTargetSlot) {
+    AddOperationForPartitions();
+
+    // Execute the update from B->A.
+    test_device->set_slot_suffix("_b");
+    ASSERT_TRUE(sm->BeginUpdate());
+    ASSERT_TRUE(sm->CreateUpdateSnapshots(manifest_));
+
+    ASSERT_TRUE(UnmapAll());
+    std::string path;
+    ASSERT_TRUE(CreateLogicalPartition(
+            CreateLogicalPartitionParams{
+                    .block_device = fake_super,
+                    .metadata_slot = 0,
+                    .partition_name = "sys_a",
+                    .timeout_ms = 1s,
+                    .partition_opener = opener_.get(),
+            },
+            &path));
+
+    // Switch back to "A", make sure we can cancel. Instead of unmapping sys_a
+    // we should simply delete the old snapshots.
+    test_device->set_slot_suffix("_a");
+    ASSERT_TRUE(sm->BeginUpdate());
+}
+
+class FlashAfterUpdateTest : public SnapshotUpdateTest,
+                             public WithParamInterface<std::tuple<uint32_t, bool>> {
+  public:
+    AssertionResult InitiateMerge(const std::string& slot_suffix) {
+        auto sm = SnapshotManager::New(new TestDeviceInfo(fake_super, slot_suffix));
+        if (!sm->CreateLogicalAndSnapshotPartitions("super", snapshot_timeout_)) {
+            return AssertionFailure() << "Cannot CreateLogicalAndSnapshotPartitions";
+        }
+        if (!sm->InitiateMerge()) {
+            return AssertionFailure() << "Cannot initiate merge";
+        }
+        return AssertionSuccess();
+    }
+};
+
+TEST_P(FlashAfterUpdateTest, FlashSlotAfterUpdate) {
+    // Execute the update.
+    ASSERT_TRUE(sm->BeginUpdate());
+    ASSERT_TRUE(sm->CreateUpdateSnapshots(manifest_));
+    ASSERT_TRUE(MapUpdateSnapshots());
+    ASSERT_TRUE(sm->FinishedSnapshotWrites(false));
+
+    // Simulate shutting down the device.
+    ASSERT_TRUE(UnmapAll());
+
+    bool after_merge = std::get<1>(GetParam());
+    if (after_merge) {
+        ASSERT_TRUE(InitiateMerge("_b"));
+        // Simulate shutting down the device after merge has initiated.
+        ASSERT_TRUE(UnmapAll());
+    }
+
+    auto flashed_slot = std::get<0>(GetParam());
+    auto flashed_slot_suffix = SlotSuffixForSlotNumber(flashed_slot);
+
+    // Simulate flashing |flashed_slot|. This clears the UPDATED flag.
+    auto flashed_builder = MetadataBuilder::New(*opener_, "super", flashed_slot);
+    ASSERT_NE(flashed_builder, nullptr);
+    flashed_builder->RemoveGroupAndPartitions(group_->name() + flashed_slot_suffix);
+    flashed_builder->RemoveGroupAndPartitions(kCowGroupName);
+    ASSERT_TRUE(FillFakeMetadata(flashed_builder.get(), manifest_, flashed_slot_suffix));
+
+    // Deliberately remove a partition from this build so that
+    // InitiateMerge do not switch state to "merging". This is possible in
+    // practice because the list of dynamic partitions may change.
+    ASSERT_NE(nullptr, flashed_builder->FindPartition("prd" + flashed_slot_suffix));
+    flashed_builder->RemovePartition("prd" + flashed_slot_suffix);
+
+    // Note that fastbootd always updates the partition table of both slots.
+    auto flashed_metadata = flashed_builder->Export();
+    ASSERT_NE(nullptr, flashed_metadata);
+    ASSERT_TRUE(UpdatePartitionTable(*opener_, "super", *flashed_metadata, 0));
+    ASSERT_TRUE(UpdatePartitionTable(*opener_, "super", *flashed_metadata, 1));
+
+    std::string path;
+    for (const auto& name : {"sys", "vnd"}) {
+        ASSERT_TRUE(CreateLogicalPartition(
+                CreateLogicalPartitionParams{
+                        .block_device = fake_super,
+                        .metadata_slot = flashed_slot,
+                        .partition_name = name + flashed_slot_suffix,
+                        .timeout_ms = 1s,
+                        .partition_opener = opener_.get(),
+                },
+                &path));
+        ASSERT_TRUE(WriteRandomData(path));
+        auto hash = GetHash(path);
+        ASSERT_TRUE(hash.has_value());
+        hashes_[name + flashed_slot_suffix] = *hash;
+    }
+
+    // Simulate shutting down the device after flash.
+    ASSERT_TRUE(UnmapAll());
+
+    // Simulate reboot. After reboot, init does first stage mount.
+    auto init = NewManagerForFirstStageMount(flashed_slot_suffix);
+    ASSERT_NE(init, nullptr);
+
+    if (flashed_slot && after_merge) {
+        ASSERT_TRUE(init->NeedSnapshotsInFirstStageMount());
+    }
+    ASSERT_TRUE(init->CreateLogicalAndSnapshotPartitions("super", snapshot_timeout_));
+
+    // Check that the target partitions have the same content.
+    for (const auto& name : {"sys", "vnd"}) {
+        ASSERT_TRUE(IsPartitionUnchanged(name + flashed_slot_suffix));
+    }
+
+    // There should be no snapshot to merge.
+    auto new_sm = SnapshotManager::New(new TestDeviceInfo(fake_super, flashed_slot_suffix));
+    if (flashed_slot == 0 && after_merge) {
+        ASSERT_EQ(UpdateState::MergeCompleted, new_sm->ProcessUpdateState());
+    } else {
+        // update_engine calls ProcessUpdateState first -- should see Cancelled.
+        ASSERT_EQ(UpdateState::Cancelled, new_sm->ProcessUpdateState());
+    }
+
+    // Next OTA calls CancelUpdate no matter what.
+    ASSERT_TRUE(new_sm->CancelUpdate());
+}
+
+INSTANTIATE_TEST_SUITE_P(Snapshot, FlashAfterUpdateTest, Combine(Values(0, 1), Bool()),
+                         [](const TestParamInfo<FlashAfterUpdateTest::ParamType>& info) {
+                             return "Flash"s + (std::get<0>(info.param) ? "New"s : "Old"s) +
+                                    "Slot"s + (std::get<1>(info.param) ? "After"s : "Before"s) +
+                                    "Merge"s;
+                         });
+
+// Test behavior of ImageManager::Create on low space scenario. These tests assumes image manager
+// uses /data as backup device.
+class ImageManagerTest : public SnapshotTest, public WithParamInterface<uint64_t> {
+  protected:
+    void SetUp() override {
+        SKIP_IF_NON_VIRTUAL_AB();
+        SnapshotTest::SetUp();
+        userdata_ = std::make_unique<LowSpaceUserdata>();
+        ASSERT_TRUE(userdata_->Init(GetParam()));
+    }
+    void TearDown() override {
+        RETURN_IF_NON_VIRTUAL_AB();
+        return;  // BUG(149738928)
+
+        EXPECT_TRUE(!image_manager_->BackingImageExists(kImageName) ||
+                    image_manager_->DeleteBackingImage(kImageName));
+    }
+    static constexpr const char* kImageName = "my_image";
+    std::unique_ptr<LowSpaceUserdata> userdata_;
+};
+
+TEST_P(ImageManagerTest, CreateImageEnoughAvailSpace) {
+    if (userdata_->available_space() == 0) {
+        GTEST_SKIP() << "/data is full (" << userdata_->available_space()
+                     << " bytes available), skipping";
+    }
+    ASSERT_TRUE(image_manager_->CreateBackingImage(kImageName, userdata_->available_space(),
+                                                   IImageManager::CREATE_IMAGE_DEFAULT))
+            << "Should be able to create image with size = " << userdata_->available_space()
+            << " bytes";
+    ASSERT_TRUE(image_manager_->DeleteBackingImage(kImageName))
+            << "Should be able to delete created image";
+}
+
+TEST_P(ImageManagerTest, CreateImageNoSpace) {
+    uint64_t to_allocate = userdata_->free_space() + userdata_->bsize();
+    auto res = image_manager_->CreateBackingImage(kImageName, to_allocate,
+                                                  IImageManager::CREATE_IMAGE_DEFAULT);
+    ASSERT_FALSE(res) << "Should not be able to create image with size = " << to_allocate
+                      << " bytes because only " << userdata_->free_space() << " bytes are free";
+    ASSERT_EQ(FiemapStatus::ErrorCode::NO_SPACE, res.error_code()) << res.string();
+}
+
+std::vector<uint64_t> ImageManagerTestParams() {
+    std::vector<uint64_t> ret;
+    for (uint64_t size = 1_MiB; size <= 512_MiB; size *= 2) {
+        ret.push_back(size);
+    }
+    return ret;
+}
+
+INSTANTIATE_TEST_SUITE_P(ImageManagerTest, ImageManagerTest, ValuesIn(ImageManagerTestParams()));
+
+bool Mkdir(const std::string& path) {
+    if (mkdir(path.c_str(), 0700) && errno != EEXIST) {
+        std::cerr << "Could not mkdir " << path << ": " << strerror(errno) << std::endl;
+        return false;
+    }
+    return true;
+}
+
+class SnapshotTestEnvironment : public ::testing::Environment {
+  public:
+    ~SnapshotTestEnvironment() override {}
+    void SetUp() override;
+    void TearDown() override;
+
+  private:
+    bool CreateFakeSuper();
+
+    std::unique_ptr<IImageManager> super_images_;
+};
+
+bool SnapshotTestEnvironment::CreateFakeSuper() {
+    // Create and map the fake super partition.
+    static constexpr int kImageFlags =
+            IImageManager::CREATE_IMAGE_DEFAULT | IImageManager::CREATE_IMAGE_ZERO_FILL;
+    if (!super_images_->CreateBackingImage("fake-super", kSuperSize, kImageFlags)) {
+        LOG(ERROR) << "Could not create fake super partition";
+        return false;
+    }
+    if (!super_images_->MapImageDevice("fake-super", 10s, &fake_super)) {
+        LOG(ERROR) << "Could not map fake super partition";
+        return false;
+    }
+    test_device->set_fake_super(fake_super);
+    return true;
+}
+
+void SnapshotTestEnvironment::SetUp() {
+    // b/163082876: GTEST_SKIP in Environment will make atest report incorrect results. Until
+    // that is fixed, don't call GTEST_SKIP here, but instead call GTEST_SKIP in individual test
+    // suites.
+    RETURN_IF_NON_VIRTUAL_AB_MSG("Virtual A/B is not enabled, skipping global setup.\n");
+
+    std::vector<std::string> paths = {
+            // clang-format off
+            "/data/gsi/ota/test",
+            "/data/gsi/ota/test/super",
+            "/metadata/gsi/ota/test",
+            "/metadata/gsi/ota/test/super",
+            "/metadata/ota/test",
+            "/metadata/ota/test/snapshots",
+            // clang-format on
+    };
+    for (const auto& path : paths) {
+        ASSERT_TRUE(Mkdir(path));
+    }
+
+    // Create this once, otherwise, gsid will start/stop between each test.
+    test_device = new TestDeviceInfo();
+    sm = SnapshotManager::New(test_device);
+    ASSERT_NE(nullptr, sm) << "Could not create snapshot manager";
+
+    // Use a separate image manager for our fake super partition.
+    super_images_ = IImageManager::Open("ota/test/super", 10s);
+    ASSERT_NE(nullptr, super_images_) << "Could not create image manager";
+
+    // Map the old image if one exists so we can safely unmap everything that
+    // depends on it.
+    bool recreate_fake_super;
+    if (super_images_->BackingImageExists("fake-super")) {
+        if (super_images_->IsImageMapped("fake-super")) {
+            ASSERT_TRUE(super_images_->GetMappedImageDevice("fake-super", &fake_super));
+        } else {
+            ASSERT_TRUE(super_images_->MapImageDevice("fake-super", 10s, &fake_super));
+        }
+        test_device->set_fake_super(fake_super);
+        recreate_fake_super = true;
+    } else {
+        ASSERT_TRUE(CreateFakeSuper());
+        recreate_fake_super = false;
+    }
+
+    // Clean up previous run.
+    MetadataMountedTest().TearDown();
+    SnapshotUpdateTest().Cleanup();
+    SnapshotTest().Cleanup();
+
+    if (recreate_fake_super) {
+        // Clean up any old copy.
+        DeleteBackingImage(super_images_.get(), "fake-super");
+        ASSERT_TRUE(CreateFakeSuper());
+    }
+}
+
+void SnapshotTestEnvironment::TearDown() {
+    RETURN_IF_NON_VIRTUAL_AB();
+    if (super_images_ != nullptr) {
+        DeleteBackingImage(super_images_.get(), "fake-super");
+    }
+}
+
+}  // namespace snapshot
+}  // namespace android
+
+int main(int argc, char** argv) {
+    ::testing::InitGoogleTest(&argc, argv);
+    ::testing::AddGlobalTestEnvironment(new ::android::snapshot::SnapshotTestEnvironment());
+
+    android::base::SetProperty("ctl.stop", "snapuserd");
+    android::base::SetProperty("snapuserd.test.dm.snapshots", "0");
+
+    return RUN_ALL_TESTS();
+}
diff --git a/fs_mgr/libsnapshot/utility.cpp b/fs_mgr/libsnapshot/utility.cpp
index 4a2af1c10..89d614599 100644
--- a/fs_mgr/libsnapshot/utility.cpp
+++ b/fs_mgr/libsnapshot/utility.cpp
@@ -22,6 +22,7 @@
 
 #include <android-base/file.h>
 #include <android-base/logging.h>
+#include <android-base/parseint.h>
 #include <android-base/properties.h>
 #include <android-base/strings.h>
 #include <fs_mgr/roots.h>
@@ -187,6 +188,10 @@ bool IsCompressionEnabled() {
     return android::base::GetBoolProperty("ro.virtual_ab.compression.enabled", false);
 }
 
+bool IsUserspaceSnapshotsEnabled() {
+    return android::base::GetBoolProperty("ro.virtual_ab.userspace.snapshots.enabled", false);
+}
+
 std::string GetOtherPartitionName(const std::string& name) {
     auto suffix = android::fs_mgr::GetPartitionSlotSuffix(name);
     CHECK(suffix == "_a" || suffix == "_b");
@@ -195,5 +200,9 @@ std::string GetOtherPartitionName(const std::string& name) {
     return name.substr(0, name.size() - suffix.size()) + other_suffix;
 }
 
+bool IsDmSnapshotTestingEnabled() {
+    return android::base::GetBoolProperty("snapuserd.test.dm.snapshots", false);
+}
+
 }  // namespace snapshot
 }  // namespace android
diff --git a/fs_mgr/libsnapshot/utility.h b/fs_mgr/libsnapshot/utility.h
index e97afed22..a032b6869 100644
--- a/fs_mgr/libsnapshot/utility.h
+++ b/fs_mgr/libsnapshot/utility.h
@@ -131,8 +131,11 @@ void AppendExtent(google::protobuf::RepeatedPtrField<chromeos_update_engine::Ext
 
 bool IsCompressionEnabled();
 
+bool IsUserspaceSnapshotsEnabled();
+
+bool IsDmSnapshotTestingEnabled();
+
 // Swap the suffix of a partition name.
 std::string GetOtherPartitionName(const std::string& name);
-
 }  // namespace snapshot
 }  // namespace android
diff --git a/init/first_stage_mount.cpp b/init/first_stage_mount.cpp
index f5c10bbd4..10a07a1cc 100644
--- a/init/first_stage_mount.cpp
+++ b/init/first_stage_mount.cpp
@@ -391,7 +391,11 @@ bool FirstStageMount::CreateSnapshotPartitions(SnapshotManager* sm) {
 
     use_snapuserd_ = sm->IsSnapuserdRequired();
     if (use_snapuserd_) {
-        LaunchFirstStageSnapuserd();
+        if (sm->UpdateUsesUserSnapshots()) {
+            LaunchFirstStageSnapuserd(SnapshotDriver::DM_USER);
+        } else {
+            LaunchFirstStageSnapuserd(SnapshotDriver::DM_SNAPSHOT);
+        }
     }
 
     sm->SetUeventRegenCallback([this](const std::string& device) -> bool {
diff --git a/init/snapuserd_transition.cpp b/init/snapuserd_transition.cpp
index b8c2fd2f3..e11510ee8 100644
--- a/init/snapuserd_transition.cpp
+++ b/init/snapuserd_transition.cpp
@@ -58,7 +58,7 @@ static constexpr char kSnapuserdFirstStageInfoVar[] = "FIRST_STAGE_SNAPUSERD_INF
 static constexpr char kSnapuserdLabel[] = "u:object_r:snapuserd_exec:s0";
 static constexpr char kSnapuserdSocketLabel[] = "u:object_r:snapuserd_socket:s0";
 
-void LaunchFirstStageSnapuserd() {
+void LaunchFirstStageSnapuserd(SnapshotDriver driver) {
     SocketDescriptor socket_desc;
     socket_desc.name = android::snapshot::kSnapuserdSocket;
     socket_desc.type = SOCK_STREAM;
@@ -80,12 +80,23 @@ void LaunchFirstStageSnapuserd() {
     }
     if (pid == 0) {
         socket->Publish();
-        char arg0[] = "/system/bin/snapuserd";
-        char* const argv[] = {arg0, nullptr};
-        if (execv(arg0, argv) < 0) {
-            PLOG(FATAL) << "Cannot launch snapuserd; execv failed";
+
+        if (driver == SnapshotDriver::DM_USER) {
+            char arg0[] = "/system/bin/snapuserd";
+            char arg1[] = "-user_snapshot";
+            char* const argv[] = {arg0, arg1, nullptr};
+            if (execv(arg0, argv) < 0) {
+                PLOG(FATAL) << "Cannot launch snapuserd; execv failed";
+            }
+            _exit(127);
+        } else {
+            char arg0[] = "/system/bin/snapuserd";
+            char* const argv[] = {arg0, nullptr};
+            if (execv(arg0, argv) < 0) {
+                PLOG(FATAL) << "Cannot launch snapuserd; execv failed";
+            }
+            _exit(127);
         }
-        _exit(127);
     }
 
     auto client = SnapuserdClient::Connect(android::snapshot::kSnapuserdSocket, 10s);
diff --git a/init/snapuserd_transition.h b/init/snapuserd_transition.h
index 62aee83f6..be22afd20 100644
--- a/init/snapuserd_transition.h
+++ b/init/snapuserd_transition.h
@@ -29,8 +29,13 @@
 namespace android {
 namespace init {
 
+enum class SnapshotDriver {
+    DM_SNAPSHOT,
+    DM_USER,
+};
+
 // Fork and exec a new copy of snapuserd.
-void LaunchFirstStageSnapuserd();
+void LaunchFirstStageSnapuserd(SnapshotDriver driver);
 
 class SnapuserdSelinuxHelper final {
     using SnapshotManager = android::snapshot::SnapshotManager;