Trigger a panic on allocation error on nightly Rust

In issue rust-lang#66342, we're seeing extremely large allocations by rustc
(4GB for the `hex` crate, which is only a few hundred lines). This is
exhausing the memory on CI, causing jobs to fail intermittently.

This PR installs a custom allocation error hook for nightly compilers,
which attempts to trigger a panic after printing the initial error
message. Hopefully, this will allow us to retrieve a backtrace when one
of these large spurious allocations occurs.

The hook is installed in `librustc_driver`, so that other compiler
frontends (e.g. clippy) will get this logic as well.

I'm unsure if this needs to be behind any kind of additional feature
gate, beyond being ngithly only. This only affects compiler frontends,
not generic users of `libstd`. While this will nake OOM errors on
nightly much more verbose, I don't think this is necessarily a bad
thing. I would expect that out of memory errors when running the
compiler are usually infrequent, so most users will probably never
notice this change. If any users are experiencing rust-lang#66342 (or something
like it) on their own crates, the extra output might even be useful to
them.

I don't know of any reasonable way of writing a test for this. I
manually verified the implementation by inserting:
`let _a: Vec<usize> = Vec::with_capacity(9999999999)` into
`librustc_driver` after the hook installation, and verified that
a backtrace was printed.

If we're very unlucky, it may turn out the large allocation on CI
happens to occur after several large successful allocations, leaving
extremely little memory left when we try to panic. If this is the case,
then we may fail to panic or print the backtrace, since panicking
currently allocates memory. However, we will still print an error
message, so the output will be no less useful (though a little more
spammy) then before.

Author: Aaron1011

src/librustc_driver/lib.rs

@@ -12,6 +12,8 @@
 #![feature(set_stdio)]
 #![feature(no_debug)]
 #![feature(integer_atomics)]
+#![feature(alloc_error_hook)]
+#![feature(alloc_internals)]
 
 #![recursion_limit="256"]
 
@@ -149,6 +151,14 @@ pub fn run_compiler(
     file_loader: Option<Box<dyn FileLoader + Send + Sync>>,
     emitter: Option<Box<dyn Write + Send>>
 ) -> interface::Result<()> {
+
+    // Install a custom alloc error hook on Nightly, to try
+    // so we can try to generate a backtrace on OOM in the compiler
+    // See https://github.com/rust-lang/rust/issues/66342 for details
+    if UnstableFeatures::from_environment().is_nightly_build() {
+        std::alloc::set_alloc_error_hook(std::alloc::rustc_alloc_error_hook);
+    }
+
     let mut args = Vec::new();
     for arg in at_args {
         match args::arg_expand(arg.clone()) {

src/libstd/alloc.rs

@@ -198,6 +198,21 @@ pub fn take_alloc_error_hook() -> fn(Layout) {
     }
 }
 
+/// A custom error hook, used by librustc_driver.
+/// This lives in `libstd` so that we can call `dumb_print`.
+///
+/// This hook prints out a message like the default error hook,
+/// but then panics to generate a backtrace. If we are completely
+/// out of memory, panicking and/or backtrace generation may fail:
+/// this is fine, since the backtrace is best-effort only. We
+/// are guaranteed to print the actual error message, though.
+#[doc(hidden)]
+#[unstable(feature = "alloc_internals", issue = "0")]
+pub fn rustc_alloc_error_hook(layout: Layout) {
+    dumb_print(format_args!("memory allocation of {} bytes failed. backtrace:", layout.size()));
+    panic!();
+}
+
 fn default_alloc_error_hook(layout: Layout) {
     dumb_print(format_args!("memory allocation of {} bytes failed", layout.size()));
 }