Merge "fs_mgr: use more unique_fd"
diff --git a/adb/transport.cpp b/adb/transport.cpp
index 76a31ce..8741654 100644
--- a/adb/transport.cpp
+++ b/adb/transport.cpp
@@ -764,6 +764,10 @@
 #if ADB_HOST
         if (t->IsTcpDevice() && !t->kicked()) {
             D("transport: %s unref (attempting reconnection)", t->serial.c_str());
+
+            // We need to clear the transport's keys, so that on the next connection, it tries
+            // again from the beginning.
+            t->ResetKeys();
             reconnect_handler.TrackTransport(t);
         } else {
             D("transport: %s unref (kicking and closing)", t->serial.c_str());
@@ -1328,10 +1332,20 @@
 
 #if ADB_HOST
 std::shared_ptr<RSA> atransport::NextKey() {
-    if (keys_.empty()) keys_ = adb_auth_get_private_keys();
+    if (keys_.empty()) {
+        LOG(INFO) << "fetching keys for transport " << this->serial_name();
+        keys_ = adb_auth_get_private_keys();
+
+        // We should have gotten at least one key: the one that's automatically generated.
+        CHECK(!keys_.empty());
+    }
 
     std::shared_ptr<RSA> result = keys_[0];
     keys_.pop_front();
     return result;
 }
+
+void atransport::ResetKeys() {
+    keys_.clear();
+}
 #endif
diff --git a/adb/transport.h b/adb/transport.h
index 9894bdf..d593700 100644
--- a/adb/transport.h
+++ b/adb/transport.h
@@ -258,6 +258,7 @@
 
 #if ADB_HOST
     std::shared_ptr<RSA> NextKey();
+    void ResetKeys();
 #endif
 
     char token[TOKEN_SIZE] = {};
diff --git a/fs_mgr/fs_mgr_dm_linear.cpp b/fs_mgr/fs_mgr_dm_linear.cpp
index fe0e039..3f075ef 100644
--- a/fs_mgr/fs_mgr_dm_linear.cpp
+++ b/fs_mgr/fs_mgr_dm_linear.cpp
@@ -51,12 +51,20 @@
 using DmTargetZero = android::dm::DmTargetZero;
 using DmTargetLinear = android::dm::DmTargetLinear;
 
-bool GetPhysicalPartitionDevicePath(const LpMetadataBlockDevice& block_device,
-                                    std::string* result) {
+static bool GetPhysicalPartitionDevicePath(const LpMetadata& metadata,
+                                           const LpMetadataBlockDevice& block_device,
+                                           const std::string& super_device,
+                                           std::string* result) {
     // Note: device-mapper will not accept symlinks, so we must use realpath
     // here.
     std::string name = GetBlockDevicePartitionName(block_device);
     std::string path = "/dev/block/by-name/" + name;
+    // If the super device is the source of this block device's metadata,
+    // make sure we use the correct super device (and not just "super",
+    // which might not exist.)
+    if (GetMetadataSuperBlockDevice(metadata) == &block_device) {
+        path = super_device;
+    }
     if (!android::base::Realpath(path, result)) {
         PERROR << "realpath: " << path;
         return false;
@@ -65,7 +73,7 @@
 }
 
 static bool CreateDmTable(const LpMetadata& metadata, const LpMetadataPartition& partition,
-                          DmTable* table) {
+                          const std::string& super_device, DmTable* table) {
     uint64_t sector = 0;
     for (size_t i = 0; i < partition.num_extents; i++) {
         const auto& extent = metadata.extents[partition.first_extent_index + i];
@@ -77,7 +85,7 @@
             case LP_TARGET_TYPE_LINEAR: {
                 const auto& block_device = metadata.block_devices[extent.target_source];
                 std::string path;
-                if (!GetPhysicalPartitionDevicePath(block_device, &path)) {
+                if (!GetPhysicalPartitionDevicePath(metadata, block_device, super_device, &path)) {
                     LOG(ERROR) << "Unable to complete device-mapper table, unknown block device";
                     return false;
                 }
@@ -102,11 +110,11 @@
 
 static bool CreateLogicalPartition(const LpMetadata& metadata, const LpMetadataPartition& partition,
                                    bool force_writable, const std::chrono::milliseconds& timeout_ms,
-                                   std::string* path) {
+                                   const std::string& super_device, std::string* path) {
     DeviceMapper& dm = DeviceMapper::Instance();
 
     DmTable table;
-    if (!CreateDmTable(metadata, partition, &table)) {
+    if (!CreateDmTable(metadata, partition, super_device, &table)) {
         return false;
     }
     if (force_writable) {
@@ -137,7 +145,7 @@
         LOG(ERROR) << "Could not read partition table.";
         return true;
     }
-    return CreateLogicalPartitions(*metadata.get());
+    return CreateLogicalPartitions(*metadata.get(), block_device);
 }
 
 std::unique_ptr<LpMetadata> ReadCurrentMetadata(const std::string& block_device) {
@@ -145,14 +153,14 @@
     return ReadMetadata(block_device.c_str(), slot);
 }
 
-bool CreateLogicalPartitions(const LpMetadata& metadata) {
+bool CreateLogicalPartitions(const LpMetadata& metadata, const std::string& super_device) {
     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, {}, &path)) {
+        if (!CreateLogicalPartition(metadata, partition, false, {}, super_device, &path)) {
             LERROR << "Could not create logical partition: " << GetPartitionName(partition);
             return false;
         }
@@ -171,7 +179,7 @@
     for (const auto& partition : metadata->partitions) {
         if (GetPartitionName(partition) == partition_name) {
             return CreateLogicalPartition(*metadata.get(), partition, force_writable, timeout_ms,
-                                          path);
+                                          block_device, path);
         }
     }
     LERROR << "Could not find any partition with name: " << partition_name;
diff --git a/fs_mgr/include/fs_mgr_dm_linear.h b/fs_mgr/include/fs_mgr_dm_linear.h
index 66abfca..f065071 100644
--- a/fs_mgr/include/fs_mgr_dm_linear.h
+++ b/fs_mgr/include/fs_mgr_dm_linear.h
@@ -43,7 +43,7 @@
 
 // Create block devices for all logical partitions in the given metadata. The
 // metadata must have been read from the current slot.
-bool CreateLogicalPartitions(const LpMetadata& metadata);
+bool CreateLogicalPartitions(const LpMetadata& metadata, const std::string& block_device);
 
 // Create block devices for all logical partitions. This is a convenience
 // method for ReadMetadata and CreateLogicalPartitions.
diff --git a/fs_mgr/liblp/builder.cpp b/fs_mgr/liblp/builder.cpp
index da86d75..699b9e7 100644
--- a/fs_mgr/liblp/builder.cpp
+++ b/fs_mgr/liblp/builder.cpp
@@ -578,6 +578,12 @@
     CHECK_NE(sectors_per_block, 0);
     CHECK(sectors_needed % sectors_per_block == 0);
 
+    if (IsABDevice() && !IsRetrofitDevice() && GetPartitionSlotSuffix(partition->name()) == "_b") {
+        // Allocate "a" partitions top-down and "b" partitions bottom-up, to
+        // minimize fragmentation during OTA.
+        free_regions = PrioritizeSecondHalfOfSuper(free_regions);
+    }
+
     // Find gaps that we can use for new extents. Note we store new extents in a
     // temporary vector, and only commit them if we are guaranteed enough free
     // space.
@@ -621,6 +627,40 @@
     return true;
 }
 
+std::vector<MetadataBuilder::Interval> MetadataBuilder::PrioritizeSecondHalfOfSuper(
+        const std::vector<Interval>& free_list) {
+    const auto& super = block_devices_[0];
+    uint64_t first_sector = super.first_logical_sector;
+    uint64_t last_sector = super.size / LP_SECTOR_SIZE;
+    uint64_t midpoint = first_sector + (last_sector - first_sector) / 2;
+
+    // Choose an aligned sector for the midpoint. This could lead to one half
+    // being slightly larger than the other, but this will not restrict the
+    // size of partitions (it might lead to one extra extent if "B" overflows).
+    midpoint = AlignSector(super, midpoint);
+
+    std::vector<Interval> first_half;
+    std::vector<Interval> second_half;
+    for (const auto& region : free_list) {
+        // Note: deprioritze if not the main super partition. Even though we
+        // don't call this for retrofit devices, we will allow adding additional
+        // block devices on non-retrofit devices.
+        if (region.device_index != 0 || region.end <= midpoint) {
+            first_half.emplace_back(region);
+            continue;
+        }
+        if (region.start < midpoint && region.end > midpoint) {
+            // Split this into two regions.
+            first_half.emplace_back(region.device_index, region.start, midpoint);
+            second_half.emplace_back(region.device_index, midpoint, region.end);
+        } else {
+            second_half.emplace_back(region);
+        }
+    }
+    second_half.insert(second_half.end(), first_half.begin(), first_half.end());
+    return second_half;
+}
+
 void MetadataBuilder::ShrinkPartition(Partition* partition, uint64_t aligned_size) {
     partition->ShrinkTo(aligned_size);
 }
@@ -930,5 +970,9 @@
     return android::base::GetBoolProperty("ro.build.ab_update", false);
 }
 
+bool MetadataBuilder::IsRetrofitDevice() const {
+    return GetBlockDevicePartitionName(block_devices_[0]) != LP_METADATA_DEFAULT_PARTITION_NAME;
+}
+
 }  // namespace fs_mgr
 }  // namespace android
diff --git a/fs_mgr/liblp/builder_test.cpp b/fs_mgr/liblp/builder_test.cpp
index 926fe12..7833a25 100644
--- a/fs_mgr/liblp/builder_test.cpp
+++ b/fs_mgr/liblp/builder_test.cpp
@@ -720,3 +720,41 @@
     ASSERT_EQ(builder->AddPartition("system", 0), nullptr);
     ASSERT_NE(builder->AddPartition("system_a", 0), nullptr);
 }
+
+TEST_F(BuilderTest, ABExtents) {
+    BlockDeviceInfo device_info("super", 10_GiB, 768 * 1024, 0, 4096);
+
+    // A and B slots should be allocated from separate halves of the partition,
+    // to mitigate allocating too many extents. (b/120433288)
+    MetadataBuilder::OverrideABForTesting(true);
+    auto builder = MetadataBuilder::New(device_info, 65536, 2);
+    ASSERT_NE(builder, nullptr);
+    Partition* system_a = builder->AddPartition("system_a", 0);
+    ASSERT_NE(system_a, nullptr);
+    Partition* system_b = builder->AddPartition("system_b", 0);
+    ASSERT_NE(system_b, nullptr);
+    ASSERT_TRUE(builder->ResizePartition(system_a, 2_GiB));
+    ASSERT_TRUE(builder->ResizePartition(system_b, 2_GiB));
+
+    builder->RemovePartition("system_a");
+    system_a = builder->AddPartition("system_a", 0);
+    ASSERT_NE(system_a, nullptr);
+    ASSERT_TRUE(builder->ResizePartition(system_a, 3_GiB));
+
+    EXPECT_EQ(system_a->extents().size(), static_cast<size_t>(1));
+    EXPECT_EQ(system_b->extents().size(), static_cast<size_t>(1));
+    ASSERT_TRUE(builder->ResizePartition(system_b, 6_GiB));
+    EXPECT_EQ(system_b->extents().size(), static_cast<size_t>(3));
+
+    unique_ptr<LpMetadata> exported = builder->Export();
+    ASSERT_NE(exported, nullptr);
+    ASSERT_EQ(exported->extents.size(), static_cast<size_t>(4));
+    EXPECT_EQ(exported->extents[0].target_data, 10487808);
+    EXPECT_EQ(exported->extents[0].num_sectors, 4194304);
+    EXPECT_EQ(exported->extents[1].target_data, 14682624);
+    EXPECT_EQ(exported->extents[1].num_sectors, 6288896);
+    EXPECT_EQ(exported->extents[2].target_data, 6292992);
+    EXPECT_EQ(exported->extents[2].num_sectors, 2099712);
+    EXPECT_EQ(exported->extents[3].target_data, 1536);
+    EXPECT_EQ(exported->extents[3].num_sectors, 6291456);
+}
diff --git a/fs_mgr/liblp/include/liblp/builder.h b/fs_mgr/liblp/include/liblp/builder.h
index 4bb38d6..f477b4b 100644
--- a/fs_mgr/liblp/include/liblp/builder.h
+++ b/fs_mgr/liblp/include/liblp/builder.h
@@ -275,6 +275,7 @@
                        const LpMetadataPartition& source);
     bool ImportPartition(const LpMetadata& metadata, const LpMetadataPartition& source);
     bool IsABDevice() const;
+    bool IsRetrofitDevice() const;
 
     struct Interval {
         uint32_t device_index;
@@ -294,6 +295,7 @@
     std::vector<Interval> GetFreeRegions() const;
     void ExtentsToFreeList(const std::vector<Interval>& extents,
                            std::vector<Interval>* free_regions) const;
+    std::vector<Interval> PrioritizeSecondHalfOfSuper(const std::vector<Interval>& free_list);
 
     static bool sABOverrideValue;
     static bool sABOverrideSet;
diff --git a/init/first_stage_mount.cpp b/init/first_stage_mount.cpp
index d35329e..2f7f79d 100644
--- a/init/first_stage_mount.cpp
+++ b/init/first_stage_mount.cpp
@@ -252,9 +252,10 @@
 }
 
 bool FirstStageMount::InitDmLinearBackingDevices(const android::fs_mgr::LpMetadata& metadata) {
-    auto partition_names = GetBlockDevicePartitionNames(metadata);
+    auto partition_names = android::fs_mgr::GetBlockDevicePartitionNames(metadata);
     for (const auto& partition_name : partition_names) {
-        if (partition_name == lp_metadata_partition_) {
+        const auto super_device = android::fs_mgr::GetMetadataSuperBlockDevice(metadata);
+        if (partition_name == android::fs_mgr::GetBlockDevicePartitionName(*super_device)) {
             continue;
         }
         required_devices_partition_names_.emplace(partition_name);
@@ -292,7 +293,7 @@
     if (!InitDmLinearBackingDevices(*metadata.get())) {
         return false;
     }
-    return android::fs_mgr::CreateLogicalPartitions(*metadata.get());
+    return android::fs_mgr::CreateLogicalPartitions(*metadata.get(), lp_metadata_partition_);
 }
 
 ListenerAction FirstStageMount::HandleBlockDevice(const std::string& name, const Uevent& uevent) {