From 0ade5a11f5705e33e031ae5b47c21fd0553d6675 Mon Sep 17 00:00:00 2001
From: Josh Stone <jistone@redhat.com>
Date: Fri, 15 Mar 2024 15:49:06 -0700
Subject: [PATCH] Register LLVM handlers for bad-alloc / OOM

LLVM's default bad-alloc handler may throw if exceptions are enabled,
and `operator new` isn't hooked at all by default. Now we register our
own handler that prints a message similar to fatal errors, then aborts.
We also call the function that registers the C++ `std::new_handler`.
---
 compiler/rustc_codegen_llvm/src/llvm/ffi.rs   |  2 +-
 compiler/rustc_codegen_llvm/src/llvm_util.rs  |  2 +-
 .../rustc_llvm/llvm-wrapper/RustWrapper.cpp   | 25 ++++++++++++++++++-
 3 files changed, 26 insertions(+), 3 deletions(-)

diff --git a/compiler/rustc_codegen_llvm/src/llvm/ffi.rs b/compiler/rustc_codegen_llvm/src/llvm/ffi.rs
index d34c289887e..284bc74d5c4 100644
--- a/compiler/rustc_codegen_llvm/src/llvm/ffi.rs
+++ b/compiler/rustc_codegen_llvm/src/llvm/ffi.rs
@@ -1519,7 +1519,7 @@ extern "C" {
 
 #[link(name = "llvm-wrapper", kind = "static")]
 extern "C" {
-    pub fn LLVMRustInstallFatalErrorHandler();
+    pub fn LLVMRustInstallErrorHandlers();
     pub fn LLVMRustDisableSystemDialogsOnCrash();
 
     // Create and destroy contexts.
diff --git a/compiler/rustc_codegen_llvm/src/llvm_util.rs b/compiler/rustc_codegen_llvm/src/llvm_util.rs
index e32c38644aa..c9e62e504ae 100644
--- a/compiler/rustc_codegen_llvm/src/llvm_util.rs
+++ b/compiler/rustc_codegen_llvm/src/llvm_util.rs
@@ -49,7 +49,7 @@ unsafe fn configure_llvm(sess: &Session) {
     let mut llvm_c_strs = Vec::with_capacity(n_args + 1);
     let mut llvm_args = Vec::with_capacity(n_args + 1);
 
-    llvm::LLVMRustInstallFatalErrorHandler();
+    llvm::LLVMRustInstallErrorHandlers();
     // On Windows, an LLVM assertion will open an Abort/Retry/Ignore dialog
     // box for the purpose of launching a debugger. However, on CI this will
     // cause it to hang until it times out, which can take several hours.
diff --git a/compiler/rustc_llvm/llvm-wrapper/RustWrapper.cpp b/compiler/rustc_llvm/llvm-wrapper/RustWrapper.cpp
index 1632b9e1249..072620c65a5 100644
--- a/compiler/rustc_llvm/llvm-wrapper/RustWrapper.cpp
+++ b/compiler/rustc_llvm/llvm-wrapper/RustWrapper.cpp
@@ -25,6 +25,13 @@
 
 #include <iostream>
 
+// for raw `write` in the bad-alloc handler
+#ifdef _MSC_VER
+#include <io.h>
+#else
+#include <unistd.h>
+#endif
+
 //===----------------------------------------------------------------------===
 //
 // This file defines alternate interfaces to core functions that are more
@@ -88,8 +95,24 @@ static void FatalErrorHandler(void *UserData,
   exit(101);
 }
 
-extern "C" void LLVMRustInstallFatalErrorHandler() {
+// Custom error handler for bad-alloc LLVM errors.
+//
+// It aborts the process without any further allocations, similar to LLVM's
+// default except that may be configured to `throw std::bad_alloc()` instead.
+static void BadAllocErrorHandler(void *UserData,
+                                 const char* Reason,
+                                 bool GenCrashDiag) {
+  const char *OOM = "rustc-LLVM ERROR: out of memory\n";
+  write(2, OOM, strlen(OOM));
+  write(2, Reason, strlen(Reason));
+  write(2, "\n", 1);
+  abort();
+}
+
+extern "C" void LLVMRustInstallErrorHandlers() {
   install_fatal_error_handler(FatalErrorHandler);
+  install_bad_alloc_error_handler(BadAllocErrorHandler);
+  install_out_of_memory_new_handler();
 }
 
 extern "C" void LLVMRustDisableSystemDialogsOnCrash() {