8318586: Explicitly handle upcall stub allocation failure

JornVernee · JornVernee · commit e96e19195287 · 2023-11-30T19:26:54.000Z
8318653: UpcallTestHelper::runInNewProcess waits for forked process without timeout

Reviewed-by: shade, mcimadamore
diff --git a/src/hotspot/cpu/aarch64/downcallLinker_aarch64.cpp b/src/hotspot/cpu/aarch64/downcallLinker_aarch64.cpp
@@ -51,19 +51,29 @@ RuntimeStub* DowncallLinker::make_downcall_stub(BasicType* signature,
   int code_size = native_invoker_code_base_size + (num_args * native_invoker_size_per_arg);
   int locs_size = 1; // must be non-zero
   CodeBuffer code("nep_invoker_blob", code_size, locs_size);
+  if (code.blob() == nullptr) {
+    return nullptr;
+  }
   StubGenerator g(&code, signature, num_args, ret_bt, abi,
                   input_registers, output_registers,
                   needs_return_buffer, captured_state_mask,
                   needs_transition);
   g.generate();
   code.log_section_sizes("nep_invoker_blob");
 
+  bool caller_must_gc_arguments = false;
+  bool alloc_fail_is_fatal = false;
   RuntimeStub* stub =
     RuntimeStub::new_runtime_stub("nep_invoker_blob",
                                   &code,
                                   g.frame_complete(),
                                   g.framesize(),
-                                  g.oop_maps(), false);
+                                  g.oop_maps(),
+                                  caller_must_gc_arguments,
+                                  alloc_fail_is_fatal);
+  if (stub == nullptr) {
+    return nullptr;
+  }
 
 #ifndef PRODUCT
   LogTarget(Trace, foreign, downcall) lt;
