Merge "Improve codegen slightly when doing FD validity checks"
diff --git a/fs_mgr/liblp/builder.cpp b/fs_mgr/liblp/builder.cpp
index f8c492d..c5d6a3b 100644
--- a/fs_mgr/liblp/builder.cpp
+++ b/fs_mgr/liblp/builder.cpp
@@ -579,12 +579,38 @@
     return true;
 }
 
-bool MetadataBuilder::GrowPartition(Partition* partition, uint64_t aligned_size) {
+Interval Interval::Intersect(const Interval& a, const Interval& b) {
+    Interval ret = a;
+    if (a.device_index != b.device_index) {
+        ret.start = ret.end = a.start;  // set length to 0 to indicate no intersection.
+        return ret;
+    }
+    ret.start = std::max(a.start, b.start);
+    ret.end = std::max(ret.start, std::min(a.end, b.end));
+    return ret;
+}
+
+std::vector<Interval> Interval::Intersect(const std::vector<Interval>& a,
+                                          const std::vector<Interval>& b) {
+    std::vector<Interval> ret;
+    for (const Interval& a_interval : a) {
+        for (const Interval& b_interval : b) {
+            auto intersect = Intersect(a_interval, b_interval);
+            if (intersect.length() > 0) ret.emplace_back(std::move(intersect));
+        }
+    }
+    return ret;
+}
+
+bool MetadataBuilder::GrowPartition(Partition* partition, uint64_t aligned_size,
+                                    const std::vector<Interval>& free_region_hint) {
     uint64_t space_needed = aligned_size - partition->size();
     uint64_t sectors_needed = space_needed / LP_SECTOR_SIZE;
     DCHECK(sectors_needed * LP_SECTOR_SIZE == space_needed);
 
     std::vector<Interval> free_regions = GetFreeRegions();
+    if (!free_region_hint.empty())
+        free_regions = Interval::Intersect(free_regions, free_region_hint);
 
     const uint64_t sectors_per_block = geometry_.logical_block_size / LP_SECTOR_SIZE;
     CHECK_NE(sectors_per_block, 0);
@@ -650,7 +676,7 @@
     return true;
 }
 
-std::vector<MetadataBuilder::Interval> MetadataBuilder::PrioritizeSecondHalfOfSuper(
+std::vector<Interval> MetadataBuilder::PrioritizeSecondHalfOfSuper(
         const std::vector<Interval>& free_list) {
     const auto& super = block_devices_[0];
     uint64_t first_sector = super.first_logical_sector;
@@ -926,7 +952,8 @@
     return true;
 }
 
-bool MetadataBuilder::ResizePartition(Partition* partition, uint64_t requested_size) {
+bool MetadataBuilder::ResizePartition(Partition* partition, uint64_t requested_size,
+                                      const std::vector<Interval>& free_region_hint) {
     // Align the space needed up to the nearest sector.
     uint64_t aligned_size = AlignTo(requested_size, geometry_.logical_block_size);
     uint64_t old_size = partition->size();
@@ -936,7 +963,7 @@
     }
 
     if (aligned_size > old_size) {
-        if (!GrowPartition(partition, aligned_size)) {
+        if (!GrowPartition(partition, aligned_size, free_region_hint)) {
             return false;
         }
     } else if (aligned_size < partition->size()) {
diff --git a/fs_mgr/liblp/builder_test.cpp b/fs_mgr/liblp/builder_test.cpp
index c5b4047..bd41f59 100644
--- a/fs_mgr/liblp/builder_test.cpp
+++ b/fs_mgr/liblp/builder_test.cpp
@@ -887,3 +887,38 @@
     std::set<std::string> partitions_to_keep{"system_a", "vendor_a", "product_a"};
     ASSERT_TRUE(builder->ImportPartitions(*on_disk.get(), partitions_to_keep));
 }
+
+// Interval has operator< defined; it is not appropriate to re-define Interval::operator== that
+// compares device index.
+namespace android {
+namespace fs_mgr {
+bool operator==(const Interval& a, const Interval& b) {
+    return a.device_index == b.device_index && a.start == b.start && a.end == b.end;
+}
+}  // namespace fs_mgr
+}  // namespace android
+
+TEST_F(BuilderTest, Interval) {
+    EXPECT_EQ(0u, Interval::Intersect(Interval(0, 100, 200), Interval(0, 50, 100)).length());
+    EXPECT_EQ(Interval(0, 100, 150),
+              Interval::Intersect(Interval(0, 100, 200), Interval(0, 50, 150)));
+    EXPECT_EQ(Interval(0, 100, 200),
+              Interval::Intersect(Interval(0, 100, 200), Interval(0, 50, 200)));
+    EXPECT_EQ(Interval(0, 100, 200),
+              Interval::Intersect(Interval(0, 100, 200), Interval(0, 50, 250)));
+    EXPECT_EQ(Interval(0, 100, 200),
+              Interval::Intersect(Interval(0, 100, 200), Interval(0, 100, 200)));
+    EXPECT_EQ(Interval(0, 150, 200),
+              Interval::Intersect(Interval(0, 100, 200), Interval(0, 150, 250)));
+    EXPECT_EQ(0u, Interval::Intersect(Interval(0, 100, 200), Interval(0, 200, 250)).length());
+
+    auto v = Interval::Intersect(std::vector<Interval>{Interval(0, 0, 50), Interval(0, 100, 150)},
+                                 std::vector<Interval>{Interval(0, 25, 125)});
+    ASSERT_EQ(2, v.size());
+    EXPECT_EQ(Interval(0, 25, 50), v[0]);
+    EXPECT_EQ(Interval(0, 100, 125), v[1]);
+
+    EXPECT_EQ(0u, Interval::Intersect(std::vector<Interval>{Interval(0, 0, 50)},
+                                      std::vector<Interval>{Interval(0, 100, 150)})
+                          .size());
+}
diff --git a/fs_mgr/liblp/include/liblp/builder.h b/fs_mgr/liblp/include/liblp/builder.h
index 5ab42f5..6f2ab75 100644
--- a/fs_mgr/liblp/include/liblp/builder.h
+++ b/fs_mgr/liblp/include/liblp/builder.h
@@ -138,6 +138,33 @@
     uint64_t size_;
 };
 
+// An interval in the metadata. This is similar to a LinearExtent with one difference.
+// LinearExtent represents a "used" region in the metadata, while Interval can also represent
+// an "unused" region.
+struct Interval {
+    uint32_t device_index;
+    uint64_t start;
+    uint64_t end;
+
+    Interval(uint32_t device_index, uint64_t start, uint64_t end)
+        : device_index(device_index), start(start), end(end) {}
+    uint64_t length() const { return end - start; }
+
+    // Note: the device index is not included in sorting (intervals are
+    // sorted in per-device lists).
+    bool operator<(const Interval& other) const {
+        return (start == other.start) ? end < other.end : start < other.start;
+    }
+
+    // Intersect |a| with |b|.
+    // If no intersection, result has 0 length().
+    static Interval Intersect(const Interval& a, const Interval& b);
+
+    // Intersect two lists of intervals, and store result to |a|.
+    static std::vector<Interval> Intersect(const std::vector<Interval>& a,
+                                           const std::vector<Interval>& b);
+};
+
 class MetadataBuilder {
   public:
     // Construct an empty logical partition table builder given the specified
@@ -244,7 +271,11 @@
     //
     // Note, this is an in-memory operation, and it does not alter the
     // underlying filesystem or contents of the partition on disk.
-    bool ResizePartition(Partition* partition, uint64_t requested_size);
+    //
+    // If |free_region_hint| is not empty, it will only try to allocate extents
+    // in regions within the list.
+    bool ResizePartition(Partition* partition, uint64_t requested_size,
+                         const std::vector<Interval>& free_region_hint = {});
 
     // Return the list of partitions belonging to a group.
     std::vector<Partition*> ListPartitionsInGroup(const std::string& group_name);
@@ -291,6 +322,9 @@
     // Return the name of the block device at |index|.
     std::string GetBlockDevicePartitionName(uint64_t index) const;
 
+    // Return the list of free regions not occupied by extents in the metadata.
+    std::vector<Interval> GetFreeRegions() const;
+
   private:
     MetadataBuilder();
     MetadataBuilder(const MetadataBuilder&) = delete;
@@ -300,7 +334,8 @@
     bool Init(const std::vector<BlockDeviceInfo>& block_devices, const std::string& super_partition,
               uint32_t metadata_max_size, uint32_t metadata_slot_count);
     bool Init(const LpMetadata& metadata);
-    bool GrowPartition(Partition* partition, uint64_t aligned_size);
+    bool GrowPartition(Partition* partition, uint64_t aligned_size,
+                       const std::vector<Interval>& free_region_hint);
     void ShrinkPartition(Partition* partition, uint64_t aligned_size);
     uint64_t AlignSector(const LpMetadataBlockDevice& device, uint64_t sector) const;
     uint64_t TotalSizeOfGroup(PartitionGroup* group) const;
@@ -323,22 +358,6 @@
 
     bool ValidatePartitionGroups() const;
 
-    struct Interval {
-        uint32_t device_index;
-        uint64_t start;
-        uint64_t end;
-
-        Interval(uint32_t device_index, uint64_t start, uint64_t end)
-            : device_index(device_index), start(start), end(end) {}
-        uint64_t length() const { return end - start; }
-
-        // Note: the device index is not included in sorting (intervals are
-        // sorted in per-device lists).
-        bool operator<(const Interval& other) const {
-            return (start == other.start) ? end < other.end : start < other.start;
-        }
-    };
-    std::vector<Interval> GetFreeRegions() const;
     bool IsAnyRegionCovered(const std::vector<Interval>& regions,
                             const LinearExtent& candidate) const;
     bool IsAnyRegionAllocated(const LinearExtent& candidate) const;
diff --git a/fs_mgr/libsnapshot/include/libsnapshot/snapshot.h b/fs_mgr/libsnapshot/include/libsnapshot/snapshot.h
index 1f3828e..6f0e804 100644
--- a/fs_mgr/libsnapshot/include/libsnapshot/snapshot.h
+++ b/fs_mgr/libsnapshot/include/libsnapshot/snapshot.h
@@ -329,6 +329,10 @@
     std::string GetSnapshotDeviceName(const std::string& snapshot_name,
                                       const SnapshotStatus& status);
 
+    // Map the base device, COW devices, and snapshot device.
+    bool MapPartitionWithSnapshot(LockedFile* lock, CreateLogicalPartitionParams params,
+                                  std::string* path);
+
     std::string gsid_dir_;
     std::string metadata_dir_;
     std::unique_ptr<IDeviceInfo> device_;
diff --git a/fs_mgr/libsnapshot/snapshot.cpp b/fs_mgr/libsnapshot/snapshot.cpp
index 3820523..7e57421 100644
--- a/fs_mgr/libsnapshot/snapshot.cpp
+++ b/fs_mgr/libsnapshot/snapshot.cpp
@@ -19,6 +19,7 @@
 #include <sys/types.h>
 #include <sys/unistd.h>
 
+#include <optional>
 #include <thread>
 #include <unordered_set>
 
@@ -217,8 +218,27 @@
     }
 
     auto cow_name = GetCowName(name);
-    int cow_flags = IImageManager::CREATE_IMAGE_ZERO_FILL;
-    return images_->CreateBackingImage(cow_name, cow_size, cow_flags);
+    int cow_flags = IImageManager::CREATE_IMAGE_DEFAULT;
+    if (!images_->CreateBackingImage(cow_name, cow_size, cow_flags)) {
+        return false;
+    }
+
+    // when the kernel creates a persistent dm-snapshot, it requires a CoW file
+    // to store the modifications. The kernel interface does not specify how
+    // the CoW is used, and there is no standard associated.
+    // By looking at the current implementation, the CoW file is treated as:
+    // - a _NEW_ snapshot if its first 32 bits are zero, so the newly created
+    // dm-snapshot device will look like a perfect copy of the origin device;
+    // - an _EXISTING_ snapshot if the first 32 bits are equal to a
+    // kernel-specified magic number and the CoW file metadata is set as valid,
+    // so it can be used to resume the last state of a snapshot device;
+    // - an _INVALID_ snapshot otherwise.
+    // To avoid zero-filling the whole CoW file when a new dm-snapshot is
+    // created, here we zero-fill only the first 32 bits. This is a temporary
+    // 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);
 }
 
 bool SnapshotManager::MapSnapshot(LockedFile* lock, const std::string& name,
@@ -1108,22 +1128,6 @@
     auto lock = LockExclusive();
     if (!lock) return false;
 
-    std::vector<std::string> snapshot_list;
-    if (!ListSnapshots(lock.get(), &snapshot_list)) {
-        return false;
-    }
-
-    std::unordered_set<std::string> live_snapshots;
-    for (const auto& snapshot : snapshot_list) {
-        SnapshotStatus status;
-        if (!ReadSnapshotStatus(lock.get(), snapshot, &status)) {
-            return false;
-        }
-        if (status.state != SnapshotState::MergeCompleted) {
-            live_snapshots.emplace(snapshot);
-        }
-    }
-
     const auto& opener = device_->GetPartitionOpener();
     uint32_t slot = SlotNumberForSlotSuffix(device_->GetSlotSuffix());
     auto metadata = android::fs_mgr::ReadMetadata(opener, super_device, slot);
@@ -1132,59 +1136,102 @@
         return false;
     }
 
-    // Map logical partitions.
-    auto& dm = DeviceMapper::Instance();
     for (const auto& partition : metadata->partitions) {
-        auto partition_name = GetPartitionName(partition);
-        if (!partition.num_extents) {
-            LOG(INFO) << "Skipping zero-length logical partition: " << partition_name;
-            continue;
-        }
-
-        if (!(partition.attributes & LP_PARTITION_ATTR_UPDATED)) {
-            LOG(INFO) << "Detected re-flashing of partition, will skip snapshot: "
-                      << partition_name;
-            live_snapshots.erase(partition_name);
-        }
-
         CreateLogicalPartitionParams params = {
                 .block_device = super_device,
                 .metadata = metadata.get(),
                 .partition = &partition,
                 .partition_opener = &opener,
         };
-
-        if (auto iter = live_snapshots.find(partition_name); iter != live_snapshots.end()) {
-            // If the device has a snapshot, it'll need to be writable, and
-            // we'll need to create the logical partition with a marked-up name
-            // (since the snapshot will use the partition name).
-            params.force_writable = true;
-            params.device_name = GetBaseDeviceName(partition_name);
-        }
-
         std::string ignore_path;
-        if (!CreateLogicalPartition(params, &ignore_path)) {
-            LOG(ERROR) << "Could not create logical partition " << partition_name << " as device "
-                       << params.GetDeviceName();
-            return false;
-        }
-        if (!params.force_writable) {
-            // No snapshot.
-            continue;
-        }
-
-        // We don't have ueventd in first-stage init, so use device major:minor
-        // strings instead.
-        std::string base_device;
-        if (!dm.GetDeviceString(params.GetDeviceName(), &base_device)) {
-            LOG(ERROR) << "Could not determine major/minor for: " << params.GetDeviceName();
-            return false;
-        }
-        if (!MapSnapshot(lock.get(), partition_name, base_device, {}, &ignore_path)) {
-            LOG(ERROR) << "Could not map snapshot for partition: " << partition_name;
+        if (!MapPartitionWithSnapshot(lock.get(), std::move(params), &ignore_path)) {
             return false;
         }
     }
+
+    LOG(INFO) << "Created logical partitions with snapshot.";
+    return true;
+}
+
+bool SnapshotManager::MapPartitionWithSnapshot(LockedFile* lock,
+                                               CreateLogicalPartitionParams params,
+                                               std::string* path) {
+    CHECK(lock);
+    path->clear();
+
+    // Fill out fields in CreateLogicalPartitionParams so that we have more information (e.g. by
+    // reading super partition metadata).
+    CreateLogicalPartitionParams::OwnedData params_owned_data;
+    if (!params.InitDefaults(&params_owned_data)) {
+        return false;
+    }
+
+    if (!params.partition->num_extents) {
+        LOG(INFO) << "Skipping zero-length logical partition: " << params.GetPartitionName();
+        return true;  // leave path empty to indicate that nothing is mapped.
+    }
+
+    // Determine if there is a live snapshot for the SnapshotStatus of the partition; i.e. if the
+    // partition still has a snapshot that needs to be mapped.  If no live snapshot or merge
+    // completed, live_snapshot_status is set to nullopt.
+    std::optional<SnapshotStatus> live_snapshot_status;
+    do {
+        if (!(params.partition->attributes & LP_PARTITION_ATTR_UPDATED)) {
+            LOG(INFO) << "Detected re-flashing of partition, will skip snapshot: "
+                      << params.GetPartitionName();
+            break;
+        }
+        auto file_path = GetSnapshotStatusFilePath(params.GetPartitionName());
+        if (access(file_path.c_str(), F_OK) != 0) {
+            if (errno != ENOENT) {
+                PLOG(INFO) << "Can't map snapshot for " << params.GetPartitionName()
+                           << ": Can't access " << file_path;
+                return false;
+            }
+            break;
+        }
+        live_snapshot_status = std::make_optional<SnapshotStatus>();
+        if (!ReadSnapshotStatus(lock, params.GetPartitionName(), &*live_snapshot_status)) {
+            return false;
+        }
+        // No live snapshot if merge is completed.
+        if (live_snapshot_status->state == SnapshotState::MergeCompleted) {
+            live_snapshot_status.reset();
+        }
+    } while (0);
+
+    if (live_snapshot_status.has_value()) {
+        // dm-snapshot requires the base device to be writable.
+        params.force_writable = true;
+        // Map the base device with a different name to avoid collision.
+        params.device_name = GetBaseDeviceName(params.GetPartitionName());
+    }
+
+    auto& dm = DeviceMapper::Instance();
+    std::string ignore_path;
+    if (!CreateLogicalPartition(params, &ignore_path)) {
+        LOG(ERROR) << "Could not create logical partition " << params.GetPartitionName()
+                   << " as device " << params.GetDeviceName();
+        return false;
+    }
+    if (!live_snapshot_status.has_value()) {
+        return true;
+    }
+
+    // We don't have ueventd in first-stage init, so use device major:minor
+    // strings instead.
+    std::string base_device;
+    if (!dm.GetDeviceString(params.GetDeviceName(), &base_device)) {
+        LOG(ERROR) << "Could not determine major/minor for: " << params.GetDeviceName();
+        return false;
+    }
+    if (!MapSnapshot(lock, params.GetPartitionName(), base_device, {}, path)) {
+        LOG(ERROR) << "Could not map snapshot for partition: " << params.GetPartitionName();
+        return false;
+    }
+
+    LOG(INFO) << "Mapped " << params.GetPartitionName() << " as snapshot device at " << *path;
+
     return true;
 }