Merge "fs_mgr: overlayfs: empty system partitions excluded."
diff --git a/fastboot/device/utility.cpp b/fastboot/device/utility.cpp
index e01e39b..b3f2d5f 100644
--- a/fastboot/device/utility.cpp
+++ b/fastboot/device/utility.cpp
@@ -56,8 +56,16 @@
     if (!path) {
         return false;
     }
+
+    CreateLogicalPartitionParams params = {
+            .block_device = *path,
+            .metadata_slot = slot_number,
+            .partition_name = partition_name,
+            .force_writable = true,
+            .timeout_ms = 5s,
+    };
     std::string dm_path;
-    if (!CreateLogicalPartition(path->c_str(), slot_number, partition_name, true, 5s, &dm_path)) {
+    if (!CreateLogicalPartition(params, &dm_path)) {
         LOG(ERROR) << "Could not map partition: " << partition_name;
         return false;
     }
diff --git a/fs_mgr/fs_mgr_dm_linear.cpp b/fs_mgr/fs_mgr_dm_linear.cpp
index 04ba0bf..eaa515a 100644
--- a/fs_mgr/fs_mgr_dm_linear.cpp
+++ b/fs_mgr/fs_mgr_dm_linear.cpp
@@ -109,26 +109,6 @@
     return true;
 }
 
