Pre-allocate necessary heap space with boot image.

Previously we relied on a subsequent mmap to succeed
without prior memory reservation.

Test: m test-art-host-gtest
Test: testrunner.py --host --optimizing
Test: Pixel 2 XL boots.
Test: m test-art-target-gtest
Test: testrunner.py --target --optimizing
Test: Repeat the above with ART_USE_READ_BARRIER=false.
Bug: 77856493
Change-Id: I1e433b1b6fecd0f908c7ea079f74361e27a78f0a
diff --git a/runtime/gc/heap.cc b/runtime/gc/heap.cc
index b74071b..589e9a4 100644
--- a/runtime/gc/heap.cc
+++ b/runtime/gc/heap.cc
@@ -312,39 +312,7 @@
   ChangeCollector(desired_collector_type_);
   live_bitmap_.reset(new accounting::HeapBitmap(this));
   mark_bitmap_.reset(new accounting::HeapBitmap(this));
-  // Requested begin for the alloc space, to follow the mapped image and oat files
-  uint8_t* requested_alloc_space_begin = nullptr;
-  if (foreground_collector_type_ == kCollectorTypeCC) {
-    // Need to use a low address so that we can allocate a contiguous 2 * Xmx space when there's no
-    // image (dex2oat for target).
-    requested_alloc_space_begin = kPreferredAllocSpaceBegin;
-  }
 
