Merge changes from topic "libsnapshot_prep"

* changes:
  libsnapshot: MapPartitionWithSnapshot cleanup itself if failed
  libsnapshot: MapPartitionWithSnapshot: fix timeout
  libsnapshot: require ex lock for {Create}{CowImage,Snapshot} / DeleteSnapshot
  libsnapshot: Refactor: add {Create,Map,Unmap}CowImage
diff --git a/fs_mgr/libsnapshot/Android.bp b/fs_mgr/libsnapshot/Android.bp
index a54db58..8df9c52 100644
--- a/fs_mgr/libsnapshot/Android.bp
+++ b/fs_mgr/libsnapshot/Android.bp
@@ -49,6 +49,7 @@
     name: "libsnapshot_sources",
     srcs: [
         "snapshot.cpp",
+        "utility.cpp",
     ],
 }
 
diff --git a/fs_mgr/libsnapshot/include/libsnapshot/snapshot.h b/fs_mgr/libsnapshot/include/libsnapshot/snapshot.h
index 6f0e804..c41a951 100644
--- a/fs_mgr/libsnapshot/include/libsnapshot/snapshot.h
+++ b/fs_mgr/libsnapshot/include/libsnapshot/snapshot.h
@@ -211,25 +211,44 @@
     std::unique_ptr<LockedFile> OpenFile(const std::string& file, int open_flags, int lock_flags);
     bool Truncate(LockedFile* file);
 
+    enum class SnapshotState : int { None, Created, Merging, MergeCompleted };
+    static std::string to_string(SnapshotState state);
+
+    // This state is persisted per-snapshot in /metadata/ota/snapshots/.
+    struct SnapshotStatus {
+        SnapshotState state = SnapshotState::None;
+        uint64_t device_size = 0;
+        uint64_t snapshot_size = 0;
+        uint64_t cow_partition_size = 0;
+        uint64_t cow_file_size = 0;
+
+        // These are non-zero when merging.
+        uint64_t sectors_allocated = 0;
+        uint64_t metadata_sectors = 0;
+    };
+
     // 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().
     //
-    // |device_size| should be the size of the base_device that will be passed
-    // via MapDevice(). |snapshot_size| should be the number of bytes in the
-    // base device, starting from 0, that will be snapshotted. The cow_size
+    // |status|.device_size should be the size of the base_device that will be passed
+    // via MapDevice(). |status|.snapshot_size should be the number of bytes in the
+    // base device, starting from 0, that will be snapshotted. |status|.cow_file_size
     // should be the amount of space that will be allocated to store snapshot
     // deltas.
     //
-    // If |snapshot_size| < device_size, then the device will always
+    // If |status|.snapshot_size < |status|.device_size, then the device will always
     // be mapped with two table entries: a dm-snapshot range covering
     // snapshot_size, and a dm-linear range covering the remainder.
     //
-    // All sizes are specified in bytes, and the device and snapshot sizes
-    // must be a multiple of the sector size (512 bytes). |cow_size| will
-    // be rounded up to the nearest sector.
-    bool CreateSnapshot(LockedFile* lock, const std::string& name, uint64_t device_size,
-                        uint64_t snapshot_size, uint64_t cow_size);
+    // All sizes are specified in bytes, and the device, snapshot and COW partition sizes
+    // must be a multiple of the sector size (512 bytes). COW file size will be rounded up
+    // to the nearest sector.
+    bool CreateSnapshot(LockedFile* lock, const std::string& name, SnapshotStatus status);
+
+    // |name| should be the base partition name (e.g. "system_a"). Create the
+    // backing COW image using the size previously passed to CreateSnapshot().
+    bool CreateCowImage(LockedFile* lock, const std::string& name);
 
     // Map a snapshot device that was previously created with CreateSnapshot.
     // If a merge was previously initiated, the device-mapper table will have a
@@ -239,15 +258,23 @@
     // timeout_ms is 0, then no wait will occur and |dev_path| may not yet
     // exist on return.
     bool MapSnapshot(LockedFile* lock, const std::string& name, const std::string& base_device,
-                     const std::chrono::milliseconds& timeout_ms, std::string* dev_path);
+                     const std::string& cow_device, const std::chrono::milliseconds& timeout_ms,
+                     std::string* dev_path);
 
