diff --git a/init/first_stage_mount.cpp b/init/first_stage_mount.cpp
index 6b4216f..4a1bc5a 100644
--- a/init/first_stage_mount.cpp
+++ b/init/first_stage_mount.cpp
@@ -77,8 +77,9 @@
     bool InitDevices();
 
   protected:
-    ListenerAction HandleBlockDevice(const std::string& name, const Uevent&);
-    bool InitRequiredDevices();
+    ListenerAction HandleBlockDevice(const std::string& name, const Uevent&,
+                                     std::set<std::string>* required_devices);
+    bool InitRequiredDevices(std::set<std::string> devices);
     bool InitMappedDevice(const std::string& verity_device);
     bool InitDeviceMapper();
     bool CreateLogicalPartitions();
@@ -89,14 +90,14 @@
     bool TrySwitchSystemAsRoot();
     bool TrySkipMountingPartitions();
     bool IsDmLinearEnabled();
-    bool GetDmLinearMetadataDevice();
+    void GetDmLinearMetadataDevice(std::set<std::string>* devices);
     bool InitDmLinearBackingDevices(const android::fs_mgr::LpMetadata& metadata);
     void UseGsiIfPresent();
 
-    ListenerAction UeventCallback(const Uevent& uevent);
+    ListenerAction UeventCallback(const Uevent& uevent, std::set<std::string>* required_devices);
 
     // Pure virtual functions.
-    virtual bool GetDmVerityDevices() = 0;
+    virtual bool GetDmVerityDevices(std::set<std::string>* devices) = 0;
     virtual bool SetUpDmVerity(FstabEntry* fstab_entry) = 0;
 
     bool need_dm_verity_;
@@ -104,7 +105,6 @@
 
     Fstab fstab_;
     std::string lp_metadata_partition_;
-    std::set<std::string> required_devices_partition_names_;
     std::string super_partition_name_;
     std::unique_ptr<DeviceHandler> device_handler_;
     UeventListener uevent_listener_;
@@ -116,7 +116,7 @@
     ~FirstStageMountVBootV1() override = default;
 
   protected:
-    bool GetDmVerityDevices() override;
+    bool GetDmVerityDevices(std::set<std::string>* devices) override;
     bool SetUpDmVerity(FstabEntry* fstab_entry) override;
 };
 
@@ -128,7 +128,7 @@
     ~FirstStageMountVBootV2() override = default;
 
   protected:
-    bool GetDmVerityDevices() override;
+    bool GetDmVerityDevices(std::set<std::string>* devices) override;
     bool SetUpDmVerity(FstabEntry* fstab_entry) override;
     bool InitAvbHandle();
 
@@ -252,7 +252,13 @@
 }
 
 bool FirstStageMount::InitDevices() {
-    return GetDmLinearMetadataDevice() && GetDmVerityDevices() && InitRequiredDevices();
+    std::set<std::string> devices;
+    GetDmLinearMetadataDevice(&devices);
+
+    if (!GetDmVerityDevices(&devices)) {
+        return false;
+    }
+    return InitRequiredDevices(std::move(devices));
 }
 
 bool FirstStageMount::IsDmLinearEnabled() {
@@ -262,45 +268,46 @@
     return false;
 }
 
