8341091: CDS: Segmented roots array misses roots

Reviewed-by: adinn, iklam

src/hotspot/share/cds/archiveHeapLoader.cpp

@@ -376,13 +376,12 @@ void ArchiveHeapLoader::finish_initialization() {
     intptr_t bottom = is_loaded() ? _loaded_heap_bottom : _mapped_heap_bottom;
 
     // The heap roots are stored in one or more segments that are laid out consecutively.
-    // The byte size of each segment (except for the last one) is max_size.
+    // The size of each segment (except for the last one) is max_size_in_{elems,bytes}.
     HeapRootSegments segments = FileMapInfo::current_info()->heap_root_segments();
-    int max_size = segments.max_size_in_bytes();
-    HeapShared::init_root_segment_sizes(max_size);
+    HeapShared::init_root_segment_sizes(segments.max_size_in_elems());
     intptr_t first_segment_addr = bottom + segments.base_offset();
     for (size_t c = 0; c < segments.count(); c++) {
-      oop segment_oop = cast_to_oop(first_segment_addr + (c * max_size));
+      oop segment_oop = cast_to_oop(first_segment_addr + (c * segments.max_size_in_bytes()));
       assert(segment_oop->is_objArray(), "Must be");
       HeapShared::add_root_segment((objArrayOop)segment_oop);
     }

src/hotspot/share/cds/archiveHeapWriter.cpp

@@ -223,6 +223,7 @@ void ArchiveHeapWriter::copy_roots_to_buffer(GrowableArrayCHeap<oop, mtClassShar
                             MIN_GC_REGION_ALIGNMENT,
                             max_elem_count);
 
+  int root_index = 0;
   for (size_t seg_idx = 0; seg_idx < segments.count(); seg_idx++) {
     int size_elems = segments.size_in_elems(seg_idx);
     size_t size_bytes = segments.size_in_bytes(seg_idx);
@@ -235,7 +236,6 @@ void ArchiveHeapWriter::copy_roots_to_buffer(GrowableArrayCHeap<oop, mtClassShar
            "Roots segment " SIZE_FORMAT " start is not aligned: " SIZE_FORMAT,
            segments.count(), oop_offset);
 
-    int root_index = 0;
     objArrayOop seg_oop = allocate_root_segment(oop_offset, size_elems);
     for (int i = 0; i < size_elems; i++) {
       root_segment_at_put(seg_oop, i, roots->at(root_index++));
@@ -245,6 +245,8 @@ void ArchiveHeapWriter::copy_roots_to_buffer(GrowableArrayCHeap<oop, mtClassShar
                         size_elems, size_bytes, p2i(seg_oop));
   }
 
+  assert(root_index == roots->length(), "Post-condition: All roots are handled");
+
   _heap_root_segments = segments;
 }
 

src/hotspot/share/cds/archiveUtils.hpp

@@ -277,7 +277,6 @@ public:
     memset(this, 0, sizeof(*this));
   }
   HeapRootSegments(size_t base_offset, int roots_count, int max_size_in_bytes, int max_size_in_elems) {
-    assert(is_power_of_2(max_size_in_bytes), "must be");
     memset(this, 0, sizeof(*this));
     _base_offset = base_offset;
     _count = (roots_count + max_size_in_elems - 1) / max_size_in_elems;

src/hotspot/share/cds/heapShared.cpp

@@ -136,8 +136,7 @@ static ArchivableStaticFieldInfo fmg_archive_subgraph_entry_fields[] = {
 KlassSubGraphInfo* HeapShared::_default_subgraph_info;
 GrowableArrayCHeap<oop, mtClassShared>* HeapShared::_pending_roots = nullptr;
 GrowableArrayCHeap<OopHandle, mtClassShared>* HeapShared::_root_segments;
-int HeapShared::_root_segment_max_size_shift;
-int HeapShared::_root_segment_max_size_mask;
+int HeapShared::_root_segment_max_size_elems;
 OopHandle HeapShared::_scratch_basic_type_mirrors[T_VOID+1];
 MetaspaceObjToOopHandleTable* HeapShared::_scratch_java_mirror_table = nullptr;
 MetaspaceObjToOopHandleTable* HeapShared::_scratch_references_table = nullptr;
@@ -244,15 +243,29 @@ objArrayOop HeapShared::root_segment(int segment_idx) {
   return segment;
 }
 
+void HeapShared::get_segment_indexes(int idx, int& seg_idx, int& int_idx) {
+  assert(_root_segment_max_size_elems > 0, "sanity");
+
+  // Try to avoid divisions for the common case.
+  if (idx < _root_segment_max_size_elems) {
+    seg_idx = 0;
+    int_idx = idx;
+  } else {
+    seg_idx = idx / _root_segment_max_size_elems;
+    int_idx = idx % _root_segment_max_size_elems;
+  }
+
+  assert(idx == seg_idx * _root_segment_max_size_elems + int_idx,
+         "sanity: %d index maps to %d segment and %d internal", idx, seg_idx, int_idx);
+}
+
 // Returns an objArray that contains all the roots of the archived objects
 oop HeapShared::get_root(int index, bool clear) {
-  assert(_root_segment_max_size_shift > 0, "sanity");
-  assert(_root_segment_max_size_mask  > 0, "sanity");
   assert(index >= 0, "sanity");
   assert(!CDSConfig::is_dumping_heap() && CDSConfig::is_using_archive(), "runtime only");
   assert(!_root_segments->is_empty(), "must have loaded shared heap");
-  int seg_idx = index >> _root_segment_max_size_shift;
-  int int_idx = index &  _root_segment_max_size_mask;
+  int seg_idx, int_idx;
+  get_segment_indexes(index, seg_idx, int_idx);
   oop result = root_segment(seg_idx)->obj_at(int_idx);
   if (clear) {
     clear_root(index);
@@ -264,10 +277,8 @@ void HeapShared::clear_root(int index) {
   assert(index >= 0, "sanity");
   assert(CDSConfig::is_using_archive(), "must be");
   if (ArchiveHeapLoader::is_in_use()) {
-    assert(_root_segment_max_size_shift > 0, "sanity");
-    assert(_root_segment_max_size_mask  > 0, "sanity");
-    int seg_idx = index >> _root_segment_max_size_shift;
-    int int_idx = index &  _root_segment_max_size_mask;
+    int seg_idx, int_idx;
+    get_segment_indexes(index, seg_idx, int_idx);
     if (log_is_enabled(Debug, cds, heap)) {
       oop old = root_segment(seg_idx)->obj_at(int_idx);
       log_debug(cds, heap)("Clearing root %d: was " PTR_FORMAT, index, p2i(old));
@@ -787,10 +798,8 @@ void HeapShared::add_root_segment(objArrayOop segment_oop) {
   _root_segments->push(OopHandle(Universe::vm_global(), segment_oop));
 }
 
-void HeapShared::init_root_segment_sizes(int max_size) {
-  assert(is_power_of_2(max_size), "must be");
-  _root_segment_max_size_shift = log2i_exact(max_size);
-  _root_segment_max_size_mask = max_size - 1;
+void HeapShared::init_root_segment_sizes(int max_size_elems) {
+  _root_segment_max_size_elems = max_size_elems;
 }
 
 void HeapShared::serialize_tables(SerializeClosure* soc) {

src/hotspot/share/cds/heapShared.hpp

@@ -291,8 +291,7 @@ private:
 
   static GrowableArrayCHeap<oop, mtClassShared>* _pending_roots;
   static GrowableArrayCHeap<OopHandle, mtClassShared>* _root_segments;
-  static int _root_segment_max_size_shift;
-  static int _root_segment_max_size_mask;
+  static int _root_segment_max_size_elems;
   static OopHandle _scratch_basic_type_mirrors[T_VOID+1];
   static MetaspaceObjToOopHandleTable* _scratch_java_mirror_table;
   static MetaspaceObjToOopHandleTable* _scratch_references_table;
@@ -407,6 +406,8 @@ private:
   // Run-time only
   static void clear_root(int index);
 
+  static void get_segment_indexes(int index, int& segment_index, int& internal_index);
+
   static void setup_test_class(const char* test_class_name) PRODUCT_RETURN;
 #endif // INCLUDE_CDS_JAVA_HEAP
 
@@ -425,7 +426,7 @@ private:
   static void init_for_dumping(TRAPS) NOT_CDS_JAVA_HEAP_RETURN;
   static void write_subgraph_info_table() NOT_CDS_JAVA_HEAP_RETURN;
   static void add_root_segment(objArrayOop segment_oop) NOT_CDS_JAVA_HEAP_RETURN;
-  static void init_root_segment_sizes(int max_size) NOT_CDS_JAVA_HEAP_RETURN;
+  static void init_root_segment_sizes(int max_size_elems) NOT_CDS_JAVA_HEAP_RETURN;
   static void serialize_tables(SerializeClosure* soc) NOT_CDS_JAVA_HEAP_RETURN;
 
 #ifndef PRODUCT