-    // Remove the backing copy-on-write image for the named snapshot. The
+    // Map a COW image that was previous created with CreateCowImage.
+    bool MapCowImage(const std::string& name, const std::chrono::milliseconds& timeout_ms,
+                     std::string* cow_image_device);
+
+    // Remove the backing copy-on-write image and snapshot states for the named snapshot. The
     // caller is responsible for ensuring that the snapshot is unmapped.
     bool DeleteSnapshot(LockedFile* lock, const std::string& name);
 
     // Unmap a snapshot device previously mapped with MapSnapshotDevice().
     bool UnmapSnapshot(LockedFile* lock, const std::string& name);
 
+    // Unmap a COW image device previously mapped with MapCowImage().
+    bool UnmapCowImage(const std::string& name);
+
     // Unmap and remove all known snapshots.
     bool RemoveAllSnapshots(LockedFile* lock);
 
@@ -270,22 +297,6 @@
     bool WriteUpdateState(LockedFile* file, UpdateState state);
     std::string GetStateFilePath() const;
 
-    enum class SnapshotState : int { Created, Merging, MergeCompleted };
-    static std::string to_string(SnapshotState state);
-
-    // This state is persisted per-snapshot in /metadata/ota/snapshots/.
-    struct SnapshotStatus {
-        SnapshotState state;
-        uint64_t device_size;
-        uint64_t snapshot_size;
-        uint64_t cow_partition_size;
-        uint64_t cow_file_size;
-
-        // These are non-zero when merging.
-        uint64_t sectors_allocated = 0;
-        uint64_t metadata_sectors = 0;
-    };
-
     // Helpers for merging.
     bool SwitchSnapshotToMerge(LockedFile* lock, const std::string& name);
     bool RewriteSnapshotDeviceTable(const std::string& dm_name);
diff --git a/fs_mgr/libsnapshot/snapshot.cpp b/fs_mgr/libsnapshot/snapshot.cpp
index 7e57421..f00129a 100644
--- a/fs_mgr/libsnapshot/snapshot.cpp
+++ b/fs_mgr/libsnapshot/snapshot.cpp
@@ -36,6 +36,8 @@
 #include <libfiemap/image_manager.h>
 #include <liblp/liblp.h>
 