-static bool CreateLogicalPartition(const LpMetadata& metadata, const LpMetadataPartition& partition,
-                                   bool force_writable, const std::chrono::milliseconds& timeout_ms,
-                                   const std::string& super_device, std::string* path) {
-    DeviceMapper& dm = DeviceMapper::Instance();
-
-    DmTable table;
-    if (!CreateDmTable(metadata, partition, super_device, &table)) {
-        return false;
-    }
-    if (force_writable) {
-        table.set_readonly(false);
-    }
-    std::string name = GetPartitionName(partition);
-    if (!dm.CreateDevice(name, table, path, timeout_ms)) {
-        return false;
-    }
-    LINFO << "Created logical partition " << name << " on device " << *path;
-    return true;
-}
-
 bool CreateLogicalPartitions(const std::string& block_device) {
     uint32_t slot = SlotNumberForSlotSuffix(fs_mgr_get_slot_suffix());
     auto metadata = ReadMetadata(block_device.c_str(), slot);
@@ -145,13 +125,20 @@
 }
 
 bool CreateLogicalPartitions(const LpMetadata& metadata, const std::string& super_device) {
+    CreateLogicalPartitionParams params = {
+            .block_device = super_device,
+            .metadata = &metadata,
+    };
     for (const auto& partition : metadata.partitions) {
         if (!partition.num_extents) {
             LINFO << "Skipping zero-length logical partition: " << GetPartitionName(partition);
             continue;
         }
-        std::string path;
-        if (!CreateLogicalPartition(metadata, partition, false, {}, super_device, &path)) {
+
+        params.partition = &partition;
+
+        std::string ignore_path;
+        if (!CreateLogicalPartition(params, &ignore_path)) {
             LERROR << "Could not create logical partition: " << GetPartitionName(partition);
             return false;
         }
@@ -159,29 +146,58 @@
     return true;
 }
 
-bool CreateLogicalPartition(const std::string& block_device, const LpMetadata& metadata,
-                            const std::string& partition_name, bool force_writable,
-                            const std::chrono::milliseconds& timeout_ms, std::string* path) {
-    for (const auto& partition : metadata.partitions) {
-        if (GetPartitionName(partition) == partition_name) {
-            return CreateLogicalPartition(metadata, partition, force_writable, timeout_ms,
-                                          block_device, path);
+bool CreateLogicalPartition(const CreateLogicalPartitionParams& params, std::string* path) {
+    const LpMetadata* metadata = params.metadata;
+
+    // Read metadata if needed.
+    std::unique_ptr<LpMetadata> local_metadata;
+    if (!metadata) {
+        if (!params.metadata_slot) {
+            LOG(ERROR) << "Either metadata or a metadata slot must be specified.";
+            return false;
+        }
+        auto slot = *params.metadata_slot;
+        if (local_metadata = ReadMetadata(params.block_device, slot); !local_metadata) {
+            LOG(ERROR) << "Could not read partition table for: " << params.block_device;
+            return false;
+        }
+        metadata = local_metadata.get();
+    }
+
+    // Find the partition by name if needed.
+    const LpMetadataPartition* partition = params.partition;
+    if (!partition) {
+        for (const auto& iter : metadata->partitions) {
+            if (GetPartitionName(iter) == params.partition_name) {
+                partition = &iter;
+                break;
+            }
+        }
+        if (!partition) {
+            LERROR << "Could not find any partition with name: " << params.partition_name;
+            return false;
         }
     }
-    LERROR << "Could not find any partition with name: " << partition_name;
-    return false;
-}
 
-bool CreateLogicalPartition(const std::string& block_device, uint32_t metadata_slot,
-                            const std::string& partition_name, bool force_writable,
-                            const std::chrono::milliseconds& timeout_ms, std::string* path) {
-    auto metadata = ReadMetadata(block_device.c_str(), metadata_slot);
-    if (!metadata) {
-        LOG(ERROR) << "Could not read partition table.";
-        return true;
+    DmTable table;
+    if (!CreateDmTable(*metadata, *partition, params.block_device, &table)) {
+        return false;
     }
-    return CreateLogicalPartition(block_device, *metadata.get(), partition_name, force_writable,
-                                  timeout_ms, path);
+    if (params.force_writable) {
+        table.set_readonly(false);
+    }
+
+    std::string device_name = params.device_name;
+    if (device_name.empty()) {
+        device_name = GetPartitionName(*partition);
+    }
+
+    DeviceMapper& dm = DeviceMapper::Instance();
+    if (!dm.CreateDevice(device_name, table, path, params.timeout_ms)) {
+        return false;
+    }
+    LINFO << "Created logical partition " << device_name << " on device " << *path;
+    return true;
 }
 
 bool UnmapDevice(const std::string& name) {
diff --git a/fs_mgr/fs_mgr_overlayfs.cpp b/fs_mgr/fs_mgr_overlayfs.cpp
index 35ff43d..d7ea81d 100644
--- a/fs_mgr/fs_mgr_overlayfs.cpp
+++ b/fs_mgr/fs_mgr_overlayfs.cpp
@@ -961,9 +961,16 @@
     }
 
     if (changed || partition_create) {
-        if (!CreateLogicalPartition(super_device, slot_number, partition_name, true, 10s,
-                                    scratch_device))
+        CreateLogicalPartitionParams params = {
+                .block_device = super_device,
+                .metadata_slot = slot_number,
+                .partition_name = partition_name,
+                .force_writable = true,
+                .timeout_ms = 10s,
+        };
+        if (!CreateLogicalPartition(params, scratch_device)) {
             return false;
+        }
 
         if (change) *change = true;
     }
@@ -1184,11 +1191,15 @@
     if ((mount_point != nullptr) && !fs_mgr_overlayfs_already_mounted(kScratchMountPoint, false)) {
         auto scratch_device = fs_mgr_overlayfs_scratch_device();
         if (scratch_device.empty()) {
-            auto slot_number = fs_mgr_overlayfs_slot_number();
-            auto super_device = fs_mgr_overlayfs_super_device(slot_number);
-            const auto partition_name = android::base::Basename(kScratchMountPoint);
-            CreateLogicalPartition(super_device, slot_number, partition_name, true, 10s,
-                                   &scratch_device);
+            auto metadata_slot = fs_mgr_overlayfs_slot_number();
+            CreateLogicalPartitionParams params = {
+                    .block_device = fs_mgr_overlayfs_super_device(metadata_slot),
+                    .metadata_slot = metadata_slot,
+                    .partition_name = android::base::Basename(kScratchMountPoint),
+                    .force_writable = true,
+                    .timeout_ms = 10s,
+            };
+            CreateLogicalPartition(params, &scratch_device);
         }
         mount_scratch = fs_mgr_overlayfs_mount_scratch(scratch_device,
                                                        fs_mgr_overlayfs_scratch_mount_type());
diff --git a/fs_mgr/include/fs_mgr_dm_linear.h b/fs_mgr/include/fs_mgr_dm_linear.h
index a1dc2dc..6eb541c 100644
--- a/fs_mgr/include/fs_mgr_dm_linear.h
+++ b/fs_mgr/include/fs_mgr_dm_linear.h
@@ -29,6 +29,7 @@
 
 #include <chrono>
 #include <memory>
+#include <optional>
 #include <string>
 #include <vector>
 
@@ -49,22 +50,32 @@
 // method for ReadMetadata and CreateLogicalPartitions.
 bool CreateLogicalPartitions(const std::string& block_device);
 
-// Create a block device for a single logical partition, given metadata and
-// the partition name. On success, a path to the partition's block device is
-// returned. If |force_writable| is true, the "readonly" flag will be ignored
-// so the partition can be flashed.
-//
-// If |timeout_ms| is non-zero, then CreateLogicalPartition will block for the
-// given amount of time until the path returned in |path| is available.
-bool CreateLogicalPartition(const std::string& block_device, uint32_t metadata_slot,
-                            const std::string& partition_name, bool force_writable,
-                            const std::chrono::milliseconds& timeout_ms, std::string* path);
+struct CreateLogicalPartitionParams {
+    // Block device of the super partition.
+    std::string block_device;
 
-// Same as above, but with a given metadata object. Care should be taken that
-// the metadata represents a valid partition layout.
-bool CreateLogicalPartition(const std::string& block_device, const LpMetadata& metadata,
-                            const std::string& partition_name, bool force_writable,
-                            const std::chrono::milliseconds& timeout_ms, std::string* path);
+    // If |metadata| is null, the slot will be read using |metadata_slot|.
+    const LpMetadata* metadata = nullptr;
+    std::optional<uint32_t> metadata_slot;
+
+    // If |partition| is not set, it will be found via |partition_name|.
+    const LpMetadataPartition* partition = nullptr;
+    std::string partition_name;
+
+    // Force the device to be read-write even if it was specified as readonly
+    // in the metadata.
+    bool force_writable = false;
+
+    // If |timeout_ms| is non-zero, then CreateLogicalPartition will block for
+    // the given amount of time until the path returned in |path| is available.
+    std::chrono::milliseconds timeout_ms = {};
+
+    // If this is non-empty, it will override the device mapper name (by
+    // default the partition name will be used).
+    std::string device_name;
+};
+
+bool CreateLogicalPartition(const CreateLogicalPartitionParams& params, std::string* path);
 
 // Destroy the block device for a logical partition, by name. If |timeout_ms|
 // is non-zero, then this will block until the device path has been unlinked.
diff --git a/fs_mgr/liblp/builder.cpp b/fs_mgr/liblp/builder.cpp
index 777743c..6dc288e 100644
--- a/fs_mgr/liblp/builder.cpp
+++ b/fs_mgr/liblp/builder.cpp
@@ -140,7 +140,7 @@
     }
     if (opener) {
         for (size_t i = 0; i < builder->block_devices_.size(); i++) {
-            std::string partition_name = GetBlockDevicePartitionName(builder->block_devices_[i]);
+            std::string partition_name = builder->GetBlockDevicePartitionName(i);
             BlockDeviceInfo device_info;
             if (opener->GetInfo(partition_name, &device_info)) {
                 builder->UpdateBlockDeviceInfo(i, device_info);
@@ -164,7 +164,7 @@
     // name and system properties.
     // See comments for UpdateMetadataForOtherSuper.
     auto super_device = GetMetadataSuperBlockDevice(*metadata.get());
-    if (GetBlockDevicePartitionName(*super_device) != "super" &&
+    if (android::fs_mgr::GetBlockDevicePartitionName(*super_device) != "super" &&
         IsRetrofitDynamicPartitionsDevice()) {
         if (!UpdateMetadataForOtherSuper(metadata.get(), source_slot_number, target_slot_number)) {
             return nullptr;
@@ -192,7 +192,8 @@
     // Translate block devices.
     auto source_block_devices = std::move(metadata->block_devices);
     for (const auto& source_block_device : source_block_devices) {
-        std::string partition_name = GetBlockDevicePartitionName(source_block_device);
+        std::string partition_name =
+                android::fs_mgr::GetBlockDevicePartitionName(source_block_device);
         std::string slot_suffix = GetPartitionSlotSuffix(partition_name);
         if (slot_suffix.empty() || slot_suffix != source_slot_suffix) {
             // This should never happen. It means that the source metadata
@@ -375,7 +376,7 @@
             block_devices_.emplace_back(out);
         }
     }
-    if (GetBlockDevicePartitionName(block_devices_[0]) != super_partition) {
+    if (GetBlockDevicePartitionName(0) != super_partition) {
         LERROR << "No super partition was specified.";
         return false;
     }
@@ -458,7 +459,7 @@
     return nullptr;
 }
 
-PartitionGroup* MetadataBuilder::FindGroup(const std::string& group_name) {
+PartitionGroup* MetadataBuilder::FindGroup(std::string_view group_name) {
     for (const auto& group : groups_) {
         if (group->name() == group_name) {
             return group.get();
@@ -582,8 +583,7 @@
     CHECK_NE(sectors_per_block, 0);
     CHECK(sectors_needed % sectors_per_block == 0);
 
-    if (IsABDevice() && !IsRetrofitMetadata() &&
-        GetPartitionSlotSuffix(partition->name()) == "_b") {
+    if (IsABDevice() && ShouldHalveSuper() && GetPartitionSlotSuffix(partition->name()) == "_b") {
         // Allocate "a" partitions top-down and "b" partitions bottom-up, to
         // minimize fragmentation during OTA.
         free_regions = PrioritizeSecondHalfOfSuper(free_regions);
@@ -849,7 +849,7 @@
 bool MetadataBuilder::FindBlockDeviceByName(const std::string& partition_name,
                                             uint32_t* index) const {
     for (size_t i = 0; i < block_devices_.size(); i++) {
-        if (GetBlockDevicePartitionName(block_devices_[i]) == partition_name) {
+        if (GetBlockDevicePartitionName(i) == partition_name) {
             *index = i;
             return true;
         }
@@ -974,7 +974,8 @@
     // Note: we don't compare alignment, since it's a performance thing and
     // won't affect whether old extents continue to work.
     return first.first_logical_sector == second.first_logical_sector && first.size == second.size &&
-           GetBlockDevicePartitionName(first) == GetBlockDevicePartitionName(second);
+           android::fs_mgr::GetBlockDevicePartitionName(first) ==
+                   android::fs_mgr::GetBlockDevicePartitionName(second);
 }
 
 bool MetadataBuilder::ImportPartitions(const LpMetadata& metadata,
@@ -1056,8 +1057,9 @@
                                                             false);
 }
 
-bool MetadataBuilder::IsRetrofitMetadata() const {
-    return GetBlockDevicePartitionName(block_devices_[0]) != LP_METADATA_DEFAULT_PARTITION_NAME;
+bool MetadataBuilder::ShouldHalveSuper() const {
+    return GetBlockDevicePartitionName(0) == LP_METADATA_DEFAULT_PARTITION_NAME &&
+           !IPropertyFetcher::GetInstance()->GetBoolProperty("ro.virtual_ab.enabled", false);
 }
 
 bool MetadataBuilder::AddLinearExtent(Partition* partition, const std::string& block_device,
@@ -1083,7 +1085,7 @@
     return partitions;
 }
 
-bool MetadataBuilder::ChangePartitionGroup(Partition* partition, const std::string& group_name) {
+bool MetadataBuilder::ChangePartitionGroup(Partition* partition, std::string_view group_name) {
     if (!FindGroup(group_name)) {
         LERROR << "Partition cannot change to unknown group: " << group_name;
         return false;
@@ -1121,5 +1123,11 @@
     return true;
 }
 
+std::string MetadataBuilder::GetBlockDevicePartitionName(uint64_t index) const {
+    return index < block_devices_.size()
+                   ? android::fs_mgr::GetBlockDevicePartitionName(block_devices_[index])
+                   : "";
+}
+
 }  // namespace fs_mgr
 }  // namespace android
diff --git a/fs_mgr/liblp/builder_test.cpp b/fs_mgr/liblp/builder_test.cpp
index 6d27873..b3c13e6 100644
--- a/fs_mgr/liblp/builder_test.cpp
+++ b/fs_mgr/liblp/builder_test.cpp
@@ -35,7 +35,7 @@
     IPropertyFetcher::OverrideForTesting(std::make_unique<NiceMock<MockPropertyFetcher>>());
 }
 
-MockPropertyFetcher* GetMockedInstance() {
+MockPropertyFetcher* GetMockedPropertyFetcher() {
     return static_cast<MockPropertyFetcher*>(IPropertyFetcher::GetInstance());
 }
 
@@ -789,7 +789,7 @@
 
     // A and B slots should be allocated from separate halves of the partition,
     // to mitigate allocating too many extents. (b/120433288)
-    ON_CALL(*GetMockedInstance(), GetProperty("ro.boot.slot_suffix", _))
+    ON_CALL(*GetMockedPropertyFetcher(), GetProperty("ro.boot.slot_suffix", _))
             .WillByDefault(Return("_a"));
 
     auto builder = MetadataBuilder::New(device_info, 65536, 2);
diff --git a/fs_mgr/liblp/include/liblp/builder.h b/fs_mgr/liblp/include/liblp/builder.h
index 3b229bd..c3f6e91 100644
--- a/fs_mgr/liblp/include/liblp/builder.h
+++ b/fs_mgr/liblp/include/liblp/builder.h
@@ -128,7 +128,7 @@
 
   private:
     void ShrinkTo(uint64_t aligned_size);
-    void set_group_name(const std::string& group_name) { group_name_ = group_name; }
+    void set_group_name(std::string_view group_name) { group_name_ = group_name; }
 
     std::string name_;
     std::string group_name_;
@@ -227,7 +227,7 @@
     Partition* FindPartition(const std::string& name);
 
     // Find a group by name. If no group is found, nullptr is returned.
-    PartitionGroup* FindGroup(const std::string& name);
+    PartitionGroup* FindGroup(std::string_view name);
 
     // Add a predetermined extent to a partition.
     bool AddLinearExtent(Partition* partition, const std::string& block_device,
@@ -252,7 +252,7 @@
     // the metadata is exported, to avoid errors during potential group and
     // size shuffling operations. This will return false if the new group does
     // not exist.
-    bool ChangePartitionGroup(Partition* partition, const std::string& group_name);
+    bool ChangePartitionGroup(Partition* partition, std::string_view group_name);
 
     // Changes the size of a partition group. Size constraints will not be
     // checked until metadata is exported, to avoid errors during group
@@ -287,6 +287,9 @@
     // Return true if a block device is found, else false.
     bool HasBlockDevice(const std::string& partition_name) const;
 
+    // Return the name of the block device at |index|.
+    std::string GetBlockDevicePartitionName(uint64_t index) const;
+
   private:
     MetadataBuilder();
     MetadataBuilder(const MetadataBuilder&) = delete;
@@ -314,8 +317,8 @@
     // Return true if the device is retrofitting dynamic partitions.
     static bool IsRetrofitDynamicPartitionsDevice();
 
-    // Return true if "this" metadata represents a metadata on a retrofit device.
-    bool IsRetrofitMetadata() const;
+    // Return true if _b partitions should be prioritized at the second half of the device.
+    bool ShouldHalveSuper() const;
 
     bool ValidatePartitionGroups() const;
 
diff --git a/fs_mgr/liblp/io_test.cpp b/fs_mgr/liblp/io_test.cpp
index 2990863..3dace25 100644
--- a/fs_mgr/liblp/io_test.cpp
+++ b/fs_mgr/liblp/io_test.cpp
@@ -668,7 +668,7 @@
 }
 
 TEST(liblp, UpdateRetrofit) {
-    ON_CALL(*GetMockedInstance(), GetBoolProperty("ro.boot.dynamic_partitions_retrofit", _))
+    ON_CALL(*GetMockedPropertyFetcher(), GetBoolProperty("ro.boot.dynamic_partitions_retrofit", _))
             .WillByDefault(Return(true));
 
     unique_ptr<MetadataBuilder> builder = CreateDefaultBuilder();
@@ -700,7 +700,7 @@
 }
 
 TEST(liblp, UpdateNonRetrofit) {
-    ON_CALL(*GetMockedInstance(), GetBoolProperty("ro.boot.dynamic_partitions_retrofit", _))
+    ON_CALL(*GetMockedPropertyFetcher(), GetBoolProperty("ro.boot.dynamic_partitions_retrofit", _))
             .WillByDefault(Return(false));
 
     unique_fd fd = CreateFlashedDisk();
diff --git a/fs_mgr/liblp/mock_property_fetcher.h b/fs_mgr/liblp/mock_property_fetcher.h
index eb91de2..0c28710 100644
--- a/fs_mgr/liblp/mock_property_fetcher.h
+++ b/fs_mgr/liblp/mock_property_fetcher.h
@@ -44,4 +44,4 @@
 }  // namespace fs_mgr
 }  // namespace android
 
-android::fs_mgr::MockPropertyFetcher* GetMockedInstance();
+android::fs_mgr::MockPropertyFetcher* GetMockedPropertyFetcher();
diff --git a/fs_mgr/libsnapshot/Android.bp b/fs_mgr/libsnapshot/Android.bp
index 746987e..128a4f9 100644
--- a/fs_mgr/libsnapshot/Android.bp
+++ b/fs_mgr/libsnapshot/Android.bp
@@ -32,9 +32,14 @@
     whole_static_libs: [
         "libext2_uuid",
         "libext4_utils",
-        "libfiemap",
         "libfstab",
     ],
+    header_libs: [
+        "libfiemap_headers",
+    ],
+    export_header_lib_headers: [
+        "libfiemap_headers",
+    ],
     export_include_dirs: ["include"],
 }
 
@@ -49,7 +54,7 @@
     name: "libsnapshot",
     defaults: ["libsnapshot_defaults"],
     srcs: [":libsnapshot_sources"],
-    static_libs: [
+    whole_static_libs: [
         "libfiemap_binder",
     ],
 }
@@ -59,6 +64,9 @@
     defaults: ["libsnapshot_defaults"],
     srcs: [":libsnapshot_sources"],
     recovery_available: true,
+    whole_static_libs: [
+        "libfiemap_passthrough",
+    ],
 }
 
 cc_test {
diff --git a/fs_mgr/libsnapshot/include/libsnapshot/snapshot.h b/fs_mgr/libsnapshot/include/libsnapshot/snapshot.h
index 4fb6808..1630ee5 100644
--- a/fs_mgr/libsnapshot/include/libsnapshot/snapshot.h
+++ b/fs_mgr/libsnapshot/include/libsnapshot/snapshot.h
@@ -219,9 +219,12 @@
     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 {
-        std::string state;
+        SnapshotState state;
         uint64_t device_size;
         uint64_t snapshot_size;
         // These are non-zero when merging.
diff --git a/fs_mgr/libsnapshot/snapshot.cpp b/fs_mgr/libsnapshot/snapshot.cpp
index 75a1f26..fef5c06 100644
--- a/fs_mgr/libsnapshot/snapshot.cpp
+++ b/fs_mgr/libsnapshot/snapshot.cpp
@@ -165,7 +165,7 @@
     // actual backing image. This is harmless, since it'll get removed when
     // CancelUpdate is called.
     SnapshotStatus status = {
-            .state = "created",
+            .state = SnapshotState::Created,
             .device_size = device_size,
             .snapshot_size = snapshot_size,
     };
@@ -190,7 +190,7 @@
     if (!ReadSnapshotStatus(lock, name, &status)) {
         return false;
     }
-    if (status.state == "merge-completed") {
+    if (status.state == SnapshotState::MergeCompleted) {
         LOG(ERROR) << "Should not create a snapshot device for " << name
                    << " after merging has completed.";
         return false;
@@ -313,7 +313,8 @@
     // There may be an extra device, since the kernel doesn't let us have a
     // snapshot and linear target in the same table.
     auto dm_name = GetSnapshotDeviceName(name, status);
-    if (name != dm_name && !dm.DeleteDevice(dm_name)) {
+    if (name != dm_name && dm.GetState(dm_name) != DmDeviceState::INVALID &&
+        !dm.DeleteDevice(dm_name)) {
         LOG(ERROR) << "Could not delete inner snapshot device: " << dm_name;
         return false;
     }
@@ -422,8 +423,8 @@
     if (!ReadSnapshotStatus(lock, name, &status)) {
         return false;
     }
-    if (status.state != "created") {
-        LOG(WARNING) << "Snapshot " << name << " has unexpected state: " << status.state;
+    if (status.state != SnapshotState::Created) {
+        LOG(WARNING) << "Snapshot " << name << " has unexpected state: " << to_string(status.state);
     }
 
     // After this, we return true because we technically did switch to a merge
@@ -433,7 +434,7 @@
         return false;
     }
 
-    status.state = "merging";
+    status.state = SnapshotState::Merging;
 
     DmTargetSnapshot::Status dm_status;
     if (!QuerySnapshotStatus(dm_name, nullptr, &dm_status)) {
@@ -657,7 +658,7 @@
     // 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)) {
+    if (snapshot_status.state == SnapshotState::MergeCompleted && !IsSnapshotDevice(dm_name)) {
         // NB: It's okay if this fails now, we gave cleanup our best effort.
         OnSnapshotMergeComplete(lock, name, snapshot_status);
         return UpdateState::MergeCompleted;
@@ -678,7 +679,7 @@
 
     // These two values are equal when merging is complete.
     if (status.sectors_allocated != status.metadata_sectors) {
-        if (snapshot_status.state == "merge-complete") {
+        if (snapshot_status.state == SnapshotState::MergeCompleted) {
             LOG(ERROR) << "Snapshot " << name << " is merging after being marked merge-complete.";
             return UpdateState::MergeFailed;
         }
@@ -693,7 +694,7 @@
     // 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";
+    snapshot_status.state = SnapshotState::MergeCompleted;
     if (!WriteSnapshotStatus(lock, name, snapshot_status)) {
         return UpdateState::MergeFailed;
     }
@@ -1068,7 +1069,16 @@
         return false;
     }
 
-    status->state = pieces[0];
+    if (pieces[0] == "created") {
+        status->state = SnapshotState::Created;
+    } else if (pieces[0] == "merging") {
+        status->state = SnapshotState::Merging;
+    } else if (pieces[0] == "merge-completed") {
+        status->state = SnapshotState::MergeCompleted;
+    } else {
+        LOG(ERROR) << "Unrecognized state " << pieces[0] << " for snapshot: " << name;
+    }
+
     if (!android::base::ParseUint(pieces[1], &status->device_size)) {
         LOG(ERROR) << "Invalid device size in status line for: " << path;
         return false;
@@ -1088,6 +1098,20 @@
     return true;
 }
 
+std::string SnapshotManager::to_string(SnapshotState state) {
+    switch (state) {
+        case SnapshotState::Created:
+            return "created";
+        case SnapshotState::Merging:
+            return "merging";
+        case SnapshotState::MergeCompleted:
+            return "merge-completed";
+        default:
+            LOG(ERROR) << "Unknown snapshot state: " << (int)state;
+            return "unknown";
+    }
+}
+
 bool SnapshotManager::WriteSnapshotStatus(LockedFile* lock, const std::string& name,
                                           const SnapshotStatus& status) {
     // The caller must take an exclusive lock to modify snapshots.
@@ -1102,7 +1126,7 @@
     }
 
     std::vector<std::string> pieces = {
-            status.state,
+            to_string(status.state),
             std::to_string(status.device_size),
             std::to_string(status.snapshot_size),
             std::to_string(status.sectors_allocated),
diff --git a/fs_mgr/libsnapshot/snapshot_test.cpp b/fs_mgr/libsnapshot/snapshot_test.cpp
index 34ea331..438ec2f 100644
--- a/fs_mgr/libsnapshot/snapshot_test.cpp
+++ b/fs_mgr/libsnapshot/snapshot_test.cpp
@@ -36,6 +36,7 @@
 using android::base::unique_fd;
 using android::dm::DeviceMapper;
 using android::dm::DmDeviceState;
+using android::fiemap::IImageManager;
 using namespace std::chrono_literals;
 using namespace std::string_literals;
 
@@ -48,9 +49,11 @@
     void set_slot_suffix(const std::string& suffix) { slot_suffix_ = suffix; }
 
   private:
-    std::string slot_suffix_;
+    std::string slot_suffix_ = "_a";
 };
 
+// These are not reset between each test because it's expensive to reconnect
+// to gsid each time.
 std::unique_ptr<SnapshotManager> sm;
 TestDeviceInfo* test_device = nullptr;
 
@@ -60,16 +63,14 @@
 
   protected:
     void SetUp() override {
+        ASSERT_TRUE(sm->EnsureImageManager());
+        image_manager_ = sm->image_manager();
+
         test_device->set_slot_suffix("_a");
 
-        if (sm->GetUpdateState() != UpdateState::None) {
-            CleanupTestArtifacts();
-        }
-        ASSERT_TRUE(sm->BeginUpdate());
-        ASSERT_TRUE(sm->EnsureImageManager());
+        CleanupTestArtifacts();
 
-        image_manager_ = sm->image_manager();
-        ASSERT_NE(image_manager_, nullptr);
+        ASSERT_TRUE(sm->BeginUpdate());
     }
 
     void TearDown() override {
@@ -81,20 +82,25 @@
     void CleanupTestArtifacts() {
         // Normally cancelling inside a merge is not allowed. Since these
         // are tests, we don't care, destroy everything that might exist.
+        // Note we hardcode this list because of an annoying quirk: when
+        // completing a merge, the snapshot stops existing, so we can't
+        // get an accurate list to remove.
+        lock_ = nullptr;
+
         std::vector<std::string> snapshots = {"test-snapshot"};
         for (const auto& snapshot : snapshots) {
             DeleteSnapshotDevice(snapshot);
-            temp_images_.emplace_back(snapshot + "-cow");
+            DeleteBackingImage(image_manager_, 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);
+        // Remove all images. We hardcode the list of names since this can run
+        // before the test (when cleaning up from a crash).
+        std::vector<std::string> temp_partitions = {"base-device"};
+        for (const auto& partition : temp_partitions) {
+            DeleteBackingImage(image_manager_, partition);
         }
 
         if (sm->GetUpdateState() != UpdateState::None) {
@@ -112,23 +118,29 @@
         if (!image_manager_->CreateBackingImage(name, size, false)) {
             return false;
         }
-        temp_images_.emplace_back(name);
         return image_manager_->MapImageDevice(name, 10s, path);
     }
 
-    bool DeleteSnapshotDevice(const std::string& snapshot) {
+    void DeleteSnapshotDevice(const std::string& snapshot) {
         if (dm_.GetState(snapshot) != DmDeviceState::INVALID) {
-            if (!dm_.DeleteDevice(snapshot)) return false;
+            ASSERT_TRUE(dm_.DeleteDevice(snapshot));
         }
         if (dm_.GetState(snapshot + "-inner") != DmDeviceState::INVALID) {
-            if (!dm_.DeleteDevice(snapshot + "-inner")) return false;
+            ASSERT_TRUE(dm_.DeleteDevice(snapshot + "-inner"));
         }
-        return true;
+    }
+
+    void DeleteBackingImage(IImageManager* manager, const std::string& name) {
+        if (manager->IsImageMapped(name)) {
+            ASSERT_TRUE(manager->UnmapImageDevice(name));
+        }
+        if (manager->BackingImageExists(name)) {
+            ASSERT_TRUE(manager->DeleteBackingImage(name));
+        }
     }
 
     DeviceMapper& dm_;
     std::unique_ptr<SnapshotManager::LockedFile> lock_;
-    std::vector<std::string> temp_images_;
     android::fiemap::IImageManager* image_manager_ = nullptr;
 };
 
@@ -148,7 +160,7 @@
     {
         SnapshotManager::SnapshotStatus status;
         ASSERT_TRUE(sm->ReadSnapshotStatus(lock_.get(), "test-snapshot", &status));
-        ASSERT_EQ(status.state, "created");
+        ASSERT_EQ(status.state, SnapshotManager::SnapshotState::Created);
         ASSERT_EQ(status.device_size, kDeviceSize);
         ASSERT_EQ(status.snapshot_size, kDeviceSize);
     }
@@ -279,7 +291,7 @@
     ASSERT_EQ(sm->WaitForMerge(), UpdateState::MergeNeedsReboot);
 
     // Forcefully delete the snapshot device, so it looks like we just rebooted.
-    ASSERT_TRUE(DeleteSnapshotDevice("test-snapshot"));
+    DeleteSnapshotDevice("test-snapshot");
 
     // Map snapshot should fail now, because we're in a merge-complete state.
     ASSERT_TRUE(AcquireLock());