diff --git a/adb/client/adb_client.cpp b/adb/client/adb_client.cpp
index 0f621eb..d91ae35 100644
--- a/adb/client/adb_client.cpp
+++ b/adb/client/adb_client.cpp
@@ -400,9 +400,15 @@
 }
 
 bool adb_get_feature_set(FeatureSet* feature_set, std::string* error) {
-    std::string result;
-    if (adb_query(format_host_command("features"), &result, error)) {
-        *feature_set = StringToFeatureSet(result);
+    static FeatureSet* features = nullptr;
+    if (!features) {
+        std::string result;
+        if (adb_query(format_host_command("features"), &result, error)) {
+            features = new FeatureSet(StringToFeatureSet(result));
+        }
+    }
+    if (features) {
+        *feature_set = *features;
         return true;
     }
     feature_set->clear();
diff --git a/adb/client/adb_install.cpp b/adb/client/adb_install.cpp
index f1f080a..a6e8998 100644
--- a/adb/client/adb_install.cpp
+++ b/adb/client/adb_install.cpp
@@ -16,6 +16,7 @@
 
 #include "adb_install.h"
 
+#include <fcntl.h>
 #include <stdio.h>
 #include <stdlib.h>
 #include <unistd.h>
@@ -40,18 +41,31 @@
 static constexpr int kFastDeployMinApi = 24;
 #endif
 
+namespace {
+
+enum InstallMode {
+    INSTALL_DEFAULT,
+    INSTALL_PUSH,
+    INSTALL_STREAM,
+};
+
+}
+
 static bool can_use_feature(const char* feature) {
     FeatureSet features;
     std::string error;
     if (!adb_get_feature_set(&features, &error)) {
         fprintf(stderr, "error: %s\n", error.c_str());
-        return true;
+        return false;
     }
     return CanUseFeature(features, feature);
 }
 
-static bool use_legacy_install() {
-    return !can_use_feature(kFeatureCmd);
+static InstallMode best_install_mode() {
+    if (can_use_feature(kFeatureCmd)) {
+        return INSTALL_STREAM;
+    }
+    return INSTALL_PUSH;
 }
 
 static bool is_apex_supported() {
@@ -112,7 +126,7 @@
 }
 
 int uninstall_app(int argc, const char** argv) {
-    if (use_legacy_install()) {
+    if (best_install_mode() == INSTALL_PUSH) {
         return uninstall_app_legacy(argc, argv);
     }
     return uninstall_app_streamed(argc, argv);
@@ -200,32 +214,49 @@
             return 1;
         }
 
+#ifdef __linux__
+        posix_fadvise(local_fd.get(), 0, 0, POSIX_FADV_SEQUENTIAL | POSIX_FADV_NOREUSE);
+#endif
+
+        const bool use_abb = can_use_feature(kFeatureAbb);
         std::string error;
-        std::string cmd = "exec:cmd package";
+        std::vector<std::string> cmd_args = {use_abb ? "package" : "exec:cmd package"};
+        cmd_args.reserve(argc + 3);
 
         // don't copy the APK name, but, copy the rest of the arguments as-is
         while (argc-- > 1) {
-            cmd += " " + escape_arg(std::string(*argv++));
+            if (use_abb) {
+                cmd_args.push_back(*argv++);
+            } else {
+                cmd_args.push_back(escape_arg(*argv++));
+            }
         }
 
         // add size parameter [required for streaming installs]
         // do last to override any user specified value
-        cmd += " " + android::base::StringPrintf("-S %" PRIu64, static_cast<uint64_t>(sb.st_size));
+        cmd_args.push_back("-S");
+        cmd_args.push_back(
+                android::base::StringPrintf("%" PRIu64, static_cast<uint64_t>(sb.st_size)));
 
         if (is_apex) {
-            cmd += " --apex";
+            cmd_args.push_back("--apex");
         }
 
-        unique_fd remote_fd(adb_connect(cmd, &error));
+        unique_fd remote_fd;
+        if (use_abb) {
+            remote_fd = send_abb_exec_command(cmd_args, &error);
+        } else {
+            remote_fd.reset(adb_connect(android::base::Join(cmd_args, " "), &error));
+        }
         if (remote_fd < 0) {
             fprintf(stderr, "adb: connect error for write: %s\n", error.c_str());
             return 1;
         }
 
-        char buf[BUFSIZ];
         copy_to_file(local_fd.get(), remote_fd.get());
-        read_status_line(remote_fd.get(), buf, sizeof(buf));
 
+        char buf[BUFSIZ];
+        read_status_line(remote_fd.get(), buf, sizeof(buf));
         if (!strncmp("Success", buf, 7)) {
             fputs(buf, stdout);
             return 0;
@@ -256,8 +287,7 @@
 
     int result = -1;
     std::vector<const char*> apk_file = {argv[last_apk]};
-    std::string apk_dest =
-            "/data/local/tmp/" + android::base::Basename(argv[last_apk]);
+    std::string apk_dest = "/data/local/tmp/" + android::base::Basename(argv[last_apk]);
 
     if (use_fastdeploy == true) {
 #if defined(ENABLE_FASTDEPLOY)
@@ -292,11 +322,7 @@
 
 int install_app(int argc, const char** argv) {
     std::vector<int> processedArgIndicies;
-    enum installMode {
-        INSTALL_DEFAULT,
-        INSTALL_PUSH,
-        INSTALL_STREAM
-    } installMode = INSTALL_DEFAULT;
+    InstallMode installMode = INSTALL_DEFAULT;
     bool use_fastdeploy = false;
     bool is_reinstall = false;
     bool use_localagent = false;
@@ -337,14 +363,10 @@
     }
 
     if (installMode == INSTALL_DEFAULT) {
-        if (use_legacy_install()) {
-            installMode = INSTALL_PUSH;
-        } else {
-            installMode = INSTALL_STREAM;
-        }
+        installMode = best_install_mode();
     }
 
-    if (installMode == INSTALL_STREAM && use_legacy_install() == true) {
+    if (installMode == INSTALL_STREAM && best_install_mode() == INSTALL_PUSH) {
         error_exit("Attempting to use streaming install on unsupported device");
     }
 
@@ -420,7 +442,7 @@
     if (first_apk == -1) error_exit("need APK file on command line");
 
     std::string install_cmd;
-    if (use_legacy_install()) {
+    if (best_install_mode() == INSTALL_PUSH) {
         install_cmd = "exec:pm";
     } else {
         install_cmd = "exec:cmd package";
@@ -545,7 +567,7 @@
 
     if (first_package == -1) error_exit("need APK or APEX files on command line");
 
-    if (use_legacy_install()) {
+    if (best_install_mode() == INSTALL_PUSH) {
         fprintf(stderr, "adb: multi-package install is not supported on this device\n");
         return EXIT_FAILURE;
     }
diff --git a/adb/client/commandline.cpp b/adb/client/commandline.cpp
index 7d0eb40..599e0e6 100644
--- a/adb/client/commandline.cpp
+++ b/adb/client/commandline.cpp
@@ -357,7 +357,7 @@
 }
 
 void copy_to_file(int inFd, int outFd) {
-    std::vector<char> buf(32 * 1024);
+    std::vector<char> buf(64 * 1024);
     int len;
     long total = 0;
     int old_stdin_mode = -1;
diff --git a/adb/client/commandline.h b/adb/client/commandline.h
index 6cfd4f7..cd5933a 100644
--- a/adb/client/commandline.h
+++ b/adb/client/commandline.h
@@ -17,7 +17,11 @@
 #ifndef COMMANDLINE_H
 #define COMMANDLINE_H
 
+#include <android-base/strings.h>
+
 #include "adb.h"
+#include "adb_client.h"
+#include "adb_unique_fd.h"
 
 // Callback used to handle the standard streams (stdout and stderr) sent by the
 // device's upon receiving a command.
@@ -105,4 +109,17 @@
         const std::string& command, bool disable_shell_protocol = false,
         StandardStreamsCallbackInterface* callback = &DEFAULT_STANDARD_STREAMS_CALLBACK);
 
+// Connects to the device "abb" service with |command| and returns the fd.
+template <typename ContainerT>
+unique_fd send_abb_exec_command(const ContainerT& command_args, std::string* error) {
+    std::string service_string = "abb_exec:" + android::base::Join(command_args, ABB_ARG_DELIMETER);
+
+    unique_fd fd(adb_connect(service_string, error));
+    if (fd < 0) {
+        fprintf(stderr, "adb: failed to run abb_exec. Error: %s\n", error->c_str());
+        return unique_fd{};
+    }
+    return fd;
+}
+
 #endif  // COMMANDLINE_H
diff --git a/adb/daemon/abb.cpp b/adb/daemon/abb.cpp
index 425438e..17c25e8 100644
--- a/adb/daemon/abb.cpp
+++ b/adb/daemon/abb.cpp
@@ -24,6 +24,7 @@
 #include "adb_io.h"
 #include "adb_utils.h"
 #include "shell_service.h"
+#include "sysdeps.h"
 
 namespace {
 
@@ -69,6 +70,11 @@
 }  // namespace
 
 static int execCmd(std::string_view args, borrowed_fd in, borrowed_fd out, borrowed_fd err) {
+    int max_buf = LINUX_MAX_SOCKET_SIZE;
+    adb_setsockopt(in, SOL_SOCKET, SO_SNDBUF, &max_buf, sizeof(max_buf));
+    adb_setsockopt(out, SOL_SOCKET, SO_SNDBUF, &max_buf, sizeof(max_buf));
+    adb_setsockopt(err, SOL_SOCKET, SO_SNDBUF, &max_buf, sizeof(max_buf));
+
     AdbFdTextOutput oin(out);
     AdbFdTextOutput oerr(err);
     return cmdMain(parseCmdArgs(args), oin, oerr, in.get(), out.get(), err.get(),
@@ -98,6 +104,8 @@
         }
 
         unique_fd result = StartCommandInProcess(std::string(name), &execCmd, protocol);
+        int max_buf = LINUX_MAX_SOCKET_SIZE;
+        adb_setsockopt(result, SOL_SOCKET, SO_SNDBUF, &max_buf, sizeof(max_buf));
         if (android::base::SendFileDescriptors(fd, "", 1, result.get()) != 1) {
             PLOG(ERROR) << "Failed to send an inprocess fd for command: " << data;
             break;
diff --git a/adb/daemon/abb_service.cpp b/adb/daemon/abb_service.cpp
index a435279..e1df4a5 100644
--- a/adb/daemon/abb_service.cpp
+++ b/adb/daemon/abb_service.cpp
@@ -53,14 +53,13 @@
             return error_fd;
         }
 
-        if (!SendProtocolString(socket_fd_, std::string(command))) {
+        if (!SendProtocolString(socket_fd_, command)) {
             PLOG(ERROR) << "failed to send command to abb";
             socket_fd_.reset();
             continue;
         }
 
         unique_fd fd;
-        std::string error;
         char buf;
         if (android::base::ReceiveFileDescriptors(socket_fd_, &buf, 1, &fd) != 1) {
             PLOG(ERROR) << "failed to receive FD from abb";
diff --git a/fs_mgr/fs_mgr_overlayfs.cpp b/fs_mgr/fs_mgr_overlayfs.cpp
index ac15ce4..4ee7db9 100644
--- a/fs_mgr/fs_mgr_overlayfs.cpp
+++ b/fs_mgr/fs_mgr_overlayfs.cpp
@@ -150,6 +150,31 @@
     return (vst.f_bfree >= (vst.f_blocks * kPercentThreshold / 100));
 }
 
+const auto kPhysicalDevice = "/dev/block/by-name/"s;
+
+bool fs_mgr_update_blk_device(FstabEntry* entry) {
+    if (entry->fs_mgr_flags.logical) {
+        fs_mgr_update_logical_partition(entry);
+    }
+    if (fs_mgr_access(entry->blk_device)) {
+        return true;
+    }
+    if (entry->blk_device != "/dev/root") {
+        return false;
+    }
+
+    // special case for system-as-root (taimen and others)
+    auto blk_device = kPhysicalDevice + "system";
+    if (!fs_mgr_access(blk_device)) {
+        blk_device += fs_mgr_get_slot_suffix();
+        if (!fs_mgr_access(blk_device)) {
+            return false;
+        }
+    }
+    entry->blk_device = blk_device;
+    return true;
+}
+
 bool fs_mgr_overlayfs_enabled(FstabEntry* entry) {
     // readonly filesystem, can not be mount -o remount,rw
     // for squashfs, erofs or if free space is (near) zero making such a remount
@@ -157,19 +182,19 @@
     if (!fs_mgr_filesystem_has_space(entry->mount_point)) {
         return true;
     }
-    if (entry->fs_mgr_flags.logical) {
-        fs_mgr_update_logical_partition(entry);
+
+    // blk_device needs to be setup so we can check superblock.
+    // If we fail here, because during init first stage and have doubts.
+    if (!fs_mgr_update_blk_device(entry)) {
+        return true;
     }
+
+    // check if ext4 de-dupe
     auto save_errno = errno;
-    errno = 0;
     auto has_shared_blocks = fs_mgr_has_shared_blocks(entry->mount_point, entry->blk_device);
     if (!has_shared_blocks && (entry->mount_point == "/system")) {
         has_shared_blocks = fs_mgr_has_shared_blocks("/", entry->blk_device);
     }
-    // special case for first stage init for system as root (taimen)
-    if (!has_shared_blocks && (errno == ENOENT) && (entry->blk_device == "/dev/root")) {
-        has_shared_blocks = true;
-    }
     errno = save_errno;
     return has_shared_blocks;
 }
@@ -388,8 +413,6 @@
     return SlotNumberForSlotSuffix(fs_mgr_get_slot_suffix());
 }
 
-const auto kPhysicalDevice = "/dev/block/by-name/"s;
-
 std::string fs_mgr_overlayfs_super_device(uint32_t slot_number) {
     return kPhysicalDevice + fs_mgr_get_super_partition_name(slot_number);
 }
diff --git a/fs_mgr/libdm/include/libdm/dm.h b/fs_mgr/libdm/include/libdm/dm.h
index f5783cb..cf306f3 100644
--- a/fs_mgr/libdm/include/libdm/dm.h
+++ b/fs_mgr/libdm/include/libdm/dm.h
@@ -197,6 +197,7 @@
     struct TargetInfo {
         struct dm_target_spec spec;
         std::string data;
+        TargetInfo() {}
         TargetInfo(const struct dm_target_spec& spec, const std::string& data)
             : spec(spec), data(data) {}
     };
diff --git a/fs_mgr/libsnapshot/include/libsnapshot/snapshot.h b/fs_mgr/libsnapshot/include/libsnapshot/snapshot.h
index 062e00b..f7608dc 100644
--- a/fs_mgr/libsnapshot/include/libsnapshot/snapshot.h
+++ b/fs_mgr/libsnapshot/include/libsnapshot/snapshot.h
@@ -22,7 +22,8 @@
 #include <vector>
 
 #include <android-base/unique_fd.h>
-#include <libdm/dm_target.h>
+#include <libdm/dm.h>
+#include <libfiemap/image_manager.h>
 
 #ifndef FRIEND_TEST
 #define FRIEND_TEST(test_set_name, individual_test) \
@@ -38,7 +39,7 @@
 
 namespace snapshot {
 
-enum class UpdateState {
+enum class UpdateState : unsigned int {
     // No update or merge is in progress.
     None,
 
@@ -51,6 +52,10 @@
     // The kernel is merging in the background.
     Merging,
 
+    // Post-merge cleanup steps could not be completed due to a transient
+    // error, but the next reboot will finish any pending operations.
+    MergeNeedsReboot,
+
     // Merging is complete, and needs to be acknowledged.
     MergeCompleted,
 
@@ -94,8 +99,23 @@
 
     // Wait for the current merge to finish, then perform cleanup when it
     // completes. It is necessary to call this after InitiateMerge(), or when
-    // a merge is detected for the first time after boot.
-    bool WaitForMerge();
+    // a merge state is detected during boot.
+    //
+    // Note that after calling WaitForMerge(), GetUpdateState() may still return
+    // that a merge is in progress:
+    //   MergeFailed indicates that a fatal error occurred. WaitForMerge() may
+    //   called any number of times again to attempt to make more progress, but
+    //   we do not expect it to succeed if a catastrophic error occurred.
+    //
+    //   MergeNeedsReboot indicates that the merge has completed, but cleanup
+    //   failed. This can happen if for some reason resources were not closed
+    //   properly. In this case another reboot is needed before we can take
+    //   another OTA. However, WaitForMerge() can be called again without
+    //   rebooting, to attempt to finish cleanup anyway.
+    //
+    //   MergeCompleted indicates that the update has fully completed.
+    //   GetUpdateState will return None, and a new update can begin.
+    UpdateState WaitForMerge();
 
     // Find the status of the current update, if any.
     //
@@ -109,9 +129,14 @@
     FRIEND_TEST(SnapshotTest, CreateSnapshot);
     FRIEND_TEST(SnapshotTest, MapSnapshot);
     FRIEND_TEST(SnapshotTest, MapPartialSnapshot);
+    FRIEND_TEST(SnapshotTest, NoMergeBeforeReboot);
+    FRIEND_TEST(SnapshotTest, Merge);
+    FRIEND_TEST(SnapshotTest, MergeCannotRemoveCow);
     friend class SnapshotTest;
 
+    using DmTargetSnapshot = android::dm::DmTargetSnapshot;
     using IImageManager = android::fiemap::IImageManager;
+    using TargetInfo = android::dm::DeviceMapper::TargetInfo;
 
     explicit SnapshotManager(IDeviceInfo* info);
 
@@ -126,16 +151,18 @@
     // this. It also serves as a proof-of-lock for some functions.
     class LockedFile final {
       public:
-        LockedFile(const std::string& path, android::base::unique_fd&& fd)
-            : path_(path), fd_(std::move(fd)) {}
+        LockedFile(const std::string& path, android::base::unique_fd&& fd, int lock_mode)
+            : path_(path), fd_(std::move(fd)), lock_mode_(lock_mode) {}
         ~LockedFile();
 
         const std::string& path() const { return path_; }
         int fd() const { return fd_; }
+        int lock_mode() const { return lock_mode_; }
 
       private:
         std::string path_;
         android::base::unique_fd fd_;
+        int lock_mode_;
     };
     std::unique_ptr<LockedFile> OpenFile(const std::string& file, int open_flags, int lock_flags);
     bool Truncate(LockedFile* file);
@@ -189,6 +216,7 @@
     std::unique_ptr<LockedFile> LockExclusive();
     UpdateState ReadUpdateState(LockedFile* file);
     bool WriteUpdateState(LockedFile* file, UpdateState state);
+    std::string GetStateFilePath() const;
 
     // This state is persisted per-snapshot in /metadata/ota/snapshots/.
     struct SnapshotStatus {
@@ -200,11 +228,38 @@
         uint64_t metadata_sectors = 0;
     };
 
+    // Helpers for merging.
+    bool SwitchSnapshotToMerge(LockedFile* lock, const std::string& name);
+    bool RewriteSnapshotDeviceTable(const std::string& dm_name);
+    bool MarkSnapshotMergeCompleted(LockedFile* snapshot_lock, const std::string& snapshot_name);
+    void AcknowledgeMergeSuccess(LockedFile* lock);
+    void AcknowledgeMergeFailure();
+
+    // Note that these require the name of the device containing the snapshot,
+    // which may be the "inner" device. Use GetsnapshotDeviecName().
+    bool QuerySnapshotStatus(const std::string& dm_name, std::string* target_type,
+                             DmTargetSnapshot::Status* status);
+    bool IsSnapshotDevice(const std::string& dm_name, TargetInfo* target = nullptr);
+
+    // 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);
+
+    // Only the following UpdateStates are used here:
+    //   UpdateState::Merging
+    //   UpdateState::MergeCompleted
+    //   UpdateState::MergeFailed
+    //   UpdateState::MergeNeedsReboot
+    UpdateState CheckMergeState();
+    UpdateState CheckMergeState(LockedFile* lock);
+    UpdateState CheckTargetMergeState(LockedFile* lock, const std::string& name);
+
     // Interact with status files under /metadata/ota/snapshots.
-    std::unique_ptr<LockedFile> OpenSnapshotStatusFile(const std::string& name, int open_flags,
-                                                       int lock_flags);
-    bool WriteSnapshotStatus(LockedFile* file, const SnapshotStatus& status);
-    bool ReadSnapshotStatus(LockedFile* file, SnapshotStatus* status);
+    bool WriteSnapshotStatus(LockedFile* lock, const std::string& name,
+                             const SnapshotStatus& status);
+    bool ReadSnapshotStatus(LockedFile* lock, const std::string& name, SnapshotStatus* status);
+    std::string GetSnapshotStatusFilePath(const std::string& name);
 
     // Return the name of the device holding the "snapshot" or "snapshot-merge"
     // target. This may not be the final device presented via MapSnapshot(), if
diff --git a/fs_mgr/libsnapshot/snapshot.cpp b/fs_mgr/libsnapshot/snapshot.cpp
index ef56179..63a01f3 100644
--- a/fs_mgr/libsnapshot/snapshot.cpp
+++ b/fs_mgr/libsnapshot/snapshot.cpp
@@ -19,6 +19,8 @@
 #include <sys/types.h>
 #include <sys/unistd.h>
 
+#include <thread>
+
 #include <android-base/file.h>
 #include <android-base/logging.h>
 #include <android-base/parseint.h>
@@ -140,9 +142,6 @@
 
     LOG(INFO) << "Snapshot " << name << " will have COW size " << cow_size;
 
-    auto status_file = OpenSnapshotStatusFile(name, O_RDWR | O_CREAT, LOCK_EX);
-    if (!status_file) return false;
-
     // Note, we leave the status file hanging around if we fail to create the
     // actual backing image. This is harmless, since it'll get removed when
     // CancelUpdate is called.
@@ -151,7 +150,7 @@
             .device_size = device_size,
             .snapshot_size = snapshot_size,
     };
-    if (!WriteSnapshotStatus(status_file.get(), status)) {
+    if (!WriteSnapshotStatus(lock, name, status)) {
         PLOG(ERROR) << "Could not write snapshot status: " << name;
         return false;
     }
@@ -168,11 +167,13 @@
     CHECK(lock);
     if (!EnsureImageManager()) return false;
 
-    auto status_file = OpenSnapshotStatusFile(name, O_RDWR, LOCK_EX);
-    if (!status_file) return false;
-
     SnapshotStatus status;
-    if (!ReadSnapshotStatus(status_file.get(), &status)) {
+    if (!ReadSnapshotStatus(lock, name, &status)) {
+        return false;
+    }
+    if (status.state == "merge-completed") {
+        LOG(ERROR) << "Should not create a snapshot device for " << name
+                   << " after merging has completed.";
         return false;
     }
 
@@ -210,19 +211,31 @@
 
     std::string cow_dev;
     if (!images_->MapImageDevice(cow_name, timeout_ms, &cow_dev)) {
+        LOG(ERROR) << "Could not map image device: " << cow_name;
         return false;
     }
 
     auto& dm = DeviceMapper::Instance();
 
-    // Merging is a global state, not per-snapshot. We do however track the
-    // progress of individual snapshots' merges.
+    // Note that merging is a global state. We do track whether individual devices
+    // have completed merging, but the start of the merge process is considered
+    // atomic.
     SnapshotStorageMode mode;
-    UpdateState update_state = ReadUpdateState(lock);
-    if (update_state == UpdateState::Merging || update_state == UpdateState::MergeCompleted) {
-        mode = SnapshotStorageMode::Merge;
-    } else {
-        mode = SnapshotStorageMode::Persistent;
+    switch (ReadUpdateState(lock)) {
+        case UpdateState::MergeCompleted:
+        case UpdateState::MergeNeedsReboot:
+            LOG(ERROR) << "Should not create a snapshot device for " << name
+                       << " after global merging has completed.";
+            return false;
+        case UpdateState::Merging:
+        case UpdateState::MergeFailed:
+            // Note: MergeFailed indicates that a merge is in progress, but
+            // is possibly stalled. We still have to honor the merge.
+            mode = SnapshotStorageMode::Merge;
+            break;
+        default:
+            mode = SnapshotStorageMode::Persistent;
+            break;
     }
 
     // The kernel (tested on 4.19) crashes horribly if a device has both a snapshot
@@ -267,11 +280,8 @@
     CHECK(lock);
     if (!EnsureImageManager()) return false;
 
-    auto status_file = OpenSnapshotStatusFile(name, O_RDWR, LOCK_EX);
-    if (!status_file) return false;
-
     SnapshotStatus status;
-    if (!ReadSnapshotStatus(status_file.get(), &status)) {
+    if (!ReadSnapshotStatus(lock, name, &status)) {
         return false;
     }
 
@@ -300,31 +310,525 @@
     CHECK(lock);
     if (!EnsureImageManager()) return false;
 
-    // Take the snapshot's lock after Unmap, since it will also try to lock.
-    auto status_file = OpenSnapshotStatusFile(name, O_RDONLY, LOCK_EX);
-    if (!status_file) return false;
-
     auto cow_name = GetCowName(name);
     if (!images_->BackingImageExists(cow_name)) {
         return true;
     }
+    if (images_->IsImageMapped(cow_name) && !images_->UnmapImageDevice(cow_name)) {
+        return false;
+    }
     if (!images_->DeleteBackingImage(cow_name)) {
         return false;
     }
 
-    if (!android::base::RemoveFileIfExists(status_file->path())) {
-        LOG(ERROR) << "Failed to remove status file: " << status_file->path();
+    std::string error;
+    auto file_path = GetSnapshotStatusFilePath(name);
+    if (!android::base::RemoveFileIfExists(file_path, &error)) {
+        LOG(ERROR) << "Failed to remove status file " << file_path << ": " << error;
         return false;
     }
     return true;
 }
 
 bool SnapshotManager::InitiateMerge() {
-    return false;
+    auto lock = LockExclusive();
+    if (!lock) return false;
+
+    UpdateState state = ReadUpdateState(lock.get());
+    if (state != UpdateState::Unverified) {
+        LOG(ERROR) << "Cannot begin a merge if an update has not been verified";
+        return false;
+    }
+    if (!device_->IsRunningSnapshot()) {
+        LOG(ERROR) << "Cannot begin a merge if the device is not booted off a snapshot";
+        return false;
+    }
+
+    std::vector<std::string> snapshots;
+    if (!ListSnapshots(lock.get(), &snapshots)) {
+        LOG(ERROR) << "Could not list snapshots";
+        return false;
+    }
+
+    auto& dm = DeviceMapper::Instance();
+    for (const auto& snapshot : snapshots) {
+        // The device has to be mapped, since everything should be merged at
+        // the same time. This is a fairly serious error. We could forcefully
+        // map everything here, but it should have been mapped during first-
+        // stage init.
+        if (dm.GetState(snapshot) == DmDeviceState::INVALID) {
+            LOG(ERROR) << "Cannot begin merge; device " << snapshot << " is not mapped.";
+            return false;
+        }
+    }
+
+    // Point of no return - mark that we're starting a merge. From now on every
+    // snapshot must be a merge target.
+    if (!WriteUpdateState(lock.get(), UpdateState::Merging)) {
+        return false;
+    }
+
+    bool rewrote_all = true;
+    for (const auto& snapshot : snapshots) {
+        // If this fails, we have no choice but to continue. Everything must
+        // be merged. This is not an ideal state to be in, but it is safe,
+        // because we the next boot will try again.
+        if (!SwitchSnapshotToMerge(lock.get(), snapshot)) {
+            LOG(ERROR) << "Failed to switch snapshot to a merge target: " << snapshot;
+            rewrote_all = false;
+        }
+    }
+
+    // If we couldn't switch everything to a merge target, pre-emptively mark
+    // this merge as failed. It will get acknowledged when WaitForMerge() is
+    // called.
+    if (!rewrote_all) {
+        WriteUpdateState(lock.get(), UpdateState::MergeFailed);
+    }
+
+    // Return true no matter what, because a merge was initiated.
+    return true;
 }
 
-bool SnapshotManager::WaitForMerge() {
-    return false;
+bool SnapshotManager::SwitchSnapshotToMerge(LockedFile* lock, const std::string& name) {
+    SnapshotStatus status;
+    if (!ReadSnapshotStatus(lock, name, &status)) {
+        return false;
+    }
+    if (status.state != "created") {
+        LOG(WARNING) << "Snapshot " << name << " has unexpected state: " << status.state;
+    }
+
+    // After this, we return true because we technically did switch to a merge
+    // target. Everything else we do here is just informational.
+    auto dm_name = GetSnapshotDeviceName(name, status);
+    if (!RewriteSnapshotDeviceTable(dm_name)) {
+        return false;
+    }
+
+    status.state = "merging";
+
+    DmTargetSnapshot::Status dm_status;
+    if (!QuerySnapshotStatus(dm_name, nullptr, &dm_status)) {
+        LOG(ERROR) << "Could not query merge status for snapshot: " << dm_name;
+    }
+    status.sectors_allocated = dm_status.sectors_allocated;
+    status.metadata_sectors = dm_status.metadata_sectors;
+    if (!WriteSnapshotStatus(lock, name, status)) {
+        LOG(ERROR) << "Could not update status file for snapshot: " << name;
+    }
+    return true;
+}
+
+bool SnapshotManager::RewriteSnapshotDeviceTable(const std::string& dm_name) {
+    auto& dm = DeviceMapper::Instance();
+
+    std::vector<DeviceMapper::TargetInfo> old_targets;
+    if (!dm.GetTableInfo(dm_name, &old_targets)) {
+        LOG(ERROR) << "Could not read snapshot device table: " << dm_name;
+        return false;
+    }
+    if (old_targets.size() != 1 || DeviceMapper::GetTargetType(old_targets[0].spec) != "snapshot") {
+        LOG(ERROR) << "Unexpected device-mapper table for snapshot: " << dm_name;
+        return false;
+    }
+
+    std::string base_device, cow_device;
+    if (!DmTargetSnapshot::GetDevicesFromParams(old_targets[0].data, &base_device, &cow_device)) {
+        LOG(ERROR) << "Could not derive underlying devices for snapshot: " << dm_name;
+        return false;
+    }
+
+    DmTable table;
+    table.Emplace<DmTargetSnapshot>(0, old_targets[0].spec.length, base_device, cow_device,
+                                    SnapshotStorageMode::Merge, kSnapshotChunkSize);
+    if (!dm.LoadTableAndActivate(dm_name, table)) {
+        LOG(ERROR) << "Could not swap device-mapper tables on snapshot device " << dm_name;
+        return false;
+    }
+    LOG(INFO) << "Successfully switched snapshot device to a merge target: " << dm_name;
+    return true;
+}
+
+enum class TableQuery {
+    Table,
+    Status,
+};
+
+static bool GetSingleTarget(const std::string& dm_name, TableQuery query,
+                            DeviceMapper::TargetInfo* target) {
+    auto& dm = DeviceMapper::Instance();
+    if (dm.GetState(dm_name) == DmDeviceState::INVALID) {
+        return false;
+    }
+
+    std::vector<DeviceMapper::TargetInfo> targets;
+    bool result;
+    if (query == TableQuery::Status) {
+        result = dm.GetTableStatus(dm_name, &targets);
+    } else {
+        result = dm.GetTableInfo(dm_name, &targets);
+    }
+    if (!result) {
+        LOG(ERROR) << "Could not query device: " << dm_name;
+        return false;
+    }
+    if (targets.size() != 1) {
+        return false;
+    }
+
+    *target = std::move(targets[0]);
+    return true;
+}
+
+bool SnapshotManager::IsSnapshotDevice(const std::string& dm_name, TargetInfo* target) {
+    DeviceMapper::TargetInfo snap_target;
+    if (!GetSingleTarget(dm_name, TableQuery::Status, &snap_target)) {
+        return false;
+    }
+    auto type = DeviceMapper::GetTargetType(snap_target.spec);
+    if (type != "snapshot" && type != "snapshot-merge") {
+        return false;
+    }
+    if (target) {
+        *target = std::move(snap_target);
+    }
+    return true;
+}
+
+bool SnapshotManager::QuerySnapshotStatus(const std::string& dm_name, std::string* target_type,
+                                          DmTargetSnapshot::Status* status) {
+    DeviceMapper::TargetInfo target;
+    if (!IsSnapshotDevice(dm_name, &target)) {
+        LOG(ERROR) << "Device " << dm_name << " is not a snapshot or snapshot-merge device";
+        return false;
+    }
+    if (!DmTargetSnapshot::ParseStatusText(target.data, status)) {
+        LOG(ERROR) << "Could not parse snapshot status text: " << dm_name;
+        return false;
+    }
+    if (target_type) {
+        *target_type = DeviceMapper::GetTargetType(target.spec);
+    }
+    return true;
+}
+
+// Note that when a merge fails, we will *always* try again to complete the
+// merge each time the device boots. There is no harm in doing so, and if
+// the problem was transient, we might manage to get a new outcome.
+UpdateState SnapshotManager::WaitForMerge() {
+    while (true) {
+        UpdateState state = CheckMergeState();
+        if (state != UpdateState::Merging) {
+            // Either there is no merge, or the merge was finished, so no need
+            // to keep waiting.
+            return state;
+        }
+
+        // This wait is not super time sensitive, so we have a relatively
+        // low polling frequency.
+        std::this_thread::sleep_for(2s);
+    }
+}
+
+UpdateState SnapshotManager::CheckMergeState() {
+    auto lock = LockExclusive();
+    if (!lock) {
+        AcknowledgeMergeFailure();
+        return UpdateState::MergeFailed;
+    }
+
+    auto state = CheckMergeState(lock.get());
+    if (state == UpdateState::MergeCompleted) {
+        AcknowledgeMergeSuccess(lock.get());
+    } else if (state == UpdateState::MergeFailed) {
+        AcknowledgeMergeFailure();
+    }
+    return state;
+}
+
+UpdateState SnapshotManager::CheckMergeState(LockedFile* lock) {
+    UpdateState state = ReadUpdateState(lock);
+    switch (state) {
+        case UpdateState::None:
+        case UpdateState::MergeCompleted:
+            // Harmless races are allowed between two callers of WaitForMerge,
+            // so in both of these cases we just propagate the state.
+            return state;
+
+        case UpdateState::Merging:
+        case UpdateState::MergeNeedsReboot:
+        case UpdateState::MergeFailed:
+            // We'll poll each snapshot below. Note that for the NeedsReboot
+            // case, we always poll once to give cleanup another opportunity to
+            // run.
+            break;
+
+        default:
+            LOG(ERROR) << "No merge exists, cannot wait. Update state: "
+                       << static_cast<uint32_t>(state);
+            return UpdateState::None;
+    }
+
+    std::vector<std::string> snapshots;
+    if (!ListSnapshots(lock, &snapshots)) {
+        return UpdateState::MergeFailed;
+    }
+
+    bool failed = false;
+    bool merging = false;
+    bool needs_reboot = false;
+    for (const auto& snapshot : snapshots) {
+        UpdateState snapshot_state = CheckTargetMergeState(lock, snapshot);
+        switch (snapshot_state) {
+            case UpdateState::MergeFailed:
+                failed = true;
+                break;
+            case UpdateState::Merging:
+                merging = true;
+                break;
+            case UpdateState::MergeNeedsReboot:
+                needs_reboot = true;
+                break;
+            case UpdateState::MergeCompleted:
+                break;
+            default:
+                LOG(ERROR) << "Unknown merge status: " << static_cast<uint32_t>(snapshot_state);
+                failed = true;
+                break;
+        }
+    }
+
+    if (merging) {
+        // Note that we handle "Merging" before we handle anything else. We
+        // want to poll until *nothing* is merging if we can, so everything has
+        // a chance to get marked as completed or failed.
+        return UpdateState::Merging;
+    }
+    if (failed) {
+        // Note: since there are many drop-out cases for failure, we acknowledge
+        // it in WaitForMerge rather than here and elsewhere.
+        return UpdateState::MergeFailed;
+    }
+    if (needs_reboot) {
+        WriteUpdateState(lock, UpdateState::MergeNeedsReboot);
+        return UpdateState::MergeNeedsReboot;
+    }
+    return UpdateState::MergeCompleted;
+}
+
+UpdateState SnapshotManager::CheckTargetMergeState(LockedFile* lock, const std::string& name) {
+    SnapshotStatus snapshot_status;
+    if (!ReadSnapshotStatus(lock, name, &snapshot_status)) {
+        return UpdateState::MergeFailed;
+    }
+
+    std::string dm_name = GetSnapshotDeviceName(name, snapshot_status);
+
+    // During a check, we decided the merge was complete, but we were unable to
+    // collapse the device-mapper stack and perform COW cleanup. If we haven't
+    // rebooted after this check, the device will still be a snapshot-merge
+    // target. If the have rebooted, the device will now be a linear target,
+    // and we can try cleanup again.
+    if (snapshot_status.state == "merge-complete" && !IsSnapshotDevice(dm_name)) {
+        // NB: It's okay if this fails now, we gave cleanup our best effort.
+        OnSnapshotMergeComplete(lock, name, snapshot_status);
+        return UpdateState::MergeCompleted;
+    }
+
+    std::string target_type;
+    DmTargetSnapshot::Status status;
+    if (!QuerySnapshotStatus(dm_name, &target_type, &status)) {
+        return UpdateState::MergeFailed;
+    }
+    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 UpdateState::MergeFailed;
+    }
+
+    // These two values are equal when merging is complete.
+    if (status.sectors_allocated != status.metadata_sectors) {
+        if (snapshot_status.state == "merge-complete") {
+            LOG(ERROR) << "Snapshot " << name << " is merging after being marked merge-complete.";
+            return UpdateState::MergeFailed;
+        }
+        return UpdateState::Merging;
+    }
+
+    // Merging is done. First, update the status file to indicate the merge
+    // is complete. We do this before calling OnSnapshotMergeComplete, even
+    // though this means the write is potentially wasted work (since in the
+    // ideal case we'll immediately delete the file).
+    //
+    // This makes it simpler to reason about the next reboot: no matter what
+    // part of cleanup failed, first-stage init won't try to create another
+    // snapshot device for this partition.
+    snapshot_status.state = "merge-complete";
+    if (!WriteSnapshotStatus(lock, name, snapshot_status)) {
+        return UpdateState::MergeFailed;
+    }
+    if (!OnSnapshotMergeComplete(lock, name, snapshot_status)) {
+        return UpdateState::MergeNeedsReboot;
+    }
+    return UpdateState::MergeCompleted;
+}
+
+void SnapshotManager::AcknowledgeMergeSuccess(LockedFile* lock) {
+    if (!WriteUpdateState(lock, UpdateState::None)) {
+        // We'll try again next reboot, ad infinitum.
+        return;
+    }
+}
+
+void SnapshotManager::AcknowledgeMergeFailure() {
+    // Log first, so worst case, we always have a record of why the calls below
+    // were being made.
+    LOG(ERROR) << "Merge could not be completed and will be marked as failed.";
+
+    auto lock = LockExclusive();
+    if (!lock) return;
+
+    // Since we released the lock in between WaitForMerge and here, it's
+    // possible (1) the merge successfully completed or (2) was already
+    // marked as a failure. So make sure to check the state again, and
+    // only mark as a failure if appropriate.
+    UpdateState state = ReadUpdateState(lock.get());
+    if (state != UpdateState::Merging && state != UpdateState::MergeNeedsReboot) {
+        return;
+    }
+
+    WriteUpdateState(lock.get(), UpdateState::MergeFailed);
+}
+
+bool SnapshotManager::OnSnapshotMergeComplete(LockedFile* lock, const std::string& name,
+                                              const SnapshotStatus& status) {
+    auto dm_name = GetSnapshotDeviceName(name, status);
+    if (IsSnapshotDevice(dm_name)) {
+        // We are extra-cautious here, to avoid deleting the wrong table.
+        std::string target_type;
+        DmTargetSnapshot::Status dm_status;
+        if (!QuerySnapshotStatus(dm_name, &target_type, &dm_status)) {
+            return false;
+        }
+        if (target_type != "snapshot-merge") {
+            LOG(ERROR) << "Unexpected target type " << target_type
+                       << " for snapshot device: " << dm_name;
+            return false;
+        }
+        if (dm_status.sectors_allocated != dm_status.metadata_sectors) {
+            LOG(ERROR) << "Merge is unexpectedly incomplete for device " << dm_name;
+            return false;
+        }
+        if (!CollapseSnapshotDevice(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.
+    }
+
+    if (!DeleteSnapshot(lock, name)) {
+        LOG(ERROR) << "Could not delete snapshot: " << name;
+        return false;
+    }
+    return true;
+}
+
+bool SnapshotManager::CollapseSnapshotDevice(const std::string& name,
+                                             const SnapshotStatus& status) {
+    // Ideally, we would complete the following steps to collapse the device:
+    //  (1) Rewrite the snapshot table to be identical to the base device table.
+    //  (2) Rewrite the verity table to use the "snapshot" (now linear) device.
+    //  (3) Delete the base device.
+    //
+    // This should be possible once libsnapshot understands LpMetadata. In the
+    // meantime, we implement a simpler solution: rewriting the snapshot table
+    // to be a single dm-linear segment against the base device. While not as
+    // ideal, it still lets us remove the COW device. We can remove this
+    // implementation once the new method has been tested.
+    auto& dm = DeviceMapper::Instance();
+    auto dm_name = GetSnapshotDeviceName(name, status);
+
+    DeviceMapper::TargetInfo target;
+    if (!GetSingleTarget(dm_name, TableQuery::Table, &target)) {
+        return false;
+    }
+    if (DeviceMapper::GetTargetType(target.spec) != "snapshot-merge") {
+        // This should be impossible, it was checked above.
+        LOG(ERROR) << "Snapshot device has invalid target type: " << dm_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 " << dm_name
+                   << " parameters: " << target.data;
+        return false;
+    }
+
+    uint64_t num_sectors = status.snapshot_size / kSectorSize;
+    if (num_sectors * kSectorSize != status.snapshot_size) {
+        LOG(ERROR) << "Snapshot " << name
+                   << " size is not sector aligned: " << status.snapshot_size;
+        return false;
+    }
+
+    if (dm_name != name) {
+        // We've derived the base device, but we actually need to replace the
+        // table of the outermost device. Do a quick verification that this
+        // device looks like we expect it to.
+        std::vector<DeviceMapper::TargetInfo> outer_table;
+        if (!dm.GetTableInfo(name, &outer_table)) {
+            LOG(ERROR) << "Could not validate outer snapshot table: " << name;
+            return false;
+        }
+        if (outer_table.size() != 2) {
+            LOG(ERROR) << "Expected 2 dm-linear targets for tabble " << name
+                       << ", got: " << outer_table.size();
+            return false;
+        }
+        for (const auto& target : outer_table) {
+            auto target_type = DeviceMapper::GetTargetType(target.spec);
+            if (target_type != "linear") {
+                LOG(ERROR) << "Outer snapshot table may only contain linear targets, but " << name
+                           << " has target: " << target_type;
+                return false;
+            }
+        }
+        uint64_t sectors = outer_table[0].spec.length + outer_table[1].spec.length;
+        if (sectors != num_sectors) {
+            LOG(ERROR) << "Outer snapshot " << name << " should have " << num_sectors
+                       << ", got: " << sectors;
+            return false;
+        }
+    }
+
+    // Note: we are replacing the OUTER table here, so we do not use dm_name.
+    DmTargetLinear new_target(0, num_sectors, base_device, 0);
+    LOG(INFO) << "Replacing snapshot device " << name
+              << " table with: " << new_target.GetParameterString();
+
+    DmTable table;
+    table.Emplace<DmTargetLinear>(new_target);
+    if (!dm.LoadTableAndActivate(name, table)) {
+        return false;
+    }
+
+    if (dm_name != name) {
+        // Attempt to delete the snapshot device. Nothing should be depending on
+        // the device, and device-mapper should have flushed remaining I/O. We
+        // could in theory replace with dm-zero (or re-use the table above), but
+        // for now it's better to know why this would fail.
+        if (!dm.DeleteDevice(dm_name)) {
+            LOG(ERROR) << "Unable to delete snapshot device " << dm_name << ", COW cannot be "
+                       << "reclaimed until after reboot.";
+            return false;
+        }
+    }
+    return true;
 }
 
 bool SnapshotManager::RemoveAllSnapshots(LockedFile* lock) {
@@ -342,6 +846,12 @@
 }
 
 UpdateState SnapshotManager::GetUpdateState(double* progress) {
+    // If we've never started an update, the state file won't exist.
+    auto state_file = GetStateFilePath();
+    if (access(state_file.c_str(), F_OK) != 0 && errno == ENOENT) {
+        return UpdateState::None;
+    }
+
     auto file = LockShared();
     if (!file) {
         return UpdateState::None;
@@ -388,7 +898,10 @@
         PLOG(ERROR) << "Acquire flock failed: " << file;
         return nullptr;
     }
-    return std::make_unique<LockedFile>(file, std::move(fd));
+    // For simplicity, we want to CHECK that lock_mode == LOCK_EX, in some
+    // calls, so strip extra flags.
+    int lock_mode = lock_flags & (LOCK_EX | LOCK_SH);
+    return std::make_unique<LockedFile>(file, std::move(fd), lock_mode);
 }
 
 SnapshotManager::LockedFile::~LockedFile() {
@@ -397,9 +910,13 @@
     }
 }
 
+std::string SnapshotManager::GetStateFilePath() const {
+    return metadata_dir_ + "/state"s;
+}
+
 std::unique_ptr<SnapshotManager::LockedFile> SnapshotManager::OpenStateFile(int open_flags,
                                                                             int lock_flags) {
-    auto state_file = metadata_dir_ + "/state"s;
+    auto state_file = GetStateFilePath();
     return OpenFile(state_file, open_flags, lock_flags);
 }
 
@@ -434,6 +951,10 @@
         return UpdateState::Merging;
     } else if (contents == "merge-completed") {
         return UpdateState::MergeCompleted;
+    } else if (contents == "merge-needs-reboot") {
+        return UpdateState::MergeNeedsReboot;
+    } else if (contents == "merge-failed") {
+        return UpdateState::MergeFailed;
     } else {
         LOG(ERROR) << "Unknown merge state in update state file";
         return UpdateState::None;
@@ -458,6 +979,12 @@
         case UpdateState::MergeCompleted:
             contents = "merge-completed";
             break;
+        case UpdateState::MergeNeedsReboot:
+            contents = "merge-needs-reboot";
+            break;
+        case UpdateState::MergeFailed:
+            contents = "merge-failed";
+            break;
         default:
             LOG(ERROR) << "Unknown update state";
             return false;
@@ -471,51 +998,66 @@
     return true;
 }
 
-auto SnapshotManager::OpenSnapshotStatusFile(const std::string& name, int open_flags,
-                                             int lock_flags) -> std::unique_ptr<LockedFile> {
+std::string SnapshotManager::GetSnapshotStatusFilePath(const std::string& name) {
     auto file = metadata_dir_ + "/snapshots/"s + name;
-    return OpenFile(file, open_flags, lock_flags);
+    return file;
 }
 
-bool SnapshotManager::ReadSnapshotStatus(LockedFile* file, SnapshotStatus* status) {
-    // Reset position since some calls read+write.
-    if (lseek(file->fd(), 0, SEEK_SET) < 0) {
-        PLOG(ERROR) << "lseek status file failed";
+bool SnapshotManager::ReadSnapshotStatus(LockedFile* lock, const std::string& name,
+                                         SnapshotStatus* status) {
+    CHECK(lock);
+    auto path = GetSnapshotStatusFilePath(name);
+
+    unique_fd fd(open(path.c_str(), O_RDONLY | O_CLOEXEC | O_NOFOLLOW));
+    if (fd < 0) {
+        PLOG(ERROR) << "Open failed: " << path;
         return false;
     }
 
     std::string contents;
-    if (!android::base::ReadFdToString(file->fd(), &contents)) {
-        PLOG(ERROR) << "read status file failed";
+    if (!android::base::ReadFdToString(fd, &contents)) {
+        PLOG(ERROR) << "read failed: " << path;
         return false;
     }
     auto pieces = android::base::Split(contents, " ");
     if (pieces.size() != 5) {
-        LOG(ERROR) << "Invalid status line for snapshot: " << file->path();
+        LOG(ERROR) << "Invalid status line for snapshot: " << path;
         return false;
     }
 
     status->state = pieces[0];
     if (!android::base::ParseUint(pieces[1], &status->device_size)) {
-        LOG(ERROR) << "Invalid device size in status line for: " << file->path();
+        LOG(ERROR) << "Invalid device size in status line for: " << path;
         return false;
     }
     if (!android::base::ParseUint(pieces[2], &status->snapshot_size)) {
-        LOG(ERROR) << "Invalid snapshot size in status line for: " << file->path();
+        LOG(ERROR) << "Invalid snapshot size in status line for: " << path;
         return false;
     }
     if (!android::base::ParseUint(pieces[3], &status->sectors_allocated)) {
-        LOG(ERROR) << "Invalid snapshot size in status line for: " << file->path();
+        LOG(ERROR) << "Invalid snapshot size in status line for: " << path;
         return false;
     }
     if (!android::base::ParseUint(pieces[4], &status->metadata_sectors)) {
-        LOG(ERROR) << "Invalid snapshot size in status line for: " << file->path();
+        LOG(ERROR) << "Invalid snapshot size in status line for: " << path;
         return false;
     }
     return true;
 }
 
-bool SnapshotManager::WriteSnapshotStatus(LockedFile* file, const SnapshotStatus& status) {
+bool SnapshotManager::WriteSnapshotStatus(LockedFile* lock, const std::string& name,
+                                          const SnapshotStatus& status) {
+    // The caller must take an exclusive lock to modify snapshots.
+    CHECK(lock);
+    CHECK(lock->lock_mode() == LOCK_EX);
+
+    auto path = GetSnapshotStatusFilePath(name);
+    unique_fd fd(open(path.c_str(), O_RDWR | O_CLOEXEC | O_NOFOLLOW | O_CREAT | O_SYNC, 0660));
+    if (fd < 0) {
+        PLOG(ERROR) << "Open failed: " << path;
+        return false;
+    }
+
     std::vector<std::string> pieces = {
             status.state,
             std::to_string(status.device_size),
@@ -525,9 +1067,8 @@
     };
     auto contents = android::base::Join(pieces, " ");
 
-    if (!Truncate(file)) return false;
-    if (!android::base::WriteStringToFd(contents, file->fd())) {
-        PLOG(ERROR) << "write to status file failed: " << file->path();
+    if (!android::base::WriteStringToFd(contents, fd)) {
+        PLOG(ERROR) << "write failed: " << path;
         return false;
     }
     return true;
diff --git a/fs_mgr/libsnapshot/snapshot_test.cpp b/fs_mgr/libsnapshot/snapshot_test.cpp
index 9cc9bd7..4903224 100644
--- a/fs_mgr/libsnapshot/snapshot_test.cpp
+++ b/fs_mgr/libsnapshot/snapshot_test.cpp
@@ -22,13 +22,20 @@
 #include <chrono>
 #include <iostream>
 
+#include <android-base/file.h>
+#include <android-base/properties.h>
 #include <android-base/strings.h>
+#include <android-base/unique_fd.h>
 #include <gtest/gtest.h>
+#include <libdm/dm.h>
 #include <libfiemap/image_manager.h>
 
 namespace android {
 namespace snapshot {
 
+using android::base::unique_fd;
+using android::dm::DeviceMapper;
+using android::dm::DmDeviceState;
 using namespace std::chrono_literals;
 using namespace std::string_literals;
 
@@ -48,12 +55,15 @@
 TestDeviceInfo* test_device = nullptr;
 
 class SnapshotTest : public ::testing::Test {
+  public:
+    SnapshotTest() : dm_(DeviceMapper::Instance()) {}
+
   protected:
     void SetUp() override {
         test_device->set_is_running_snapshot(false);
 
         if (sm->GetUpdateState() != UpdateState::None) {
-            ASSERT_TRUE(sm->CancelUpdate());
+            CleanupTestArtifacts();
         }
         ASSERT_TRUE(sm->BeginUpdate());
         ASSERT_TRUE(sm->EnsureImageManager());
@@ -65,13 +75,32 @@
     void TearDown() override {
         lock_ = nullptr;
 
-        if (sm->GetUpdateState() != UpdateState::None) {
-            ASSERT_TRUE(sm->CancelUpdate());
+        CleanupTestArtifacts();
+    }
+
+    void CleanupTestArtifacts() {
+        // Normally cancelling inside a merge is not allowed. Since these
+        // are tests, we don't care, destroy everything that might exist.
+        std::vector<std::string> snapshots = {"test-snapshot"};
+        for (const auto& snapshot : snapshots) {
+            DeleteSnapshotDevice(snapshot);
+            temp_images_.emplace_back(snapshot + "-cow");
+
+            auto status_file = sm->GetSnapshotStatusFilePath(snapshot);
+            android::base::RemoveFileIfExists(status_file);
         }
+
+        // Remove all images.
+        temp_images_.emplace_back("test-snapshot-cow");
         for (const auto& temp_image : temp_images_) {
             image_manager_->UnmapImageDevice(temp_image);
             image_manager_->DeleteBackingImage(temp_image);
         }
+
+        if (sm->GetUpdateState() != UpdateState::None) {
+            auto state_file = sm->GetStateFilePath();
+            unlink(state_file.c_str());
+        }
     }
 
     bool AcquireLock() {
@@ -87,6 +116,17 @@
         return image_manager_->MapImageDevice(name, 10s, path);
     }
 
+    bool DeleteSnapshotDevice(const std::string& snapshot) {
+        if (dm_.GetState(snapshot) != DmDeviceState::INVALID) {
+            if (!dm_.DeleteDevice(snapshot)) return false;
+        }
+        if (dm_.GetState(snapshot + "-inner") != DmDeviceState::INVALID) {
+            if (!dm_.DeleteDevice(snapshot + "-inner")) return false;
+        }
+        return true;
+    }
+
+    DeviceMapper& dm_;
     std::unique_ptr<SnapshotManager::LockedFile> lock_;
     std::vector<std::string> temp_images_;
     android::fiemap::IImageManager* image_manager_ = nullptr;
@@ -106,11 +146,8 @@
 
     // Scope so delete can re-acquire the snapshot file lock.
     {
-        auto file = sm->OpenSnapshotStatusFile("test-snapshot", O_RDONLY, LOCK_SH);
-        ASSERT_NE(file, nullptr);
-
         SnapshotManager::SnapshotStatus status;
-        ASSERT_TRUE(sm->ReadSnapshotStatus(file.get(), &status));
+        ASSERT_TRUE(sm->ReadSnapshotStatus(lock_.get(), "test-snapshot", &status));
         ASSERT_EQ(status.state, "created");
         ASSERT_EQ(status.device_size, kDeviceSize);
         ASSERT_EQ(status.snapshot_size, kDeviceSize);
@@ -151,6 +188,117 @@
     ASSERT_TRUE(android::base::StartsWith(snap_device, "/dev/block/dm-"));
 }
 
+TEST_F(SnapshotTest, NoMergeBeforeReboot) {
+    ASSERT_TRUE(AcquireLock());
+
+    // Set the state to Unverified, as if we finished an update.
+    ASSERT_TRUE(sm->WriteUpdateState(lock_.get(), UpdateState::Unverified));
+
+    // Release the lock.
+    lock_ = nullptr;
+
+    // Merge should fail, since we didn't mark the device as rebooted.
+    ASSERT_FALSE(sm->InitiateMerge());
+}
+
+TEST_F(SnapshotTest, Merge) {
+    ASSERT_TRUE(AcquireLock());
+
+    static const uint64_t kDeviceSize = 1024 * 1024;
+    ASSERT_TRUE(sm->CreateSnapshot(lock_.get(), "test-snapshot", kDeviceSize, kDeviceSize,
+                                   kDeviceSize));
+
+    std::string base_device, snap_device;
+    ASSERT_TRUE(CreateTempDevice("base-device", kDeviceSize, &base_device));
+    ASSERT_TRUE(sm->MapSnapshot(lock_.get(), "test-snapshot", base_device, 10s, &snap_device));
+
+    std::string test_string = "This is a test string.";
+    {
+        unique_fd fd(open(snap_device.c_str(), O_RDWR | O_CLOEXEC | O_SYNC));
+        ASSERT_GE(fd, 0);
+        ASSERT_TRUE(android::base::WriteFully(fd, test_string.data(), test_string.size()));
+    }
+
+    // Note: we know the name of the device is test-snapshot because we didn't
+    // request a linear segment.
+    DeviceMapper::TargetInfo target;
+    ASSERT_TRUE(sm->IsSnapshotDevice("test-snapshot", &target));
+    ASSERT_EQ(DeviceMapper::GetTargetType(target.spec), "snapshot");
+
+    // Set the state to Unverified, as if we finished an update.
+    ASSERT_TRUE(sm->WriteUpdateState(lock_.get(), UpdateState::Unverified));
+
+    // Release the lock.
+    lock_ = nullptr;
+
+    test_device->set_is_running_snapshot(true);
+    ASSERT_TRUE(sm->InitiateMerge());
+
+    // The device should have been switched to a snapshot-merge target.
+    ASSERT_TRUE(sm->IsSnapshotDevice("test-snapshot", &target));
+    ASSERT_EQ(DeviceMapper::GetTargetType(target.spec), "snapshot-merge");
+
+    // We should not be able to cancel an update now.
+    ASSERT_FALSE(sm->CancelUpdate());
+
+    ASSERT_EQ(sm->WaitForMerge(), UpdateState::MergeCompleted);
+    ASSERT_EQ(sm->GetUpdateState(), UpdateState::None);
+
+    // The device should no longer be a snapshot or snapshot-merge.
+    ASSERT_FALSE(sm->IsSnapshotDevice("test-snapshot"));
+
+    // Test that we can read back the string we wrote to the snapshot.
+    unique_fd fd(open(base_device.c_str(), O_RDONLY | O_CLOEXEC));
+    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, MergeCannotRemoveCow) {
+    ASSERT_TRUE(AcquireLock());
+
+    static const uint64_t kDeviceSize = 1024 * 1024;
+    ASSERT_TRUE(sm->CreateSnapshot(lock_.get(), "test-snapshot", kDeviceSize, kDeviceSize,
+                                   kDeviceSize));
+
+    std::string base_device, snap_device;
+    ASSERT_TRUE(CreateTempDevice("base-device", kDeviceSize, &base_device));
+    ASSERT_TRUE(sm->MapSnapshot(lock_.get(), "test-snapshot", base_device, 10s, &snap_device));
+
+    // 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.test-snapshot-cow", "");
+    unique_fd fd(open(cow_path.c_str(), O_RDONLY | O_CLOEXEC));
+    ASSERT_GE(fd, 0);
+
+    // Set the state to Unverified, as if we finished an update.
+    ASSERT_TRUE(sm->WriteUpdateState(lock_.get(), UpdateState::Unverified));
+
+    // Release the lock.
+    lock_ = nullptr;
+
+    test_device->set_is_running_snapshot(true);
+    ASSERT_TRUE(sm->InitiateMerge());
+
+    // COW cannot be removed due to open fd, so expect a soft failure.
+    ASSERT_EQ(sm->WaitForMerge(), UpdateState::MergeNeedsReboot);
+
+    // Forcefully delete the snapshot device, so it looks like we just rebooted.
+    ASSERT_TRUE(DeleteSnapshotDevice("test-snapshot"));
+
+    // Map snapshot should fail now, because we're in a merge-complete state.
+    ASSERT_TRUE(AcquireLock());
+    ASSERT_FALSE(sm->MapSnapshot(lock_.get(), "test-snapshot", base_device, 10s, &snap_device));
+
+    // Release everything and now the merge should complete.
+    fd = {};
+    lock_ = nullptr;
+
+    ASSERT_EQ(sm->WaitForMerge(), UpdateState::MergeCompleted);
+}
+
 }  // namespace snapshot
 }  // namespace android