|
15 | 15 | use core::intrinsics::{min_align_of_val, size_of_val};
|
16 | 16 | use core::ptr::{NonNull, Unique};
|
17 | 17 | use core::usize;
|
| 18 | +#[cfg(target_has_atomic = "ptr")] use core::{mem, ptr}; |
| 19 | +#[cfg(target_has_atomic = "ptr")] use core::sync::atomic::{AtomicPtr, Ordering}; |
18 | 20 |
|
19 | 21 | #[stable(feature = "alloc_module", since = "1.28.0")]
|
20 | 22 | #[doc(inline)]
|
@@ -191,12 +193,89 @@ pub(crate) unsafe fn box_free<T: ?Sized>(ptr: Unique<T>) {
|
191 | 193 | #[stable(feature = "global_alloc", since = "1.28.0")]
|
192 | 194 | #[rustc_allocator_nounwind]
|
193 | 195 | pub fn handle_alloc_error(layout: Layout) -> ! {
|
194 |
| - #[allow(improper_ctypes)] |
195 |
| - extern "Rust" { |
196 |
| - #[lang = "oom"] |
197 |
| - fn oom_impl(layout: Layout) -> !; |
| 196 | + #[cfg(target_has_atomic = "ptr")] |
| 197 | + { |
| 198 | + let hook = HOOK.load(Ordering::SeqCst); |
| 199 | + let hook: fn(Layout) = if hook.is_null() { |
| 200 | + default_alloc_error_hook |
| 201 | + } else { |
| 202 | + unsafe { mem::transmute(hook) } |
| 203 | + }; |
| 204 | + hook(layout); |
| 205 | + } |
| 206 | + #[cfg(not(target_has_atomic = "ptr"))] |
| 207 | + { |
| 208 | + default_alloc_error_hook(layout) |
| 209 | + } |
| 210 | + |
| 211 | + unsafe { |
| 212 | + #[cfg(not(stage0))] |
| 213 | + { |
| 214 | + #[allow(improper_ctypes)] |
| 215 | + extern "Rust" { |
| 216 | + // This function is generated by the compiler |
| 217 | + // and calls __rust_abort_internal in src/libstd/alloc.rs |
| 218 | + // if the `std` crate is linked. |
| 219 | + fn __rust_maybe_abort_internal(); |
| 220 | + } |
| 221 | + __rust_maybe_abort_internal(); |
| 222 | + } |
| 223 | + // The above does nothing and returns when `std` is not linked. |
| 224 | + // In that case we abort "naively" here: |
| 225 | + ::core::intrinsics::abort() |
| 226 | + } |
| 227 | +} |
| 228 | + |
| 229 | +fn default_alloc_error_hook(layout: Layout) { |
| 230 | + #[cfg(stage0)] |
| 231 | + let _ = layout; |
| 232 | + #[cfg(not(stage0))] |
| 233 | + unsafe { |
| 234 | + #[allow(improper_ctypes)] |
| 235 | + extern "Rust" { |
| 236 | + // This function is generated by the compiler |
| 237 | + // and calls __rust_default_alloc_error_hook in src/libstd/alloc.rs |
| 238 | + // if the `std` crate is linked. |
| 239 | + fn __rust_maybe_default_alloc_error_hook(size: usize, align: usize); |
| 240 | + } |
| 241 | + __rust_maybe_default_alloc_error_hook(layout.size(), layout.align()) |
| 242 | + } |
| 243 | +} |
| 244 | + |
| 245 | +#[cfg(target_has_atomic = "ptr")] |
| 246 | +static HOOK: AtomicPtr<()> = AtomicPtr::new(ptr::null_mut()); |
| 247 | + |
| 248 | +/// Registers a custom allocation error hook, replacing any that was previously registered. |
| 249 | +/// |
| 250 | +/// The allocation error hook is invoked when an infallible memory allocation fails, before |
| 251 | +/// the runtime aborts. The default hook prints a message to standard error, |
| 252 | +/// but this behavior can be customized with the [`set_alloc_error_hook`] and |
| 253 | +/// [`take_alloc_error_hook`] functions. |
| 254 | +/// |
| 255 | +/// The hook is provided with a `Layout` struct which contains information |
| 256 | +/// about the allocation that failed. |
| 257 | +/// |
| 258 | +/// The allocation error hook is a global resource. |
| 259 | +#[unstable(feature = "alloc_error_hook", issue = "51245")] |
| 260 | +#[cfg(target_has_atomic = "ptr")] |
| 261 | +pub fn set_alloc_error_hook(hook: fn(Layout)) { |
| 262 | + HOOK.store(hook as *mut (), Ordering::SeqCst); |
| 263 | +} |
| 264 | + |
| 265 | +/// Unregisters the current allocation error hook, returning it. |
| 266 | +/// |
| 267 | +/// *See also the function [`set_alloc_error_hook`].* |
| 268 | +/// |
| 269 | +/// If no custom hook is registered, the default hook will be returned. |
| 270 | +#[unstable(feature = "alloc_error_hook", issue = "51245")] |
| 271 | +#[cfg(target_has_atomic = "ptr")] |
| 272 | +pub fn take_alloc_error_hook() -> fn(Layout) { |
| 273 | + let hook = HOOK.swap(ptr::null_mut(), Ordering::SeqCst); |
| 274 | + if hook.is_null() { |
| 275 | + default_alloc_error_hook |
| 276 | + } else { |
| 277 | + unsafe { mem::transmute(hook) } |
198 | 278 | }
|
199 |
| - unsafe { oom_impl(layout) } |
200 | 279 | }
|
201 | 280 |
|
202 | 281 | #[cfg(test)]
|
|
0 commit comments