+#include "utility.h"
+
 namespace android {
 namespace snapshot {
 
@@ -54,6 +56,7 @@
 using android::fs_mgr::GetPartitionName;
 using android::fs_mgr::LpMetadata;
 using android::fs_mgr::SlotNumberForSlotSuffix;
+using std::chrono::duration_cast;
 using namespace std::chrono_literals;
 using namespace std::string_literals;
 
@@ -99,10 +102,14 @@
     metadata_dir_ = device_->GetMetadataDir();
 }
 
-static std::string GetCowName(const std::string& snapshot_name) {
+[[maybe_unused]] static std::string GetCowName(const std::string& snapshot_name) {
     return snapshot_name + "-cow";
 }
 
+static std::string GetCowImageDeviceName(const std::string& snapshot_name) {
+    return snapshot_name + "-cow-img";
+}
+
 static std::string GetBaseDeviceName(const std::string& partition_name) {
     return partition_name + "-base";
 }
@@ -177,49 +184,60 @@
 }
 
 bool SnapshotManager::CreateSnapshot(LockedFile* lock, const std::string& name,
-                                     uint64_t device_size, uint64_t snapshot_size,
-                                     uint64_t cow_size) {
+                                     SnapshotManager::SnapshotStatus status) {
     CHECK(lock);
-    if (!EnsureImageManager()) return false;
-
+    CHECK(lock->lock_mode() == LOCK_EX);
     // Sanity check these sizes. Like liblp, we guarantee the partition size
     // is respected, which means it has to be sector-aligned. (This guarantee
     // is useful for locating avb footers correctly). The COW size, however,
     // can be arbitrarily larger than specified, so we can safely round it up.
-    if (device_size % kSectorSize != 0) {
+    if (status.device_size % kSectorSize != 0) {
         LOG(ERROR) << "Snapshot " << name
-                   << " device size is not a multiple of the sector size: " << device_size;
+                   << " device size is not a multiple of the sector size: " << status.device_size;
         return false;
     }
-    if (snapshot_size % kSectorSize != 0) {
-        LOG(ERROR) << "Snapshot " << name
-                   << " snapshot size is not a multiple of the sector size: " << snapshot_size;
+    if (status.snapshot_size % kSectorSize != 0) {
+        LOG(ERROR) << "Snapshot " << name << " snapshot size is not a multiple of the sector size: "
+                   << status.snapshot_size;
         return false;
     }
 
     // Round the COW size up to the nearest sector.
-    cow_size += kSectorSize - 1;
-    cow_size &= ~(kSectorSize - 1);
+    status.cow_file_size += kSectorSize - 1;
+    status.cow_file_size &= ~(kSectorSize - 1);
 
-    LOG(INFO) << "Snapshot " << name << " will have COW size " << cow_size;
+    status.state = SnapshotState::Created;
+    status.sectors_allocated = 0;
+    status.metadata_sectors = 0;
 
-    // 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.
-    SnapshotStatus status = {
-            .state = SnapshotState::Created,
-            .device_size = device_size,
-            .snapshot_size = snapshot_size,
-            .cow_file_size = cow_size,
-    };
     if (!WriteSnapshotStatus(lock, name, status)) {
         PLOG(ERROR) << "Could not write snapshot status: " << name;
         return false;
     }
+    return true;
+}
 
-    auto cow_name = GetCowName(name);
+bool SnapshotManager::CreateCowImage(LockedFile* lock, const std::string& name) {
+    CHECK(lock);
+    CHECK(lock->lock_mode() == LOCK_EX);
+    if (!EnsureImageManager()) return false;
+
+    SnapshotStatus status;
+    if (!ReadSnapshotStatus(lock, name, &status)) {
+        return false;
+    }
+
+    // The COW file size should have been rounded up to the nearest sector in CreateSnapshot.
+    // Sanity check this.
+    if (status.cow_file_size % kSectorSize != 0) {
+        LOG(ERROR) << "Snapshot " << name << " COW file size is not a multiple of the sector size: "
+                   << status.cow_file_size;
+        return false;
+    }
+
+    std::string cow_image_name = GetCowImageDeviceName(name);
     int cow_flags = IImageManager::CREATE_IMAGE_DEFAULT;
-    if (!images_->CreateBackingImage(cow_name, cow_size, cow_flags)) {
+    if (!images_->CreateBackingImage(cow_image_name, status.cow_file_size, cow_flags)) {
         return false;
     }
 
@@ -238,11 +256,11 @@
     // workaround that will be discussed again when the kernel API gets
     // consolidated.
     ssize_t dm_snap_magic_size = 4;  // 32 bit
-    return images_->ZeroFillNewImage(cow_name, dm_snap_magic_size);
+    return images_->ZeroFillNewImage(cow_image_name, dm_snap_magic_size);
 }
 
 bool SnapshotManager::MapSnapshot(LockedFile* lock, const std::string& name,
-                                  const std::string& base_device,
+                                  const std::string& base_device, const std::string& cow_device,
                                   const std::chrono::milliseconds& timeout_ms,
                                   std::string* dev_path) {
     CHECK(lock);
@@ -288,22 +306,7 @@
     uint64_t snapshot_sectors = status.snapshot_size / kSectorSize;
     uint64_t linear_sectors = (status.device_size - status.snapshot_size) / kSectorSize;
 
-    auto cow_name = GetCowName(name);
 
-    bool ok;
-    std::string cow_dev;
-    if (has_local_image_manager_) {
-        // If we forced a local image manager, it means we don't have binder,
-        // which means first-stage init. We must use device-mapper.
-        const auto& opener = device_->GetPartitionOpener();
-        ok = images_->MapImageWithDeviceMapper(opener, cow_name, &cow_dev);
-    } else {
-        ok = images_->MapImageDevice(cow_name, timeout_ms, &cow_dev);
-    }
-    if (!ok) {
-        LOG(ERROR) << "Could not map image device: " << cow_name;
-        return false;
-    }
 
     auto& dm = DeviceMapper::Instance();
 
@@ -335,11 +338,10 @@
     auto snap_name = (linear_sectors > 0) ? GetSnapshotExtraDeviceName(name) : name;
 
     DmTable table;
-    table.Emplace<DmTargetSnapshot>(0, snapshot_sectors, base_device, cow_dev, mode,
+    table.Emplace<DmTargetSnapshot>(0, snapshot_sectors, base_device, cow_device, mode,
                                     kSnapshotChunkSize);
     if (!dm.CreateDevice(snap_name, table, dev_path, timeout_ms)) {
         LOG(ERROR) << "Could not create snapshot device: " << snap_name;
-        images_->UnmapImageDevice(cow_name);
         return false;
     }
 
@@ -355,7 +357,6 @@
         if (!dm.CreateDevice(name, table, dev_path, timeout_ms)) {
             LOG(ERROR) << "Could not create outer snapshot device: " << name;
             dm.DeleteDevice(snap_name);
-            images_->UnmapImageDevice(cow_name);
             return false;
         }
     }
@@ -366,9 +367,29 @@
     return true;
 }
 
+bool SnapshotManager::MapCowImage(const std::string& name,
+                                  const std::chrono::milliseconds& timeout_ms,
+                                  std::string* cow_dev) {
+    if (!EnsureImageManager()) return false;
+    auto cow_image_name = GetCowImageDeviceName(name);
+
+    bool ok;
+    if (has_local_image_manager_) {
+        // If we forced a local image manager, it means we don't have binder,
+        // which means first-stage init. We must use device-mapper.
+        const auto& opener = device_->GetPartitionOpener();
+        ok = images_->MapImageWithDeviceMapper(opener, cow_image_name, cow_dev);
+    } else {
+        ok = images_->MapImageDevice(cow_image_name, timeout_ms, cow_dev);
+    }
+    if (!ok) {
+        LOG(ERROR) << "Could not map image device: " << cow_image_name;
+    }
+    return ok;
+}
+
 bool SnapshotManager::UnmapSnapshot(LockedFile* lock, const std::string& name) {
     CHECK(lock);
-    if (!EnsureImageManager()) return false;
 
     SnapshotStatus status;
     if (!ReadSnapshotStatus(lock, name, &status)) {
@@ -389,23 +410,25 @@
         return false;
     }
 
-    auto cow_name = GetCowName(name);
-    if (images_->IsImageMapped(cow_name) && !images_->UnmapImageDevice(cow_name)) {
-        return false;
-    }
     return true;
 }
 
+bool SnapshotManager::UnmapCowImage(const std::string& name) {
+    if (!EnsureImageManager()) return false;
+    return images_->UnmapImageIfExists(GetCowImageDeviceName(name));
+}
+
 bool SnapshotManager::DeleteSnapshot(LockedFile* lock, const std::string& name) {
     CHECK(lock);
+    CHECK(lock->lock_mode() == LOCK_EX);
     if (!EnsureImageManager()) return false;
 
-    auto cow_name = GetCowName(name);
-    if (images_->BackingImageExists(cow_name)) {
-        if (images_->IsImageMapped(cow_name) && !images_->UnmapImageDevice(cow_name)) {
+    auto cow_image_name = GetCowImageDeviceName(name);
+    if (images_->BackingImageExists(cow_image_name)) {
+        if (!images_->UnmapImageIfExists(cow_image_name)) {
             return false;
         }
-        if (!images_->DeleteBackingImage(cow_name)) {
+        if (!images_->DeleteBackingImage(cow_image_name)) {
             return false;
         }
     }
@@ -1153,9 +1176,28 @@
     return true;
 }
 
+static std::chrono::milliseconds GetRemainingTime(
+        const std::chrono::milliseconds& timeout,
+        const std::chrono::time_point<std::chrono::steady_clock>& begin) {
+    // If no timeout is specified, execute all commands without specifying any timeout.
+    if (timeout.count() == 0) return std::chrono::milliseconds(0);
+    auto passed_time = std::chrono::steady_clock::now() - begin;
+    auto remaining_time = timeout - duration_cast<std::chrono::milliseconds>(passed_time);
+    if (remaining_time.count() <= 0) {
+        LOG(ERROR) << "MapPartitionWithSnapshot has reached timeout " << timeout.count() << "ms ("
+                   << remaining_time.count() << "ms remaining)";
+        // Return min() instead of remaining_time here because 0 is treated as a special value for
+        // no timeout, where the rest of the commands will still be executed.
+        return std::chrono::milliseconds::min();
+    }
+    return remaining_time;
+}
+
 bool SnapshotManager::MapPartitionWithSnapshot(LockedFile* lock,
                                                CreateLogicalPartitionParams params,
                                                std::string* path) {
+    auto begin = std::chrono::steady_clock::now();
+
     CHECK(lock);
     path->clear();
 
@@ -1207,6 +1249,11 @@
         params.device_name = GetBaseDeviceName(params.GetPartitionName());
     }
 
+    AutoDeviceList created_devices;
+
+    // Create the base device for the snapshot, or if there is no snapshot, the
+    // device itself. This device consists of the real blocks in the super
+    // partition that this logical partition occupies.
     auto& dm = DeviceMapper::Instance();
     std::string ignore_path;
     if (!CreateLogicalPartition(params, &ignore_path)) {
@@ -1214,7 +1261,10 @@
                    << " as device " << params.GetDeviceName();
         return false;
     }
+    created_devices.EmplaceBack<AutoUnmapDevice>(&dm, params.GetDeviceName());
+
     if (!live_snapshot_status.has_value()) {
+        created_devices.Release();
         return true;
     }
 
@@ -1225,10 +1275,35 @@
         LOG(ERROR) << "Could not determine major/minor for: " << params.GetDeviceName();
         return false;
     }
-    if (!MapSnapshot(lock, params.GetPartitionName(), base_device, {}, path)) {
+
+    // If there is a timeout specified, compute the remaining time to call Map* functions.
+    // init calls CreateLogicalAndSnapshotPartitions, which has no timeout specified. Still call
+    // Map* functions in this case.
+    auto remaining_time = GetRemainingTime(params.timeout_ms, begin);
+    if (remaining_time.count() < 0) return false;
+
+    std::string cow_image_device;
+    if (!MapCowImage(params.GetPartitionName(), remaining_time, &cow_image_device)) {
+        LOG(ERROR) << "Could not map cow image for partition: " << params.GetPartitionName();
+        return false;
+    }
+    created_devices.EmplaceBack<AutoUnmapImage>(images_.get(),
+                                                GetCowImageDeviceName(params.partition_name));
+
+    // TODO: map cow linear device here
+    std::string cow_device = cow_image_device;
+
+    remaining_time = GetRemainingTime(params.timeout_ms, begin);
+    if (remaining_time.count() < 0) return false;
+
+    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.
+
+    created_devices.Release();
 
     LOG(INFO) << "Mapped " << params.GetPartitionName() << " as snapshot device at " << *path;
 
@@ -1373,7 +1448,9 @@
         return false;
     }
 
-    if (pieces[0] == "created") {
+    if (pieces[0] == "none") {
+        status->state = SnapshotState::None;
+    } else if (pieces[0] == "created") {
         status->state = SnapshotState::Created;
     } else if (pieces[0] == "merging") {
         status->state = SnapshotState::Merging;
@@ -1381,6 +1458,7 @@
         status->state = SnapshotState::MergeCompleted;
     } else {
         LOG(ERROR) << "Unrecognized state " << pieces[0] << " for snapshot: " << name;
+        return false;
     }
 
     if (!android::base::ParseUint(pieces[1], &status->device_size)) {
@@ -1412,6 +1490,8 @@
 
 std::string SnapshotManager::to_string(SnapshotState state) {
     switch (state) {
+        case SnapshotState::None:
+            return "none";
         case SnapshotState::Created:
             return "created";
         case SnapshotState::Merging:
diff --git a/fs_mgr/libsnapshot/snapshot_test.cpp b/fs_mgr/libsnapshot/snapshot_test.cpp
index 8487339..429fd8e 100644
--- a/fs_mgr/libsnapshot/snapshot_test.cpp
+++ b/fs_mgr/libsnapshot/snapshot_test.cpp
@@ -106,7 +106,7 @@
                                               "test_partition_b"};
         for (const auto& snapshot : snapshots) {
             DeleteSnapshotDevice(snapshot);
-            DeleteBackingImage(image_manager_, snapshot + "-cow");
+            DeleteBackingImage(image_manager_, snapshot + "-cow-img");
 
             auto status_file = sm->GetSnapshotStatusFilePath(snapshot);
             android::base::RemoveFileIfExists(status_file);
@@ -214,6 +214,7 @@
     void DeleteSnapshotDevice(const std::string& snapshot) {
         DeleteDevice(snapshot);
         DeleteDevice(snapshot + "-inner");
+        ASSERT_TRUE(image_manager_->UnmapImageIfExists(snapshot + "-cow-img"));
     }
     void DeleteDevice(const std::string& device) {
         if (dm_.GetState(device) != DmDeviceState::INVALID) {
@@ -231,8 +232,11 @@
     ASSERT_TRUE(AcquireLock());
 
     static const uint64_t kDeviceSize = 1024 * 1024;
-    ASSERT_TRUE(sm->CreateSnapshot(lock_.get(), "test-snapshot", kDeviceSize, kDeviceSize,
-                                   kDeviceSize));
+    ASSERT_TRUE(sm->CreateSnapshot(lock_.get(), "test-snapshot",
+                                   {.device_size = kDeviceSize,
+                                    .snapshot_size = kDeviceSize,
+                                    .cow_file_size = kDeviceSize}));
+    ASSERT_TRUE(sm->CreateCowImage(lock_.get(), "test-snapshot"));
 
     std::vector<std::string> snapshots;
     ASSERT_TRUE(sm->ListSnapshots(lock_.get(), &snapshots));
@@ -249,6 +253,7 @@
     }
 
     ASSERT_TRUE(sm->UnmapSnapshot(lock_.get(), "test-snapshot"));
+    ASSERT_TRUE(sm->UnmapCowImage("test-snapshot"));
     ASSERT_TRUE(sm->DeleteSnapshot(lock_.get(), "test-snapshot"));
 }
 
@@ -256,14 +261,21 @@
     ASSERT_TRUE(AcquireLock());
 
     static const uint64_t kDeviceSize = 1024 * 1024;
-    ASSERT_TRUE(sm->CreateSnapshot(lock_.get(), "test-snapshot", kDeviceSize, kDeviceSize,
-                                   kDeviceSize));
+    ASSERT_TRUE(sm->CreateSnapshot(lock_.get(), "test-snapshot",
+                                   {.device_size = kDeviceSize,
+                                    .snapshot_size = kDeviceSize,
+                                    .cow_file_size = kDeviceSize}));
+    ASSERT_TRUE(sm->CreateCowImage(lock_.get(), "test-snapshot"));
 
     std::string base_device;
     ASSERT_TRUE(CreatePartition("base-device", kDeviceSize, &base_device));
 
+    std::string cow_device;
+    ASSERT_TRUE(sm->MapCowImage("test-snapshot", 10s, &cow_device));
+
     std::string snap_device;
-    ASSERT_TRUE(sm->MapSnapshot(lock_.get(), "test-snapshot", base_device, 10s, &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-"));
 }
 
@@ -272,14 +284,21 @@
 
     static const uint64_t kSnapshotSize = 1024 * 1024;
     static const uint64_t kDeviceSize = 1024 * 1024 * 2;
-    ASSERT_TRUE(sm->CreateSnapshot(lock_.get(), "test-snapshot", kDeviceSize, kSnapshotSize,
-                                   kSnapshotSize));
+    ASSERT_TRUE(sm->CreateSnapshot(lock_.get(), "test-snapshot",
+                                   {.device_size = kDeviceSize,
+                                    .snapshot_size = kSnapshotSize,
+                                    .cow_file_size = kSnapshotSize}));
+    ASSERT_TRUE(sm->CreateCowImage(lock_.get(), "test-snapshot"));
 
     std::string base_device;
     ASSERT_TRUE(CreatePartition("base-device", kDeviceSize, &base_device));
 
+    std::string cow_device;
+    ASSERT_TRUE(sm->MapCowImage("test-snapshot", 10s, &cow_device));
+
     std::string snap_device;
-    ASSERT_TRUE(sm->MapSnapshot(lock_.get(), "test-snapshot", base_device, 10s, &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-"));
 }
 
@@ -317,13 +336,18 @@
 
     static const uint64_t kDeviceSize = 1024 * 1024;
 
-    std::string base_device, snap_device;
+    std::string base_device, cow_device, snap_device;
     ASSERT_TRUE(CreatePartition("test_partition_a", kDeviceSize));
     ASSERT_TRUE(MapUpdatePartitions());
     ASSERT_TRUE(dm_.GetDmDevicePathByName("test_partition_b-base", &base_device));
-    ASSERT_TRUE(sm->CreateSnapshot(lock_.get(), "test_partition_b", kDeviceSize, kDeviceSize,
-                                   kDeviceSize));
-    ASSERT_TRUE(sm->MapSnapshot(lock_.get(), "test_partition_b", base_device, 10s, &snap_device));
+    ASSERT_TRUE(sm->CreateSnapshot(lock_.get(), "test_partition_b",
+                                   {.device_size = kDeviceSize,
+                                    .snapshot_size = kDeviceSize,
+                                    .cow_file_size = kDeviceSize}));
+    ASSERT_TRUE(sm->CreateCowImage(lock_.get(), "test_partition_b"));
+    ASSERT_TRUE(sm->MapCowImage("test_partition_b", 10s, &cow_device));
+    ASSERT_TRUE(sm->MapSnapshot(lock_.get(), "test_partition_b", base_device, cow_device, 10s,
+                                &snap_device));
 
     std::string test_string = "This is a test string.";
     {
@@ -375,16 +399,21 @@
     ASSERT_TRUE(AcquireLock());
 
     static const uint64_t kDeviceSize = 1024 * 1024;
-    ASSERT_TRUE(sm->CreateSnapshot(lock_.get(), "test-snapshot", kDeviceSize, kDeviceSize,
-                                   kDeviceSize));
+    ASSERT_TRUE(sm->CreateSnapshot(lock_.get(), "test-snapshot",
+                                   {.device_size = kDeviceSize,
+                                    .snapshot_size = kDeviceSize,
+                                    .cow_file_size = kDeviceSize}));
+    ASSERT_TRUE(sm->CreateCowImage(lock_.get(), "test-snapshot"));
 
-    std::string base_device, snap_device;
+    std::string base_device, cow_device, snap_device;
     ASSERT_TRUE(CreatePartition("base-device", kDeviceSize, &base_device));
-    ASSERT_TRUE(sm->MapSnapshot(lock_.get(), "test-snapshot", base_device, 10s, &snap_device));
+    ASSERT_TRUE(sm->MapCowImage("test-snapshot", 10s, &cow_device));
+    ASSERT_TRUE(sm->MapSnapshot(lock_.get(), "test-snapshot", base_device, cow_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", "");
+    auto cow_path = android::base::GetProperty("gsid.mapped_image.test-snapshot-cow-img", "");
     unique_fd fd(open(cow_path.c_str(), O_RDONLY | O_CLOEXEC));
     ASSERT_GE(fd, 0);
 
@@ -399,12 +428,18 @@
     // COW cannot be removed due to open fd, so expect a soft failure.
     ASSERT_EQ(sm->ProcessUpdateState(), UpdateState::MergeNeedsReboot);
 
+    // Release the handle to the COW device to fake a reboot.
+    fd.reset();
+    // Wait 1s, otherwise DeleteSnapshotDevice may fail with EBUSY.
+    sleep(1);
     // Forcefully delete the snapshot device, so it looks like we just rebooted.
     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));
+    ASSERT_TRUE(sm->MapCowImage("test-snapshot", 10s, &cow_device));
+    ASSERT_FALSE(sm->MapSnapshot(lock_.get(), "test-snapshot", base_device, cow_device, 10s,
+                                 &snap_device));
 
     // Release everything and now the merge should complete.
     fd = {};
@@ -423,8 +458,11 @@
 
     ASSERT_TRUE(CreatePartition("test_partition_a", kDeviceSize));
     ASSERT_TRUE(MapUpdatePartitions());
-    ASSERT_TRUE(sm->CreateSnapshot(lock_.get(), "test_partition_b", kDeviceSize, kDeviceSize,
-                                   kDeviceSize));
+    ASSERT_TRUE(sm->CreateSnapshot(lock_.get(), "test_partition_b",
+                                   {.device_size = kDeviceSize,
+                                    .snapshot_size = kDeviceSize,
+                                    .cow_file_size = kDeviceSize}));
+    ASSERT_TRUE(sm->CreateCowImage(lock_.get(), "test_partition_b"));
 
     // Simulate a reboot into the new slot.
     lock_ = nullptr;
@@ -462,8 +500,11 @@
 
     ASSERT_TRUE(CreatePartition("test_partition_a", kDeviceSize));
     ASSERT_TRUE(MapUpdatePartitions());
-    ASSERT_TRUE(sm->CreateSnapshot(lock_.get(), "test_partition_b", kDeviceSize, kDeviceSize,
-                                   kDeviceSize));
+    ASSERT_TRUE(sm->CreateSnapshot(lock_.get(), "test_partition_b",
+                                   {.device_size = kDeviceSize,
+                                    .snapshot_size = kDeviceSize,
+                                    .cow_file_size = kDeviceSize}));
+    ASSERT_TRUE(sm->CreateCowImage(lock_.get(), "test_partition_b"));
 
     // Simulate a reboot into the new slot.
     lock_ = nullptr;
@@ -507,8 +548,11 @@
 
     ASSERT_TRUE(CreatePartition("test_partition_a", kDeviceSize));
     ASSERT_TRUE(MapUpdatePartitions());
-    ASSERT_TRUE(sm->CreateSnapshot(lock_.get(), "test_partition_b", kDeviceSize, kDeviceSize,
-                                   kDeviceSize));
+    ASSERT_TRUE(sm->CreateSnapshot(lock_.get(), "test_partition_b",
+                                   {.device_size = kDeviceSize,
+                                    .snapshot_size = kDeviceSize,
+                                    .cow_file_size = kDeviceSize}));
+    ASSERT_TRUE(sm->CreateCowImage(lock_.get(), "test_partition_b"));
 
     // Simulate a reboot into the new slot.
     lock_ = nullptr;
@@ -527,7 +571,7 @@
     // Now, reflash super. Note that we haven't called ProcessUpdateState, so the
     // status is still Merging.
     DeleteSnapshotDevice("test_partition_b");
-    ASSERT_TRUE(init->image_manager()->UnmapImageDevice("test_partition_b-cow"));
+    ASSERT_TRUE(init->image_manager()->UnmapImageIfExists("test_partition_b-cow-img"));
     FormatFakeSuper();
     ASSERT_TRUE(CreatePartition("test_partition_b", kDeviceSize));
     ASSERT_TRUE(init->NeedSnapshotsInFirstStageMount());
diff --git a/fs_mgr/libsnapshot/utility.cpp b/fs_mgr/libsnapshot/utility.cpp
new file mode 100644
index 0000000..164b472
--- /dev/null
+++ b/fs_mgr/libsnapshot/utility.cpp
@@ -0,0 +1,56 @@
+// Copyright (C) 2019 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 "utility.h"
+
+#include <android-base/logging.h>
+#include <android-base/strings.h>
+
+namespace android {
+namespace snapshot {
+
+void AutoDevice::Release() {
+    name_.clear();
+}
+
+AutoDeviceList::~AutoDeviceList() {
+    // Destroy devices in the reverse order because newer devices may have dependencies
+    // on older devices.
+    for (auto it = devices_.rbegin(); it != devices_.rend(); ++it) {
+        it->reset();
+    }
+}
+
+void AutoDeviceList::Release() {
+    for (auto&& p : devices_) {
+        p->Release();
+    }
+}
+
+AutoUnmapDevice::~AutoUnmapDevice() {
+    if (name_.empty()) return;
+    if (!dm_->DeleteDeviceIfExists(name_)) {
+        LOG(ERROR) << "Failed to auto unmap device " << name_;
+    }
+}
+
+AutoUnmapImage::~AutoUnmapImage() {
+    if (name_.empty()) return;
+    if (!images_->UnmapImageIfExists(name_)) {
+        LOG(ERROR) << "Failed to auto unmap cow image " << name_;
+    }
+}
+
+}  // namespace snapshot
+}  // namespace android
diff --git a/fs_mgr/libsnapshot/utility.h b/fs_mgr/libsnapshot/utility.h
new file mode 100644
index 0000000..cbab472
--- /dev/null
+++ b/fs_mgr/libsnapshot/utility.h
@@ -0,0 +1,85 @@
+// Copyright (C) 2019 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.
+
+#pragma once
+
+#include <string>
+
+#include <android-base/macros.h>
+#include <libdm/dm.h>
+#include <libfiemap/image_manager.h>
+
+namespace android {
+namespace snapshot {
+
+struct AutoDevice {
+    virtual ~AutoDevice(){};
+    void Release();
+
+  protected:
+    AutoDevice(const std::string& name) : name_(name) {}
+    std::string name_;
+
+  private:
+    DISALLOW_COPY_AND_ASSIGN(AutoDevice);
+    AutoDevice(AutoDevice&& other) = delete;
+};
+
+// A list of devices we created along the way.
+// - Whenever a device is created that is subject to GC'ed at the end of
+//   this function, add it to this list.
+// - If any error has occurred, the list is destroyed, and all these devices
+//   are cleaned up.
+// - Upon success, Release() should be called so that the created devices
+//   are kept.
+struct AutoDeviceList {
+    ~AutoDeviceList();
+    template <typename T, typename... Args>
+    void EmplaceBack(Args&&... args) {
+        devices_.emplace_back(std::make_unique<T>(std::forward<Args>(args)...));
+    }
+    void Release();
+
+  private:
+    std::vector<std::unique_ptr<AutoDevice>> devices_;
+};
+
+// Automatically unmap a device upon deletion.
+struct AutoUnmapDevice : AutoDevice {
+    // On destruct, delete |name| from device mapper.
+    AutoUnmapDevice(android::dm::DeviceMapper* dm, const std::string& name)
+        : AutoDevice(name), dm_(dm) {}
+    AutoUnmapDevice(AutoUnmapDevice&& other) = default;
+    ~AutoUnmapDevice();
+
+  private:
+    DISALLOW_COPY_AND_ASSIGN(AutoUnmapDevice);
+    android::dm::DeviceMapper* dm_ = nullptr;
+};
+
+// Automatically unmap an image upon deletion.
+struct AutoUnmapImage : AutoDevice {
+    // On destruct, delete |name| from image manager.
+    AutoUnmapImage(android::fiemap::IImageManager* images, const std::string& name)
+        : AutoDevice(name), images_(images) {}
+    AutoUnmapImage(AutoUnmapImage&& other) = default;
+    ~AutoUnmapImage();
+
+  private:
+    DISALLOW_COPY_AND_ASSIGN(AutoUnmapImage);
+    android::fiemap::IImageManager* images_ = nullptr;
+};
+
+}  // namespace snapshot
+}  // namespace android