Implement -Z oom=panic

compiler/rustc_codegen_cranelift/src/allocator.rs

@@ -4,6 +4,7 @@
 use crate::prelude::*;
 
 use rustc_ast::expand::allocator::{AllocatorKind, AllocatorTy, ALLOCATOR_METHODS};
+use rustc_session::config::OomStrategy;
 
 /// Returns whether an allocator shim was created
 pub(crate) fn codegen(
@@ -18,7 +19,13 @@ pub(crate) fn codegen(
     if any_dynamic_crate {
         false
     } else if let Some(kind) = tcx.allocator_kind(()) {
-        codegen_inner(module, unwind_context, kind, tcx.lang_items().oom().is_some());
+        codegen_inner(
+            module,
+            unwind_context,
+            kind,
+            tcx.lang_items().oom().is_some(),
+            tcx.sess.opts.debugging_opts.oom,
+        );
         true
     } else {
         false
@@ -30,6 +37,7 @@ fn codegen_inner(
     unwind_context: &mut UnwindContext,
     kind: AllocatorKind,
     has_alloc_error_handler: bool,
+    oom_strategy: OomStrategy,
 ) {
     let usize_ty = module.target_config().pointer_type();
 
@@ -129,4 +137,11 @@ fn codegen_inner(
     }
     module.define_function(func_id, &mut ctx).unwrap();
     unwind_context.add_function(func_id, &ctx, module.isa());
+
+    let data_id = module.declare_data(OomStrategy::SYMBOL, Linkage::Export, false, false).unwrap();
+    let mut data_ctx = DataContext::new();
+    data_ctx.set_align(1);
+    let val = oom_strategy.should_panic();
+    data_ctx.define(Box::new([val]));
+    module.define_data(data_id, &data_ctx).unwrap();
 }

compiler/rustc_codegen_gcc/src/allocator.rs

@@ -1,7 +1,8 @@
-use gccjit::{FunctionType, ToRValue};
+use gccjit::{FunctionType, GlobalKind, ToRValue};
 use rustc_ast::expand::allocator::{AllocatorKind, AllocatorTy, ALLOCATOR_METHODS};
 use rustc_middle::bug;
 use rustc_middle::ty::TyCtxt;
+use rustc_session::config::OomStrategy;
 use rustc_span::symbol::sym;
 
 use crate::GccContext;
@@ -113,4 +114,10 @@ pub(crate) unsafe fn codegen(tcx: TyCtxt<'_>, mods: &mut GccContext, _module_nam
     let _ret = context.new_call(None, callee, &args);
     //llvm::LLVMSetTailCall(ret, True);
     block.end_with_void_return(None);
+
+    let name = OomStrategy::SYMBOL.to_string();
+    let global = context.new_global(None, GlobalKind::Exported, i8, name);
+    let value = tcx.sess.opts.debugging_opts.oom.should_panic();
+    let value = context.new_rvalue_from_int(i8, value as i32);
+    global.global_set_initializer_rvalue(value);
 }

compiler/rustc_codegen_llvm/src/allocator.rs

@@ -3,7 +3,7 @@ use libc::c_uint;
 use rustc_ast::expand::allocator::{AllocatorKind, AllocatorTy, ALLOCATOR_METHODS};
 use rustc_middle::bug;
 use rustc_middle::ty::TyCtxt;
-use rustc_session::config::DebugInfo;
+use rustc_session::config::{DebugInfo, OomStrategy};
 use rustc_span::symbol::sym;
 
 use crate::debuginfo;
@@ -136,6 +136,16 @@ pub(crate) unsafe fn codegen(
     llvm::LLVMBuildRetVoid(llbuilder);
     llvm::LLVMDisposeBuilder(llbuilder);
 
+    // __rust_alloc_error_handler_should_panic
+    let name = OomStrategy::SYMBOL;
+    let ll_g = llvm::LLVMRustGetOrInsertGlobal(llmod, name.as_ptr().cast(), name.len(), i8);
+    if tcx.sess.target.default_hidden_visibility {
+        llvm::LLVMRustSetVisibility(ll_g, llvm::Visibility::Hidden);
+    }
+    let val = tcx.sess.opts.debugging_opts.oom.should_panic();
+    let llval = llvm::LLVMConstInt(i8, val as u64, False);
+    llvm::LLVMSetInitializer(ll_g, llval);
+
     if tcx.sess.opts.debuginfo != DebugInfo::None {
         let dbg_cx = debuginfo::CrateDebugContext::new(llmod);
         debuginfo::metadata::compile_unit_metadata(tcx, module_name, &dbg_cx);

compiler/rustc_interface/src/tests.rs

@@ -9,8 +9,8 @@ use rustc_session::config::{
     rustc_optgroups, ErrorOutputType, ExternLocation, LocationDetail, Options, Passes,
 };
 use rustc_session::config::{
-    BranchProtection, Externs, OutputType, OutputTypes, PAuthKey, PacRet, SymbolManglingVersion,
-    WasiExecModel,
+    BranchProtection, Externs, OomStrategy, OutputType, OutputTypes, PAuthKey, PacRet,
+    SymbolManglingVersion, WasiExecModel,
 };
 use rustc_session::config::{CFGuard, ExternEntry, LinkerPluginLto, LtoCli, SwitchWithOptPath};
 use rustc_session::lint::Level;
@@ -755,6 +755,7 @@ fn test_debugging_options_tracking_hash() {
     tracked!(no_link, true);
     tracked!(no_unique_section_names, true);
     tracked!(no_profiler_runtime, true);
+    tracked!(oom, OomStrategy::Panic);
     tracked!(osx_rpath_install_name, true);
     tracked!(panic_abort_tests, true);
     tracked!(panic_in_drop, PanicStrategy::Abort);

compiler/rustc_session/src/config.rs

@@ -2735,9 +2735,9 @@ impl PpMode {
 crate mod dep_tracking {
     use super::{
         BranchProtection, CFGuard, CFProtection, CrateType, DebugInfo, ErrorOutputType,
-        InstrumentCoverage, LdImpl, LinkerPluginLto, LocationDetail, LtoCli, OptLevel, OutputType,
-        OutputTypes, Passes, SourceFileHashAlgorithm, SwitchWithOptPath, SymbolManglingVersion,
-        TrimmedDefPaths,
+        InstrumentCoverage, LdImpl, LinkerPluginLto, LocationDetail, LtoCli, OomStrategy, OptLevel,
+        OutputType, OutputTypes, Passes, SourceFileHashAlgorithm, SwitchWithOptPath,
+        SymbolManglingVersion, TrimmedDefPaths,
     };
     use crate::lint;
     use crate::options::WasiExecModel;
@@ -2833,6 +2833,7 @@ crate mod dep_tracking {
         RealFileName,
         LocationDetail,
         BranchProtection,
+        OomStrategy,
     );
 
     impl<T1, T2> DepTrackingHash for (T1, T2)
@@ -2922,3 +2923,24 @@ crate mod dep_tracking {
         }
     }
 }
+
+/// Default behavior to use in out-of-memory situations.
+#[derive(Clone, Copy, PartialEq, Hash, Debug, Encodable, Decodable, HashStable_Generic)]
+pub enum OomStrategy {
+    /// Generate a panic that can be caught by `catch_unwind`.
+    Panic,
+
+    /// Abort the process immediately.
+    Abort,
+}
+
+impl OomStrategy {
+    pub const SYMBOL: &'static str = "__rust_alloc_error_handler_should_panic";
+
+    pub fn should_panic(self) -> u8 {
+        match self {
+            OomStrategy::Panic => 1,
+            OomStrategy::Abort => 0,
+        }
+    }
+}

compiler/rustc_session/src/options.rs

@@ -375,6 +375,7 @@ mod desc {
     pub const parse_passes: &str = "a space-separated list of passes, or `all`";
     pub const parse_panic_strategy: &str = "either `unwind` or `abort`";
     pub const parse_opt_panic_strategy: &str = parse_panic_strategy;
+    pub const parse_oom_strategy: &str = "either `panic` or `abort`";
     pub const parse_relro_level: &str = "one of: `full`, `partial`, or `off`";
     pub const parse_sanitizers: &str = "comma separated list of sanitizers: `address`, `cfi`, `hwaddress`, `leak`, `memory`, `memtag`, or `thread`";
     pub const parse_sanitizer_memory_track_origins: &str = "0, 1, or 2";
@@ -620,6 +621,15 @@ mod parse {
         true
     }
 
+    crate fn parse_oom_strategy(slot: &mut OomStrategy, v: Option<&str>) -> bool {
+        match v {
+            Some("panic") => *slot = OomStrategy::Panic,
+            Some("abort") => *slot = OomStrategy::Abort,
+            _ => return false,
+        }
+        true
+    }
+
     crate fn parse_relro_level(slot: &mut Option<RelroLevel>, v: Option<&str>) -> bool {
         match v {
             Some(s) => match s.parse::<RelroLevel>() {
@@ -1328,6 +1338,8 @@ options! {
         "prevent automatic injection of the profiler_builtins crate"),
     normalize_docs: bool = (false, parse_bool, [TRACKED],
         "normalize associated items in rustdoc when generating documentation"),
+    oom: OomStrategy = (OomStrategy::Abort, parse_oom_strategy, [TRACKED],
+        "panic strategy for out-of-memory handling"),
     osx_rpath_install_name: bool = (false, parse_bool, [TRACKED],
         "pass `-install_name @rpath/...` to the macOS linker (default: no)"),
     panic_abort_tests: bool = (false, parse_bool, [TRACKED],

library/std/src/alloc.rs

@@ -315,7 +315,21 @@ pub fn take_alloc_error_hook() -> fn(Layout) {
 }
 
 fn default_alloc_error_hook(layout: Layout) {
-    rtprintpanic!("memory allocation of {} bytes failed\n", layout.size());
+    #[cfg(not(bootstrap))]
+    extern "Rust" {
+        // This symbol is emitted by rustc next to __rust_alloc_error_handler.
+        // Its value depends on the -Zoom={panic,abort} compiler option.
+        static __rust_alloc_error_handler_should_panic: u8;
+    }
+    #[cfg(bootstrap)]
+    let __rust_alloc_error_handler_should_panic = 0;
+
+    #[allow(unused_unsafe)]
+    if unsafe { __rust_alloc_error_handler_should_panic != 0 } {
+        panic!("memory allocation of {} bytes failed\n", layout.size());
+    } else {
+        rtprintpanic!("memory allocation of {} bytes failed\n", layout.size());
+    }
 }
 
 #[cfg(not(test))]

src/test/ui/oom_unwind.rs

@@ -0,0 +1,23 @@
+// compile-flags: -Z oom=panic
+// run-pass
+// no-prefer-dynamic
+// needs-unwind
+// only-linux
+
+#![feature(bench_black_box)]
+
+use std::hint::black_box;
+use std::mem::forget;
+use std::panic::catch_unwind;
+
+fn main() {
+    let panic = catch_unwind(|| {
+        // This is guaranteed to exceed even the size of the address space
+        for _ in 0..16 {
+            // Truncates to a suitable value for both 32-bit and 64-bit targets.
+            let alloc_size = 0x1000_0000_1000_0000u64 as usize;
+            forget(black_box(vec![0u8; alloc_size]));
+        }
+    });
+    assert!(panic.is_err());
+}

src/tools/tidy/src/ui_tests.rs

@@ -7,7 +7,7 @@ use std::path::Path;
 
 const ENTRY_LIMIT: usize = 1000;
 // FIXME: The following limits should be reduced eventually.
-const ROOT_ENTRY_LIMIT: usize = 983;
+const ROOT_ENTRY_LIMIT: usize = 984;
 const ISSUES_ENTRY_LIMIT: usize = 2310;
 
 fn check_entries(path: &Path, bad: &mut bool) {