8359683: ZGC: NUMA-Aware Relocation

Reviewed-by: aboldtch, sjohanss

Author: jsikstro

src/hotspot/share/gc/z/zArray.hpp

@@ -91,6 +91,10 @@ class ZArrayIteratorImpl : public StackObj {
   ZArrayIteratorImpl(const ZArray<T>* array);
 
   bool next(T* elem);
+
+  template <typename Function, typename... Args>
+  bool next_if(T* elem, Function predicate, Args&&... args);
+
   bool next_index(size_t* index);
 
   T index_to_elem(size_t index);

src/hotspot/share/gc/z/zArray.inline.hpp

@@ -161,6 +161,20 @@ inline bool ZArrayIteratorImpl<T, Parallel>::next(T* elem) {
   return false;
 }
 
+template <typename T, bool Parallel>
+template <typename Function, typename... Args>
+inline bool ZArrayIteratorImpl<T, Parallel>::next_if(T* elem, Function predicate, Args&&... args) {
+  size_t index;
+  while (next_index(&index)) {
+    if (predicate(index_to_elem(index), args...)) {
+      *elem = index_to_elem(index);
+      return true;
+    }
+  }
+
+  return false;
+}
+
 template <typename T, bool Parallel>
 inline bool ZArrayIteratorImpl<T, Parallel>::next_index(size_t* index) {
   if (Parallel) {

src/hotspot/share/gc/z/zForwarding.hpp

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2015, 2024, Oracle and/or its affiliates. All rights reserved.
+ * Copyright (c) 2015, 2025, Oracle and/or its affiliates. All rights reserved.
  * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
  *
  * This code is free software; you can redistribute it and/or modify it
@@ -59,6 +59,7 @@ class ZForwarding {
   const size_t           _object_alignment_shift;
   const AttachedArray    _entries;
   ZPage* const           _page;
+  const uint32_t         _partition_id;
   const ZPageAge         _from_age;
   const ZPageAge         _to_age;
   volatile bool          _claimed;
@@ -108,6 +109,8 @@ class ZForwarding {
   size_t size() const;
   size_t object_alignment_shift() const;
 
+  uint32_t partition_id() const;
+
   bool is_promotion() const;
 
   // Visit from-objects

src/hotspot/share/gc/z/zForwarding.inline.hpp

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2015, 2024, Oracle and/or its affiliates. All rights reserved.
+ * Copyright (c) 2015, 2025, Oracle and/or its affiliates. All rights reserved.
  * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
  *
  * This code is free software; you can redistribute it and/or modify it
@@ -61,6 +61,7 @@ inline ZForwarding::ZForwarding(ZPage* page, ZPageAge to_age, size_t nentries)
     _object_alignment_shift(page->object_alignment_shift()),
     _entries(nentries),
     _page(page),
+    _partition_id(page->single_partition_id()),
     _from_age(page->age()),
     _to_age(to_age),
     _claimed(false),
@@ -102,6 +103,10 @@ inline size_t ZForwarding::object_alignment_shift() const {
   return _object_alignment_shift;
 }
 
+inline uint32_t ZForwarding::partition_id() const {
+  return _partition_id;
+}
+
 inline bool ZForwarding::is_promotion() const {
   return _from_age != ZPageAge::old &&
          _to_age == ZPageAge::old;

src/hotspot/share/gc/z/zHeap.cpp

@@ -250,8 +250,8 @@ void ZHeap::account_undo_alloc_page(ZPage* page) {
                 p2i(Thread::current()), ZUtils::thread_name(), p2i(page), page->size());
 }
 
-ZPage* ZHeap::alloc_page(ZPageType type, size_t size, ZAllocationFlags flags, ZPageAge age) {
-  ZPage* const page = _page_allocator.alloc_page(type, size, flags, age);
+ZPage* ZHeap::alloc_page(ZPageType type, size_t size, ZAllocationFlags flags, ZPageAge age, uint32_t preferred_partition) {
+  ZPage* const page = _page_allocator.alloc_page(type, size, flags, age, preferred_partition);
   if (page != nullptr) {
     // Insert page table entry
     _page_table.insert(page);

src/hotspot/share/gc/z/zHeap.hpp

@@ -108,7 +108,7 @@ class ZHeap {
   void mark_flush(Thread* thread);
 
   // Page allocation
-  ZPage* alloc_page(ZPageType type, size_t size, ZAllocationFlags flags, ZPageAge age);
+  ZPage* alloc_page(ZPageType type, size_t size, ZAllocationFlags flags, ZPageAge age, uint32_t preferred_partition);
   void undo_alloc_page(ZPage* page);
   void free_page(ZPage* page);
   size_t free_empty_pages(ZGenerationId id, const ZArray<ZPage*>* pages);

src/hotspot/share/gc/z/zHeuristics.cpp

@@ -26,6 +26,7 @@
 #include "gc/z/zCPU.inline.hpp"
 #include "gc/z/zGlobals.hpp"
 #include "gc/z/zHeuristics.hpp"
+#include "gc/z/zNUMA.inline.hpp"
 #include "runtime/globals.hpp"
 #include "runtime/os.hpp"
 #include "utilities/align.hpp"
@@ -58,9 +59,11 @@ void ZHeuristics::set_medium_page_size() {
 }
 
 size_t ZHeuristics::relocation_headroom() {
-  // Calculate headroom needed to avoid in-place relocation. Each worker will try
-  // to allocate a small page, and all workers will share a single medium page.
-  return (ConcGCThreads * ZPageSizeSmall) + ZPageSizeMediumMax;
+  // Calculate headroom needed to avoid in-place relocation. For each NUMA node,
+  // each worker will try to allocate a small page, and all workers will share a
+  // single medium page.
+  const size_t per_numa_headroom = (ConcGCThreads * ZPageSizeSmall) + ZPageSizeMediumMax;
+  return per_numa_headroom * ZNUMA::count();
 }
 
 bool ZHeuristics::use_per_cpu_shared_small_pages() {

src/hotspot/share/gc/z/zObjectAllocator.cpp

@@ -53,7 +53,7 @@ ZPage* const* ZObjectAllocator::PerAge::shared_small_page_addr() const {
 }
 
 ZPage* ZObjectAllocator::PerAge::alloc_page(ZPageType type, size_t size, ZAllocationFlags flags) {
-  return ZHeap::heap()->alloc_page(type, size, flags, _age);
+  return ZHeap::heap()->alloc_page(type, size, flags, _age, ZNUMA::id());
 }
 
 void ZObjectAllocator::PerAge::undo_alloc_page(ZPage* page) {

src/hotspot/share/gc/z/zPage.inline.hpp

@@ -157,6 +157,7 @@ inline const ZVirtualMemory& ZPage::virtual_memory() const {
 }
 
 inline uint32_t ZPage::single_partition_id() const {
+  assert(!is_multi_partition(), "Don't fetch single partition id if page is multi-partition");
   return _single_partition_id;
 }
 

src/hotspot/share/gc/z/zPageAllocator.cpp

@@ -412,28 +412,30 @@ class ZPageAllocation : public StackObj {
   const Ticks                _start_timestamp;
   const uint32_t             _young_seqnum;
   const uint32_t             _old_seqnum;
-  const uint32_t             _initiating_numa_id;
+  const uint32_t             _preferred_partition;
   bool                       _is_multi_partition;
   ZSinglePartitionAllocation _single_partition_allocation;
   ZMultiPartitionAllocation  _multi_partition_allocation;
   ZListNode<ZPageAllocation> _node;
   ZFuture<bool>              _stall_result;
 
 public:
-  ZPageAllocation(ZPageType type, size_t size, ZAllocationFlags flags, ZPageAge age)
+  ZPageAllocation(ZPageType type, size_t size, ZAllocationFlags flags, ZPageAge age, uint32_t preferred_partition)
     : _type(type),
       _requested_size(size),
       _flags(flags),
       _age(age),
       _start_timestamp(Ticks::now()),
       _young_seqnum(ZGeneration::young()->seqnum()),
       _old_seqnum(ZGeneration::old()->seqnum()),
-      _initiating_numa_id(ZNUMA::id()),
+      _preferred_partition(preferred_partition),
       _is_multi_partition(false),
       _single_partition_allocation(size),
       _multi_partition_allocation(size),
       _node(),
-      _stall_result() {}
+      _stall_result() {
+    assert(_preferred_partition < ZNUMA::count(), "Preferred partition out-of-bounds (0 <= %d < %d)", _preferred_partition, ZNUMA::count());
+  }
 
   void reset_for_retry() {
     _is_multi_partition = false;
@@ -474,8 +476,8 @@ class ZPageAllocation : public StackObj {
     return _old_seqnum;
   }
 
-  uint32_t initiating_numa_id() const {
-    return _initiating_numa_id;
+  uint32_t preferred_partition() const {
+    return _preferred_partition;
   }
 
   bool is_multi_partition() const {
@@ -1397,10 +1399,10 @@ static void check_out_of_memory_during_initialization() {
   }
 }
 
-ZPage* ZPageAllocator::alloc_page(ZPageType type, size_t size, ZAllocationFlags flags, ZPageAge age) {
+ZPage* ZPageAllocator::alloc_page(ZPageType type, size_t size, ZAllocationFlags flags, ZPageAge age, uint32_t preferred_partition) {
   EventZPageAllocation event;
 
-  ZPageAllocation allocation(type, size, flags, age);
+  ZPageAllocation allocation(type, size, flags, age, preferred_partition);
 
   // Allocate the page
   ZPage* const page = alloc_page_inner(&allocation);
@@ -1548,7 +1550,7 @@ bool ZPageAllocator::claim_capacity(ZPageAllocation* allocation) {
   }
 
   // Round robin single-partition claiming
-  const uint32_t start_numa_id = allocation->initiating_numa_id();
+  const uint32_t start_numa_id = allocation->preferred_partition();
   const uint32_t start_partition = start_numa_id;
   const uint32_t num_partitions = _partitions.count();
 
@@ -1560,7 +1562,7 @@ bool ZPageAllocator::claim_capacity(ZPageAllocation* allocation) {
     }
   }
 
-  if (!is_multi_partition_enabled() || sum_available() < allocation->size()) {
+  if (!is_multi_partition_allowed(allocation)) {
     // Multi-partition claiming is not possible
     return false;
   }
@@ -1578,7 +1580,7 @@ bool ZPageAllocator::claim_capacity(ZPageAllocation* allocation) {
 }
 
 bool ZPageAllocator::claim_capacity_fast_medium(ZPageAllocation* allocation) {
-  const uint32_t start_node = allocation->initiating_numa_id();
+  const uint32_t start_node = allocation->preferred_partition();
   const uint32_t numa_nodes = ZNUMA::count();
 
   for (uint32_t i = 0; i < numa_nodes; ++i) {
@@ -2191,6 +2193,12 @@ bool ZPageAllocator::is_multi_partition_enabled() const {
   return _virtual.is_multi_partition_enabled();
 }
 
+bool ZPageAllocator::is_multi_partition_allowed(const ZPageAllocation* allocation) const {
+  return is_multi_partition_enabled() &&
+         allocation->type() == ZPageType::large &&
+         allocation->size() <= sum_available();
+}
+
 const ZPartition& ZPageAllocator::partition_from_partition_id(uint32_t numa_id) const {
   return _partitions.get(numa_id);
 }