-  // Load image space(s).
-  std::vector<std::unique_ptr<space::ImageSpace>> boot_image_spaces;
-  if (space::ImageSpace::LoadBootImage(image_file_name,
-                                       image_instruction_set,
-                                       &boot_image_spaces,
-                                       &requested_alloc_space_begin)) {
-    for (std::unique_ptr<space::ImageSpace>& space : boot_image_spaces) {
-      boot_image_spaces_.push_back(space.get());
-      AddSpace(space.release());
-    }
-  }
-
-  /*
-  requested_alloc_space_begin ->     +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-
-                                     +-  nonmoving space (non_moving_space_capacity)+-
-                                     +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-
-                                     +-????????????????????????????????????????????+-
-                                     +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-
-                                     +-main alloc space / bump space 1 (capacity_) +-
-                                     +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-
-                                     +-????????????????????????????????????????????+-
-                                     +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-
-                                     +-main alloc space2 / bump space 2 (capacity_)+-
-                                     +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-
-  */
   // We don't have hspace compaction enabled with GSS or CC.
   if (foreground_collector_type_ == kCollectorTypeGSS ||
       foreground_collector_type_ == kCollectorTypeCC) {
@@ -363,21 +331,63 @@
   if (foreground_collector_type_ == kCollectorTypeGSS) {
     separate_non_moving_space = false;
   }
+
+  // Requested begin for the alloc space, to follow the mapped image and oat files
+  uint8_t* request_begin = nullptr;
+  // Calculate the extra space required after the boot image, see allocations below.
+  size_t heap_reservation_size = separate_non_moving_space
+      ? non_moving_space_capacity
+      : ((is_zygote && foreground_collector_type_ != kCollectorTypeCC) ? capacity_ : 0u);
+  heap_reservation_size = RoundUp(heap_reservation_size, kPageSize);
+  // Load image space(s).
+  std::vector<std::unique_ptr<space::ImageSpace>> boot_image_spaces;
+  MemMap heap_reservation;
+  if (space::ImageSpace::LoadBootImage(image_file_name,
+                                       image_instruction_set,
+                                       heap_reservation_size,
+                                       &boot_image_spaces,
+                                       &heap_reservation)) {
+    DCHECK_EQ(heap_reservation_size, heap_reservation.IsValid() ? heap_reservation.Size() : 0u);
+    DCHECK(!boot_image_spaces.empty());
+    request_begin = boot_image_spaces.back()->GetImageHeader().GetOatFileEnd();
+    DCHECK(!heap_reservation.IsValid() || request_begin == heap_reservation.Begin())
+        << "request_begin=" << static_cast<const void*>(request_begin)
+        << " heap_reservation.Begin()=" << static_cast<const void*>(heap_reservation.Begin());
+    for (std::unique_ptr<space::ImageSpace>& space : boot_image_spaces) {
+      boot_image_spaces_.push_back(space.get());
+      AddSpace(space.release());
+    }
+  } else {
+    if (foreground_collector_type_ == kCollectorTypeCC) {
+      // Need to use a low address so that we can allocate a contiguous 2 * Xmx space
+      // when there's no image (dex2oat for target).
+      request_begin = kPreferredAllocSpaceBegin;
+    }
+    // Gross hack to make dex2oat deterministic.
+    if (foreground_collector_type_ == kCollectorTypeMS && Runtime::Current()->IsAotCompiler()) {
+      // Currently only enabled for MS collector since that is what the deterministic dex2oat uses.
+      // b/26849108
+      request_begin = reinterpret_cast<uint8_t*>(kAllocSpaceBeginForDeterministicAoT);
+    }
+  }
+
+  /*
+  requested_alloc_space_begin ->     +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-
+                                     +-  nonmoving space (non_moving_space_capacity)+-
+                                     +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-
+                                     +-????????????????????????????????????????????+-
+                                     +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-
+                                     +-main alloc space / bump space 1 (capacity_) +-
+                                     +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-
+                                     +-????????????????????????????????????????????+-
+                                     +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-
+                                     +-main alloc space2 / bump space 2 (capacity_)+-
+                                     +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-
+  */
+
   MemMap main_mem_map_1;
   MemMap main_mem_map_2;
 
-  // Gross hack to make dex2oat deterministic.
-  if (foreground_collector_type_ == kCollectorTypeMS &&
-      requested_alloc_space_begin == nullptr &&
-      Runtime::Current()->IsAotCompiler()) {
-    // Currently only enabled for MS collector since that is what the deterministic dex2oat uses.
-    // b/26849108
-    requested_alloc_space_begin = reinterpret_cast<uint8_t*>(kAllocSpaceBeginForDeterministicAoT);
-  }
-  uint8_t* request_begin = requested_alloc_space_begin;
-  if (request_begin != nullptr && separate_non_moving_space) {
-    request_begin += non_moving_space_capacity;
-  }
   std::string error_str;
   MemMap non_moving_space_mem_map;
   if (separate_non_moving_space) {
@@ -388,9 +398,16 @@
     const char* space_name = is_zygote ? kZygoteSpaceName : kNonMovingSpaceName;
     // Reserve the non moving mem map before the other two since it needs to be at a specific
     // address.
-    non_moving_space_mem_map = MapAnonymousPreferredAddress(
-        space_name, requested_alloc_space_begin, non_moving_space_capacity, &error_str);
+    DCHECK_EQ(heap_reservation.IsValid(), !boot_image_spaces_.empty());
+    if (heap_reservation.IsValid()) {
+      non_moving_space_mem_map = heap_reservation.RemapAtEnd(
+          heap_reservation.Begin(), space_name, PROT_READ | PROT_WRITE, &error_str);
+    } else {
+      non_moving_space_mem_map = MapAnonymousPreferredAddress(
+          space_name, request_begin, non_moving_space_capacity, &error_str);
+    }
     CHECK(non_moving_space_mem_map.IsValid()) << error_str;
+    DCHECK(!heap_reservation.IsValid());
     // Try to reserve virtual memory at a lower address if we have a separate non moving space.
     request_begin = kPreferredAllocSpaceBegin + non_moving_space_capacity;
   }
@@ -404,14 +421,19 @@
       // If no separate non-moving space and we are the zygote, the main space must come right
       // after the image space to avoid a gap. This is required since we want the zygote space to
       // be adjacent to the image space.
-      main_mem_map_1 = MemMap::MapAnonymous(kMemMapSpaceName[0],
-                                            request_begin,
-                                            capacity_,
-                                            PROT_READ | PROT_WRITE,
-                                            /* low_4gb */ true,
-                                            &error_str);
+      DCHECK_EQ(heap_reservation.IsValid(), !boot_image_spaces_.empty());
+      main_mem_map_1 = MemMap::MapAnonymous(
+          kMemMapSpaceName[0],
+          request_begin,
+          capacity_,
+          PROT_READ | PROT_WRITE,
+          /* low_4gb */ true,
+          /* reuse */ false,
+          heap_reservation.IsValid() ? &heap_reservation : nullptr,
+          &error_str);
     }
     CHECK(main_mem_map_1.IsValid()) << error_str;
+    DCHECK(!heap_reservation.IsValid());
   }
   if (support_homogeneous_space_compaction ||
       background_collector_type_ == kCollectorTypeSS ||
@@ -437,7 +459,7 @@
                                                                /* can_move_objects */ false);
     non_moving_space_->SetFootprintLimit(non_moving_space_->Capacity());
     CHECK(non_moving_space_ != nullptr) << "Failed creating non moving space "
-        << requested_alloc_space_begin;
+        << non_moving_space_mem_map.Begin();
     AddSpace(non_moving_space_);
   }
   // Create other spaces based on whether or not we have a moving GC.
diff --git a/runtime/gc/space/image_space.cc b/runtime/gc/space/image_space.cc
index 0af049a..f308f63 100644
--- a/runtime/gc/space/image_space.cc
+++ b/runtime/gc/space/image_space.cc
@@ -1461,9 +1461,10 @@
     return cache_filename_;
   }
 
