Move heap lock acquisition and old-gen check into allocator

src/hotspot/share/gc/shenandoah/shenandoahHeap.cpp

@@ -944,7 +944,7 @@ HeapWord* ShenandoahHeap::allocate_memory(ShenandoahAllocRequest& req) {
   if (req.is_mutator_alloc()) {
 
     if (!ShenandoahAllocFailureALot || !should_inject_alloc_failure()) {
-      result = allocate_memory_under_lock(req, in_new_region);
+      result = allocate_memory_work(req, in_new_region);
     }
 
     // Check that gc overhead is not exceeded.
@@ -976,7 +976,7 @@ HeapWord* ShenandoahHeap::allocate_memory(ShenandoahAllocRequest& req) {
       const size_t original_count = shenandoah_policy()->full_gc_count();
       while (result == nullptr && should_retry_allocation(original_count)) {
         control_thread()->handle_alloc_failure(req, true);
-        result = allocate_memory_under_lock(req, in_new_region);
+        result = allocate_memory_work(req, in_new_region);
       }
       if (result != nullptr) {
         // If our allocation request has been satisfied after it initially failed, we count this as good gc progress
@@ -992,7 +992,7 @@ HeapWord* ShenandoahHeap::allocate_memory(ShenandoahAllocRequest& req) {
     }
   } else {
     assert(req.is_gc_alloc(), "Can only accept GC allocs here");
-    result = allocate_memory_under_lock(req, in_new_region);
+    result = allocate_memory_work(req, in_new_region);
     // Do not call handle_alloc_failure() here, because we cannot block.
     // The allocation failure would be handled by the LRB slowpath with handle_alloc_failure_evac().
   }
@@ -1022,20 +1022,7 @@ inline bool ShenandoahHeap::should_retry_allocation(size_t original_full_gc_coun
       && !shenandoah_policy()->is_at_shutdown();
 }
 
-HeapWord* ShenandoahHeap::allocate_memory_under_lock(ShenandoahAllocRequest& req, bool& in_new_region) {
-  // If we are dealing with mutator allocation, then we may need to block for safepoint.
-  // We cannot block for safepoint for GC allocations, because there is a high chance
-  // we are already running at safepoint or from stack watermark machinery, and we cannot
-  // block again.
-  ShenandoahHeapLocker locker(lock(), req.is_mutator_alloc());
-
-  // Make sure the old generation has room for either evacuations or promotions before trying to allocate.
-  if (req.is_old() && !old_generation()->can_allocate(req)) {
-    return nullptr;
-  }
-
-  // If TLAB request size is greater than available, allocate() will attempt to downsize request to fit within available
-  // memory.
+HeapWord* ShenandoahHeap::allocate_memory_work(ShenandoahAllocRequest& req, bool& in_new_region) {
   HeapWord* result = _allocator->allocate(req, in_new_region);
 
   if (result != nullptr) {

src/hotspot/share/gc/shenandoah/shenandoahHeap.hpp

@@ -705,7 +705,7 @@ protected:
   inline HeapWord* allocate_from_gclab(Thread* thread, size_t size);
 
 private:
-  HeapWord* allocate_memory_under_lock(ShenandoahAllocRequest& request, bool& in_new_region);
+  HeapWord* allocate_memory_work(ShenandoahAllocRequest& request, bool& in_new_region);
   HeapWord* allocate_from_gclab_slow(Thread* thread, size_t size);
   HeapWord* allocate_new_gclab(size_t min_size, size_t word_size, size_t* actual_size);
 

src/hotspot/share/gc/shenandoah/shenandoahPartitionAllocator.cpp

@@ -28,6 +28,7 @@
 #include "gc/shenandoah/shenandoahHeap.inline.hpp"
 #include "gc/shenandoah/shenandoahHeapRegion.hpp"
 #include "gc/shenandoah/shenandoahMarkingContext.inline.hpp"
+#include "gc/shenandoah/shenandoahOldGeneration.hpp"
 #include "gc/shenandoah/shenandoahPartitionAllocator.hpp"
 #include "logging/log.hpp"
 
@@ -40,6 +41,13 @@ template<ShenandoahFreeSetPartitionId PARTITION>
 HeapWord* ShenandoahPartitionAllocator<PARTITION>::allocate(ShenandoahAllocRequest& req, bool& in_new_region) {
   shenandoah_assert_heaplocked();
 
+  // OldCollector: verify old generation has room before attempting allocation.
+  if constexpr (PARTITION == ShenandoahFreeSetPartitionId::OldCollector) {
+    if (!ShenandoahHeap::heap()->old_generation()->can_allocate(req)) {
+      return nullptr;
+    }
+  }
+
   // Fast path: try the retained region first.
   if (_retained_region != nullptr) {
     size_t min_size = req.is_lab_alloc() ? req.min_size() : req.size();

src/hotspot/share/gc/shenandoah/shenandoahSerialAllocator.cpp

@@ -24,6 +24,7 @@
 
 #include "gc/shenandoah/shenandoahAllocRequest.hpp"
 #include "gc/shenandoah/shenandoahFreeSet.hpp"
+#include "gc/shenandoah/shenandoahHeap.inline.hpp"
 #include "gc/shenandoah/shenandoahHeapRegion.hpp"
 #include "gc/shenandoah/shenandoahSerialAllocator.hpp"
 
@@ -34,7 +35,9 @@ ShenandoahSerialAllocator::ShenandoahSerialAllocator(ShenandoahFreeSet* free_set
     _old_collector_alloc(free_set) {}
 
 HeapWord* ShenandoahSerialAllocator::allocate(ShenandoahAllocRequest& req, bool& in_new_region) {
-  shenandoah_assert_heaplocked();
+  // Serial allocator acquires heap lock for all allocations.
+  // Mutator allocations may yield to safepoint; GC allocations cannot.
+  ShenandoahHeapLocker locker(ShenandoahHeap::heap()->lock(), req.is_mutator_alloc());
 
   if (ShenandoahHeapRegion::requires_humongous(req.size())) {
     switch (req.type()) {