-bool FirstStageMount::GetDmLinearMetadataDevice() {
+void FirstStageMount::GetDmLinearMetadataDevice(std::set<std::string>* devices) {
     // Add any additional devices required for dm-linear mappings.
     if (!IsDmLinearEnabled()) {
-        return true;
+        return;
     }
 
-    required_devices_partition_names_.emplace(super_partition_name_);
-    return true;
+    devices->emplace(super_partition_name_);
 }
 
-// Creates devices with uevent->partition_name matching one in the member variable
-// required_devices_partition_names_. Found partitions will then be removed from it
-// for the subsequent member function to check which devices are NOT created.
-bool FirstStageMount::InitRequiredDevices() {
+// Creates devices with uevent->partition_name matching ones in the given set.
+// Found partitions will then be removed from it for the subsequent member
+// function to check which devices are NOT created.
+bool FirstStageMount::InitRequiredDevices(std::set<std::string> devices) {
     if (!InitDeviceMapper()) {
         return false;
     }
 
-    if (required_devices_partition_names_.empty()) {
+    if (devices.empty()) {
         return true;
     }
 
-    auto uevent_callback = [this](const Uevent& uevent) { return UeventCallback(uevent); };
+    auto uevent_callback = [&, this](const Uevent& uevent) {
+        return UeventCallback(uevent, &devices);
+    };
     uevent_listener_.RegenerateUevents(uevent_callback);
 
-    // UeventCallback() will remove found partitions from required_devices_partition_names_.
-    // So if it isn't empty here, it means some partitions are not found.
-    if (!required_devices_partition_names_.empty()) {
+    // UeventCallback() will remove found partitions from |devices|. So if it
+    // isn't empty here, it means some partitions are not found.
+    if (!devices.empty()) {
         LOG(INFO) << __PRETTY_FUNCTION__
                   << ": partition(s) not found in /sys, waiting for their uevent(s): "
-                  << android::base::Join(required_devices_partition_names_, ", ");
+                  << android::base::Join(devices, ", ");
         Timer t;
         uevent_listener_.Poll(uevent_callback, 10s);
         LOG(INFO) << "Wait for partitions returned after " << t;
     }
 
-    if (!required_devices_partition_names_.empty()) {
+    if (!devices.empty()) {
         LOG(ERROR) << __PRETTY_FUNCTION__ << ": partition(s) not found after polling timeout: "
-                   << android::base::Join(required_devices_partition_names_, ", ");
+                   << android::base::Join(devices, ", ");
         return false;
     }
 
@@ -333,27 +340,20 @@
 }
 
 bool FirstStageMount::InitDmLinearBackingDevices(const android::fs_mgr::LpMetadata& metadata) {
+    std::set<std::string> devices;
+
     auto partition_names = android::fs_mgr::GetBlockDevicePartitionNames(metadata);
     for (const auto& partition_name : partition_names) {
         // The super partition was found in the earlier pass.
         if (partition_name == super_partition_name_) {
             continue;
         }
-        required_devices_partition_names_.emplace(partition_name);
+        devices.emplace(partition_name);
     }
-    if (required_devices_partition_names_.empty()) {
+    if (devices.empty()) {
         return true;
     }
-
-    auto uevent_callback = [this](const Uevent& uevent) { return UeventCallback(uevent); };
-    uevent_listener_.RegenerateUevents(uevent_callback);
-
-    if (!required_devices_partition_names_.empty()) {
-        LOG(ERROR) << __PRETTY_FUNCTION__ << ": partition(s) not found after polling timeout: "
-                   << android::base::Join(required_devices_partition_names_, ", ");
-        return false;
-    }
-    return true;
+    return InitRequiredDevices(std::move(devices));
 }
 
 bool FirstStageMount::CreateLogicalPartitions() {
@@ -372,6 +372,11 @@
             return false;
         }
         if (sm->NeedSnapshotsInFirstStageMount()) {
+            // When COW images are present for snapshots, they are stored on
+            // the data partition.
+            if (!InitRequiredDevices({"userdata"})) {
+                return false;
+            }
             return sm->CreateLogicalAndSnapshotPartitions(lp_metadata_partition_);
         }
     }
@@ -387,20 +392,21 @@
     return android::fs_mgr::CreateLogicalPartitions(*metadata.get(), lp_metadata_partition_);
 }
 
-ListenerAction FirstStageMount::HandleBlockDevice(const std::string& name, const Uevent& uevent) {
+ListenerAction FirstStageMount::HandleBlockDevice(const std::string& name, const Uevent& uevent,
+                                                  std::set<std::string>* required_devices) {
     // Matches partition name to create device nodes.
     // Both required_devices_partition_names_ and uevent->partition_name have A/B
     // suffix when A/B is used.
-    auto iter = required_devices_partition_names_.find(name);
-    if (iter != required_devices_partition_names_.end()) {
+    auto iter = required_devices->find(name);
+    if (iter != required_devices->end()) {
         LOG(VERBOSE) << __PRETTY_FUNCTION__ << ": found partition: " << *iter;
         if (IsDmLinearEnabled() && name == super_partition_name_) {
             std::vector<std::string> links = device_handler_->GetBlockDeviceSymlinks(uevent);
             lp_metadata_partition_ = links[0];
         }
-        required_devices_partition_names_.erase(iter);
+        required_devices->erase(iter);
         device_handler_->HandleUevent(uevent);
-        if (required_devices_partition_names_.empty()) {
+        if (required_devices->empty()) {
             return ListenerAction::kStop;
         } else {
             return ListenerAction::kContinue;
@@ -409,18 +415,19 @@
     return ListenerAction::kContinue;
 }
 
-ListenerAction FirstStageMount::UeventCallback(const Uevent& uevent) {
+ListenerAction FirstStageMount::UeventCallback(const Uevent& uevent,
+                                               std::set<std::string>* required_devices) {
     // Ignores everything that is not a block device.
     if (uevent.subsystem != "block") {
         return ListenerAction::kContinue;
     }
 
     if (!uevent.partition_name.empty()) {
-        return HandleBlockDevice(uevent.partition_name, uevent);
+        return HandleBlockDevice(uevent.partition_name, uevent, required_devices);
     } else {
         size_t base_idx = uevent.path.rfind('/');
         if (base_idx != std::string::npos) {
-            return HandleBlockDevice(uevent.path.substr(base_idx + 1), uevent);
+            return HandleBlockDevice(uevent.path.substr(base_idx + 1), uevent, required_devices);
         }
     }
     // Not found a partition or find an unneeded partition, continue to find others.
@@ -577,17 +584,7 @@
     const auto devices = fs_mgr_overlayfs_required_devices(&fstab_);
     for (auto const& device : devices) {
         if (android::base::StartsWith(device, "/dev/block/by-name/")) {
-            required_devices_partition_names_.emplace(basename(device.c_str()));
-            auto uevent_callback = [this](const Uevent& uevent) { return UeventCallback(uevent); };
-            uevent_listener_.RegenerateUevents(uevent_callback);
-            if (!required_devices_partition_names_.empty()) {
-                uevent_listener_.Poll(uevent_callback, 10s);
-                if (!required_devices_partition_names_.empty()) {
-                    LOG(ERROR) << __PRETTY_FUNCTION__
-                               << ": partition(s) not found after polling timeout: "
-                               << android::base::Join(required_devices_partition_names_, ", ");
-                }
-            }
+            InitRequiredDevices({basename(device.c_str())});
         } else {
             InitMappedDevice(device);
         }
@@ -641,7 +638,7 @@
     gsi_not_on_userdata_ = (super_name != "userdata");
 }
 
-bool FirstStageMountVBootV1::GetDmVerityDevices() {
+bool FirstStageMountVBootV1::GetDmVerityDevices(std::set<std::string>* devices) {
     need_dm_verity_ = false;
 
     for (const auto& fstab_entry : fstab_) {
@@ -660,7 +657,7 @@
     // Notes that fstab_rec->blk_device has A/B suffix updated by fs_mgr when A/B is used.
     for (const auto& fstab_entry : fstab_) {
         if (!fstab_entry.fs_mgr_flags.logical) {
-            required_devices_partition_names_.emplace(basename(fstab_entry.blk_device.c_str()));
+            devices->emplace(basename(fstab_entry.blk_device.c_str()));
         }
     }
 
@@ -711,7 +708,7 @@
     }
 }
 
-bool FirstStageMountVBootV2::GetDmVerityDevices() {
+bool FirstStageMountVBootV2::GetDmVerityDevices(std::set<std::string>* devices) {
     need_dm_verity_ = false;
 
     std::set<std::string> logical_partitions;
@@ -725,7 +722,7 @@
             // Don't try to find logical partitions via uevent regeneration.
             logical_partitions.emplace(basename(fstab_entry.blk_device.c_str()));
         } else {
-            required_devices_partition_names_.emplace(basename(fstab_entry.blk_device.c_str()));
+            devices->emplace(basename(fstab_entry.blk_device.c_str()));
         }
     }
 
@@ -742,11 +739,11 @@
             if (logical_partitions.count(partition_name)) {
                 continue;
             }
-            // required_devices_partition_names_ is of type std::set so it's not an issue
-            // to emplace a partition twice. e.g., /vendor might be in both places:
+            // devices is of type std::set so it's not an issue to emplace a
+            // partition twice. e.g., /vendor might be in both places:
             //   - device_tree_vbmeta_parts_ = "vbmeta,boot,system,vendor"
             //   - mount_fstab_recs_: /vendor_a
-            required_devices_partition_names_.emplace(partition_name);
+            devices->emplace(partition_name);
         }
     }
     return true;