-  bool LoadFromSystem(/*out*/ std::vector<std::unique_ptr<space::ImageSpace>>* boot_image_spaces,
-                      /*out*/ uint8_t** oat_file_end,
-                      /*out*/ std::string* error_msg) REQUIRES_SHARED(Locks::mutator_lock_) {
+  bool LoadFromSystem(size_t extra_reservation_size,
+                      /*out*/std::vector<std::unique_ptr<space::ImageSpace>>* boot_image_spaces,
+                      /*out*/MemMap* extra_reservation,
+                      /*out*/std::string* error_msg) REQUIRES_SHARED(Locks::mutator_lock_) {
     std::string filename = GetSystemImageFilename(image_location_.c_str(), image_isa_);
     std::vector<std::string> locations;
     if (!GetBootClassPathImageLocations(image_location_, filename, &locations, error_msg)) {
@@ -1484,11 +1485,14 @@
     }
     MemMap image_reservation;
     MemMap oat_reservation;
+    MemMap local_extra_reservation;
     if (!ReserveBootImageMemory(image_start,
                                 image_end,
                                 oat_end,
+                                extra_reservation_size,
                                 &image_reservation,
                                 &oat_reservation,
+                                &local_extra_reservation,
                                 error_msg)) {
       return false;
     }
@@ -1511,7 +1515,7 @@
       return false;
     }
 
-    *oat_file_end = GetOatFileEnd(spaces);
+    *extra_reservation = std::move(local_extra_reservation);
     boot_image_spaces->swap(spaces);
     return true;
   }
@@ -1519,9 +1523,10 @@
   bool LoadFromDalvikCache(
       bool validate_system_checksums,
       bool validate_oat_file,
-      /*out*/ std::vector<std::unique_ptr<space::ImageSpace>>* boot_image_spaces,
-      /*out*/ uint8_t** oat_file_end,
-      /*out*/ std::string* error_msg) REQUIRES_SHARED(Locks::mutator_lock_) {
+      size_t extra_reservation_size,
+      /*out*/std::vector<std::unique_ptr<space::ImageSpace>>* boot_image_spaces,
+      /*out*/MemMap* extra_reservation,
+      /*out*/std::string* error_msg) REQUIRES_SHARED(Locks::mutator_lock_) {
     DCHECK(DalvikCacheExists());
     std::vector<std::string> locations;
     if (!GetBootClassPathImageLocations(image_location_, cache_filename_, &locations, error_msg)) {
@@ -1548,11 +1553,14 @@
     }
     MemMap image_reservation;
     MemMap oat_reservation;
+    MemMap local_extra_reservation;
     if (!ReserveBootImageMemory(image_start,
                                 image_end,
                                 oat_end,
+                                extra_reservation_size,
                                 &image_reservation,
                                 &oat_reservation,
+                                &local_extra_reservation,
                                 error_msg)) {
       return false;
     }
@@ -1594,7 +1602,7 @@
       return false;
     }
 
-    *oat_file_end = GetOatFileEnd(spaces);
+    *extra_reservation = std::move(local_extra_reservation);
     boot_image_spaces->swap(spaces);
     return true;
   }