diff --git a/src/hotspot/cpu/aarch64/upcallLinker_aarch64.cpp b/src/hotspot/cpu/aarch64/upcallLinker_aarch64.cpp
@@ -127,6 +127,9 @@ address UpcallLinker::make_upcall_stub(jobject receiver, Method* entry,
   const CallRegs call_regs = ForeignGlobals::parse_call_regs(jconv);
   int code_size = upcall_stub_code_base_size + (total_out_args * upcall_stub_size_per_arg);
   CodeBuffer buffer("upcall_stub", code_size, /* locs_size = */ 1);
+  if (buffer.blob() == nullptr) {
+    return nullptr;
+  }
 
   GrowableArray<VMStorage> unfiltered_out_regs;
   int out_arg_bytes = ForeignGlobals::java_calling_convention(out_sig_bt, total_out_args, unfiltered_out_regs);
@@ -324,6 +327,9 @@ address UpcallLinker::make_upcall_stub(jobject receiver, Method* entry,
                          &buffer,
                          receiver,
                          in_ByteSize(frame_data_offset));
+  if (blob == nullptr) {
+    return nullptr;
+  }
 
 #ifndef PRODUCT
   if (lt.is_enabled()) {
diff --git a/src/hotspot/cpu/ppc/downcallLinker_ppc.cpp b/src/hotspot/cpu/ppc/downcallLinker_ppc.cpp
@@ -51,19 +51,29 @@ RuntimeStub* DowncallLinker::make_downcall_stub(BasicType* signature,
   int code_size = native_invoker_code_base_size + (num_args * native_invoker_size_per_arg);
   int locs_size = 1; // must be non-zero
   CodeBuffer code("nep_invoker_blob", code_size, locs_size);
+  if (code.blob() == nullptr) {
+    return nullptr;
+  }
   StubGenerator g(&code, signature, num_args, ret_bt, abi,
                   input_registers, output_registers,
                   needs_return_buffer, captured_state_mask,
                   needs_transition);
   g.generate();
   code.log_section_sizes("nep_invoker_blob");
 
+  bool caller_must_gc_arguments = false;
+  bool alloc_fail_is_fatal = false;
   RuntimeStub* stub =
     RuntimeStub::new_runtime_stub("nep_invoker_blob",
                                   &code,
                                   g.frame_complete(),
                                   g.framesize(),
-                                  g.oop_maps(), false);
+                                  g.oop_maps(),
+                                  caller_must_gc_arguments,
+                                  alloc_fail_is_fatal);
+  if (stub == nullptr) {
+    return nullptr;
+  }
 
 #ifndef PRODUCT
   LogTarget(Trace, foreign, downcall) lt;
diff --git a/src/hotspot/cpu/ppc/upcallLinker_ppc.cpp b/src/hotspot/cpu/ppc/upcallLinker_ppc.cpp
@@ -128,6 +128,9 @@ address UpcallLinker::make_upcall_stub(jobject receiver, Method* entry,
   const CallRegs call_regs = ForeignGlobals::parse_call_regs(jconv);
   int code_size = upcall_stub_code_base_size + (total_out_args * upcall_stub_size_per_arg);
   CodeBuffer buffer("upcall_stub", code_size, /* locs_size = */ 1);
+  if (buffer.blob() == nullptr) {
+    return nullptr;
+  }
 
   Register callerSP            = R2, // C/C++ uses R2 as TOC, but we can reuse it here
            tmp                 = R11_scratch1, // same as shuffle_reg
@@ -332,6 +335,10 @@ address UpcallLinker::make_upcall_stub(jobject receiver, Method* entry,
                          &buffer,
                          receiver,
                          in_ByteSize(frame_data_offset));
+  if (blob == nullptr) {
+    return nullptr;
+  }
+
 #ifndef ABI_ELFv2
   // Need to patch the FunctionDescriptor after relocating.
   address fd_addr = blob->code_begin();
diff --git a/src/hotspot/cpu/riscv/downcallLinker_riscv.cpp b/src/hotspot/cpu/riscv/downcallLinker_riscv.cpp
@@ -52,19 +52,29 @@ RuntimeStub* DowncallLinker::make_downcall_stub(BasicType* signature,
   int code_size = native_invoker_code_base_size + (num_args * native_invoker_size_per_arg);
   int locs_size = 1; // must be non-zero
   CodeBuffer code("nep_invoker_blob", code_size, locs_size);
+  if (code.blob() == nullptr) {
+    return nullptr;
+  }
   StubGenerator g(&code, signature, num_args, ret_bt, abi,
                   input_registers, output_registers,
                   needs_return_buffer, captured_state_mask,
                   needs_transition);
   g.generate();
   code.log_section_sizes("nep_invoker_blob");
 
+  bool caller_must_gc_arguments = false;
+  bool alloc_fail_is_fatal = false;
   RuntimeStub* stub =
     RuntimeStub::new_runtime_stub("nep_invoker_blob",
                                   &code,
                                   g.frame_complete(),
                                   g.framesize(),
-                                  g.oop_maps(), false);
+                                  g.oop_maps(),
+                                  caller_must_gc_arguments,
+                                  alloc_fail_is_fatal);
+  if (stub == nullptr) {
+    return nullptr;
+  }
 
 #ifndef PRODUCT
   LogTarget(Trace, foreign, downcall) lt;
diff --git a/src/hotspot/cpu/riscv/upcallLinker_riscv.cpp b/src/hotspot/cpu/riscv/upcallLinker_riscv.cpp
@@ -128,6 +128,9 @@ address UpcallLinker::make_upcall_stub(jobject receiver, Method* entry,
   const CallRegs call_regs = ForeignGlobals::parse_call_regs(jconv);
   int code_size = upcall_stub_code_base_size + (total_out_args * upcall_stub_size_per_arg);
   CodeBuffer buffer("upcall_stub", code_size, /* locs_size = */ 1);
+  if (buffer.blob() == nullptr) {
+    return nullptr;
+  }
 
   GrowableArray<VMStorage> unfiltered_out_regs;
   int out_arg_bytes = ForeignGlobals::java_calling_convention(out_sig_bt, total_out_args, unfiltered_out_regs);
@@ -344,6 +347,10 @@ address UpcallLinker::make_upcall_stub(jobject receiver, Method* entry,
                          &buffer,
                          receiver,
                          in_ByteSize(frame_data_offset));
+  if (blob == nullptr) {
+    return nullptr;
+  }
+
 #ifndef PRODUCT
   if (lt.is_enabled()) {
     ResourceMark rm;
diff --git a/src/hotspot/cpu/s390/downcallLinker_s390.cpp b/src/hotspot/cpu/s390/downcallLinker_s390.cpp
@@ -52,6 +52,9 @@ RuntimeStub* DowncallLinker::make_downcall_stub(BasicType* signature,
   int code_size = native_invoker_code_base_size + (num_args * native_invoker_size_per_args);
   int locs_size = 1; //must be non zero
   CodeBuffer code("nep_invoker_blob", code_size, locs_size);
+  if (code.blob() == nullptr) {
+    return nullptr;
+  }
 
   StubGenerator g(&code, signature, num_args, ret_bt, abi,
                   input_registers, output_registers,
@@ -60,12 +63,19 @@ RuntimeStub* DowncallLinker::make_downcall_stub(BasicType* signature,
   g.generate();
   code.log_section_sizes("nep_invoker_blob");
 
+  bool caller_must_gc_arguments = false;
+  bool alloc_fail_is_fatal = false;
   RuntimeStub* stub =
     RuntimeStub::new_runtime_stub("nep_invoker_blob",
                                   &code,
                                   g.frame_complete(),
                                   g.framesize(),
-                                  g.oop_maps(), false);
+                                  g.oop_maps(),
+                                  caller_must_gc_arguments,
+                                  alloc_fail_is_fatal);
+  if (stub == nullptr) {
+    return nullptr;
+  }
 
 #ifndef PRODUCT
   LogTarget(Trace, foreign, downcall) lt;
diff --git a/src/hotspot/cpu/s390/upcallLinker_s390.cpp b/src/hotspot/cpu/s390/upcallLinker_s390.cpp
@@ -126,6 +126,9 @@ address UpcallLinker::make_upcall_stub(jobject receiver, Method* entry,
   const CallRegs call_regs = ForeignGlobals::parse_call_regs(jconv);
   int code_size = upcall_stub_code_base_size + (total_out_args * upcall_stub_size_per_arg);
   CodeBuffer buffer("upcall_stub", code_size, /* locs_size = */ 0);
+  if (buffer.blob() == nullptr) {
+    return nullptr;
+  }
 
   Register call_target_address = Z_R1_scratch;
 
@@ -283,6 +286,10 @@ address UpcallLinker::make_upcall_stub(jobject receiver, Method* entry,
                          &buffer,
                          receiver,
                          in_ByteSize(frame_data_offset));
+  if (blob == nullptr) {
+    return nullptr;
+  }
+
 #ifndef PRODUCT
   if (lt.is_enabled()) {
     ResourceMark rm;
diff --git a/src/hotspot/cpu/x86/downcallLinker_x86_64.cpp b/src/hotspot/cpu/x86/downcallLinker_x86_64.cpp
@@ -49,19 +49,29 @@ RuntimeStub* DowncallLinker::make_downcall_stub(BasicType* signature,
   int code_size = native_invoker_code_base_size + (num_args * native_invoker_size_per_arg);
   int locs_size = 1; // can not be zero
   CodeBuffer code("nep_invoker_blob", code_size, locs_size);
+  if (code.blob() == nullptr) {
+    return nullptr;
+  }
   StubGenerator g(&code, signature, num_args, ret_bt, abi,
                   input_registers, output_registers,
                   needs_return_buffer, captured_state_mask,
                   needs_transition);
   g.generate();
   code.log_section_sizes("nep_invoker_blob");
 
+  bool caller_must_gc_arguments = false;
+  bool alloc_fail_is_fatal = false;
   RuntimeStub* stub =
     RuntimeStub::new_runtime_stub("nep_invoker_blob",
                                   &code,
                                   g.frame_complete(),
                                   g.framesize(),
-                                  g.oop_maps(), false);
+                                  g.oop_maps(),
+                                  caller_must_gc_arguments,
+                                  alloc_fail_is_fatal);
+  if (stub == nullptr) {
+    return nullptr;
+  }
 
 #ifndef PRODUCT
   LogTarget(Trace, foreign, downcall) lt;
diff --git a/src/hotspot/cpu/x86/upcallLinker_x86_64.cpp b/src/hotspot/cpu/x86/upcallLinker_x86_64.cpp
@@ -177,6 +177,9 @@ address UpcallLinker::make_upcall_stub(jobject receiver, Method* entry,
   const CallRegs call_regs = ForeignGlobals::parse_call_regs(jconv);
   int code_size = upcall_stub_code_base_size + (total_out_args * upcall_stub_size_per_arg);
   CodeBuffer buffer("upcall_stub", code_size, /* locs_size = */ 1);
+  if (buffer.blob() == nullptr) {
+    return nullptr;
+  }
 
   GrowableArray<VMStorage> unfiltered_out_regs;
   int out_arg_bytes = ForeignGlobals::java_calling_convention(out_sig_bt, total_out_args, unfiltered_out_regs);
@@ -379,6 +382,9 @@ address UpcallLinker::make_upcall_stub(jobject receiver, Method* entry,
                          &buffer,
                          receiver,
                          in_ByteSize(frame_data_offset));
+  if (blob == nullptr) {
+    return nullptr;
+  }
 
 #ifndef PRODUCT
   if (lt.is_enabled()) {
diff --git a/src/hotspot/share/code/codeBlob.cpp b/src/hotspot/share/code/codeBlob.cpp
@@ -759,6 +759,10 @@ UpcallStub* UpcallStub::create(const char* name, CodeBuffer* cb, jobject receive
     MutexLocker mu(CodeCache_lock, Mutex::_no_safepoint_check_flag);
     blob = new (size) UpcallStub(name, cb, size, receiver, frame_data_offset);
   }
+  if (blob == nullptr) {
+    return nullptr; // caller must handle this
+  }
+
   // Track memory usage statistic after releasing CodeCache_lock
   MemoryService::track_code_cache_memory_usage();
 
diff --git a/src/hotspot/share/prims/nativeEntryPoint.cpp b/src/hotspot/share/prims/nativeEntryPoint.cpp
@@ -78,10 +78,14 @@ JNI_ENTRY(jlong, NEP_makeDowncallStub(JNIEnv* env, jclass _unused, jobject metho
     output_regs.push(ForeignGlobals::parse_vmstorage(ret_moves_oop->obj_at(i)));
   }
 
-  return (jlong) DowncallLinker::make_downcall_stub(basic_type, pslots, ret_bt, abi,
-                                                    input_regs, output_regs,
-                                                    needs_return_buffer, captured_state_mask,
-                                                    needs_transition)->code_begin();
+  RuntimeStub* stub = DowncallLinker::make_downcall_stub(basic_type, pslots, ret_bt, abi,
+                                                         input_regs, output_regs,
+                                                         needs_return_buffer, captured_state_mask,
+                                                         needs_transition);
+  if (stub == nullptr) {
+    return 0;
+  }
+  return (jlong) stub->code_begin();
 JNI_END
 
 JNI_ENTRY(jboolean, NEP_freeDowncallStub(JNIEnv* env, jclass _unused, jlong invoker))
diff --git a/src/java.base/share/classes/jdk/internal/foreign/abi/NativeEntryPoint.java b/src/java.base/share/classes/jdk/internal/foreign/abi/NativeEntryPoint.java
@@ -71,6 +71,9 @@ public static NativeEntryPoint make(ABIDescriptor abi,
         return NEP_CACHE.get(key, k -> {
             long downcallStub = makeDowncallStub(methodType, abi, argMoves, returnMoves, needsReturnBuffer,
                                                  capturedStateMask, needsTransition);
+            if (downcallStub == 0) {
+                throw new OutOfMemoryError("Failed to allocate downcall stub");
+            }
             NativeEntryPoint nep = new NativeEntryPoint(methodType, downcallStub);
             CLEANER.register(nep, () -> freeDowncallStub(downcallStub));
             return nep;
diff --git a/src/java.base/share/classes/jdk/internal/foreign/abi/UpcallLinker.java b/src/java.base/share/classes/jdk/internal/foreign/abi/UpcallLinker.java
@@ -106,6 +106,9 @@ public static UpcallStubFactory makeFactory(MethodType targetType, ABIDescriptor
             doBindings = insertArguments(exactInvoker(doBindings.type()), 0, doBindings);
             long entryPoint = makeUpcallStub(doBindings, abi, conv,
                     callingSequence.needsReturnBuffer(), callingSequence.returnBufferSize());
+            if (entryPoint == 0) {
+                throw new OutOfMemoryError("Failed to allocate upcall stub");
+            }
             return UpcallStubs.makeUpcall(entryPoint, scope);
         };
     }
diff --git a/test/jdk/java/foreign/NativeTestHelper.java b/test/jdk/java/foreign/NativeTestHelper.java
@@ -175,9 +175,13 @@ public static MethodHandle downcallHandle(String symbol, FunctionDescriptor desc
     }
 
     public static MemorySegment upcallStub(Class<?> holder, String name, FunctionDescriptor descriptor) {
+        return upcallStub(holder, name, descriptor, Arena.ofAuto());
+    }
+
+    public static MemorySegment upcallStub(Class<?> holder, String name, FunctionDescriptor descriptor, Arena arena) {
         try {
             MethodHandle target = MethodHandles.lookup().findStatic(holder, name, descriptor.toMethodType());
-            return LINKER.upcallStub(target, descriptor, Arena.ofAuto());
+            return LINKER.upcallStub(target, descriptor, arena);
         } catch (ReflectiveOperationException e) {
             throw new RuntimeException(e);
         }
diff --git a/test/jdk/java/foreign/TestAddressDereference.java b/test/jdk/java/foreign/TestAddressDereference.java
@@ -128,6 +128,7 @@ public void testNativeUpcallArgNeg(long alignment, ValueLayout layout) throws Th
         if (!badAlign) return;
         runInNewProcess(UpcallTestRunner.class, true,
                 new String[] {Long.toString(alignment), layout.toString() })
+                .assertFailed()
                 .assertStdErrContains("alignment constraint for address");
     }
 
diff --git a/test/jdk/java/foreign/TestStubAllocFailure.java b/test/jdk/java/foreign/TestStubAllocFailure.java
@@ -0,0 +1,66 @@
+/*
+ * Copyright (c) 2023, 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
+ * under the terms of the GNU General Public License version 2 only, as
+ * published by the Free Software Foundation.
+ *
+ * This code is distributed in the hope that it will be useful, but WITHOUT
+ * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
+ * FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License
+ * version 2 for more details (a copy is included in the LICENSE file that
+ * accompanied this code).
+ *
+ * You should have received a copy of the GNU General Public License version
+ * 2 along with this work; if not, write to the Free Software Foundation,
+ * Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301 USA.
+ *
+ * Please contact Oracle, 500 Oracle Parkway, Redwood Shores, CA 94065 USA
+ * or visit www.oracle.com if you need additional information or have any
+ * questions.
+ */
+
+/*
+ * @test
+ * @library ../ /test/lib
+ * @requires jdk.foreign.linker != "FALLBACK"
+ * @run testng/othervm/native
+ *   --enable-native-access=ALL-UNNAMED
+ *   TestStubAllocFailure
+ */
+
+import java.lang.foreign.*;
+import java.io.IOException;
+import java.util.List;
+
+import org.testng.annotations.Test;
+
+import static org.testng.Assert.assertTrue;
+import static org.testng.Assert.fail;
+
+public class TestStubAllocFailure extends UpcallTestHelper {
+
+    @Test
+    public void testUpcallAllocFailure() throws IOException, InterruptedException {
+        runInNewProcess(UpcallRunner.class, true, List.of("-XX:ReservedCodeCacheSize=3M"), List.of())
+                .assertSuccess();
+    }
+
+    public static class UpcallRunner extends NativeTestHelper {
+        public static void main(String[] args) throws Throwable {
+            try (Arena arena = Arena.ofConfined()) {
+                while (true) {
+                    // allocate stubs until we crash
+                    upcallStub(UpcallRunner.class, "target", FunctionDescriptor.ofVoid(), arena);
+                }
+            } catch (OutOfMemoryError e) {
+                assertTrue(e.getMessage().contains("Failed to allocate upcall stub"));
+            }
+        }
+
+        public static void target() {
+            fail("Should not get here");
+        }
+    }
+}
diff --git a/test/jdk/java/foreign/TestUpcallException.java b/test/jdk/java/foreign/TestUpcallException.java
diff --git a/test/jdk/java/foreign/UpcallTestHelper.java b/test/jdk/java/foreign/UpcallTestHelper.java
diff --git a/test/jdk/java/foreign/critical/TestCriticalUpcall.java b/test/jdk/java/foreign/critical/TestCriticalUpcall.java
diff --git a/test/jdk/java/foreign/passheapsegment/TestPassHeapSegment.java b/test/jdk/java/foreign/passheapsegment/TestPassHeapSegment.java