@@ -1692,14 +1700,18 @@
   bool ReserveBootImageMemory(uint32_t image_start,
                               uint32_t image_end,
                               uint32_t oat_end,
+                              size_t extra_reservation_size,
                               /*out*/MemMap* image_reservation,
                               /*out*/MemMap* oat_reservation,
+                              /*out*/MemMap* extra_reservation,
                               /*out*/std::string* error_msg) {
     DCHECK(!image_reservation->IsValid());
+    size_t total_size =
+        dchecked_integral_cast<size_t>(oat_end - image_start) + extra_reservation_size;
     *image_reservation =
         MemMap::MapAnonymous("Boot image reservation",
                              reinterpret_cast32<uint8_t*>(image_start),
-                             oat_end - image_start,
+                             total_size,
                              PROT_NONE,
                              /* low_4gb */ true,
                              /* reuse */ false,
@@ -1708,6 +1720,19 @@
     if (!image_reservation->IsValid()) {
       return false;
     }
+    DCHECK(!extra_reservation->IsValid());
+    if (extra_reservation_size != 0u) {
+      DCHECK_ALIGNED(extra_reservation_size, kPageSize);
+      DCHECK_LT(extra_reservation_size, image_reservation->Size());
+      uint8_t* split = image_reservation->End() - extra_reservation_size;
+      *extra_reservation = image_reservation->RemapAtEnd(split,
+                                                         "Boot image extra reservation",
+                                                         PROT_NONE,
+                                                         error_msg);
+      if (!extra_reservation->IsValid()) {
+        return false;
+      }
+    }
     DCHECK(!oat_reservation->IsValid());
     *oat_reservation = image_reservation->RemapAtEnd(reinterpret_cast32<uint8_t*>(image_end),
                                                      "Boot image oat reservation",
@@ -1738,16 +1763,6 @@
     return true;
   }
 
-  uint8_t* GetOatFileEnd(const std::vector<std::unique_ptr<ImageSpace>>& spaces) {
-    DCHECK(std::is_sorted(
-        spaces.begin(),
-        spaces.end(),
-        [](const std::unique_ptr<ImageSpace>& lhs, const std::unique_ptr<ImageSpace>& rhs) {
-          return lhs->GetOatFileEnd() < rhs->GetOatFileEnd();
-        }));
-    return AlignUp(spaces.back()->GetOatFileEnd(), kPageSize);
-  }
-
   const std::string& image_location_;
   InstructionSet image_isa_;
   bool is_zygote_;
@@ -1797,13 +1812,15 @@
 bool ImageSpace::LoadBootImage(
     const std::string& image_location,
     const InstructionSet image_isa,
-    /*out*/ std::vector<std::unique_ptr<space::ImageSpace>>* boot_image_spaces,
-    /*out*/ uint8_t** oat_file_end) {
+    size_t extra_reservation_size,
+    /*out*/std::vector<std::unique_ptr<space::ImageSpace>>* boot_image_spaces,
+    /*out*/MemMap* extra_reservation) {
   ScopedTrace trace(__FUNCTION__);
 
   DCHECK(boot_image_spaces != nullptr);
   DCHECK(boot_image_spaces->empty());
-  DCHECK(oat_file_end != nullptr);
+  DCHECK_ALIGNED(extra_reservation_size, kPageSize);
+  DCHECK(extra_reservation != nullptr);
   DCHECK_NE(image_isa, InstructionSet::kNone);
 
   if (image_location.empty()) {
@@ -1860,8 +1877,9 @@
     // If we have system image, validate system image checksums, otherwise validate the oat file.
     if (loader.LoadFromDalvikCache(/* validate_system_checksums */ loader.HasSystem(),
                                    /* validate_oat_file */ !loader.HasSystem(),
+                                   extra_reservation_size,
                                    boot_image_spaces,
-                                   oat_file_end,
+                                   extra_reservation,
                                    &local_error_msg)) {
       return true;
     }
@@ -1875,7 +1893,10 @@
 
   if (loader.HasSystem() && !relocate) {
     std::string local_error_msg;
-    if (loader.LoadFromSystem(boot_image_spaces, oat_file_end, &local_error_msg)) {
+    if (loader.LoadFromSystem(extra_reservation_size,
+                              boot_image_spaces,
+                              extra_reservation,
+                              &local_error_msg)) {
       return true;
     }
     error_msgs.push_back(local_error_msg);
@@ -1892,8 +1913,9 @@
       if (patch_success) {
         if (loader.LoadFromDalvikCache(/* validate_system_checksums */ false,
                                        /* validate_oat_file */ false,
+                                       extra_reservation_size,
                                        boot_image_spaces,
-                                       oat_file_end,
+                                       extra_reservation,
                                        &local_error_msg)) {
           return true;
         }
@@ -1917,8 +1939,9 @@
       if (compilation_success) {
         if (loader.LoadFromDalvikCache(/* validate_system_checksums */ false,
                                        /* validate_oat_file */ false,
+                                       extra_reservation_size,
                                        boot_image_spaces,
-                                       oat_file_end,
+                                       extra_reservation,
                                        &local_error_msg)) {
           return true;
         }
diff --git a/runtime/gc/space/image_space.h b/runtime/gc/space/image_space.h
index bd686be..a2490ac 100644
--- a/runtime/gc/space/image_space.h
+++ b/runtime/gc/space/image_space.h
@@ -44,8 +44,9 @@
   static bool LoadBootImage(
       const std::string& image_location,
       const InstructionSet image_isa,
-      /*out*/ std::vector<std::unique_ptr<space::ImageSpace>>* boot_image_spaces,
-      /*out*/ uint8_t** oat_file_end) REQUIRES_SHARED(Locks::mutator_lock_);
+      size_t extra_reservation_size,
+      /*out*/std::vector<std::unique_ptr<space::ImageSpace>>* boot_image_spaces,
+      /*out*/MemMap* extra_reservation) REQUIRES_SHARED(Locks::mutator_lock_);
 
   // Try to open an existing app image space.
   static std::unique_ptr<ImageSpace> CreateFromAppImage(const char* image,