src/libstd/alloc.rs

   1 // Copyright 2017 The Rust Project Developers. See the COPYRIGHT
   2 // file at the top-level directory of this distribution and at
   3 // http://rust-lang.org/COPYRIGHT.
   4 //
   5 // Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
   6 // http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
   7 // <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
   8 // option. This file may not be copied, modified, or distributed
   9 // except according to those terms.
  10
  11 //! Memory allocation APIs
  12 //!
  13 //! In a given program, the standard library has one “global” memory allocator
  14 //! that is used for example by `Box<T>` and `Vec<T>`.
  15 //!
  16 //! Currently the default global allocator is unspecified.
  17 //! The compiler may link to a version of [jemalloc] on some platforms,
  18 //! but this is not guaranteed.
  19 //! Libraries, however, like `cdylib`s and `staticlib`s are guaranteed
  20 //! to use the [`System`] by default.
  21 //!
  22 //! [jemalloc]: https://github.com/jemalloc/jemalloc
  23 //! [`System`]: struct.System.html
  24 //!
  25 //! # The `#[global_allocator]` attribute
  26 //!
  27 //! This attribute allows configuring the choice of global allocator.
  28 //! You can use this to implement a completely custom global allocator
  29 //! to route all default allocation requests to a custom object.
  30 //!
  31 //! ```rust
  32 //! use std::alloc::{GlobalAlloc, System, Layout};
  33 //!
  34 //! struct MyAllocator;
  35 //!
  36 //! unsafe impl GlobalAlloc for MyAllocator {
  37 //!     unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
  38 //!         System.alloc(layout)
  39 //!     }
  40 //!
  41 //!     unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
  42 //!         System.dealloc(ptr, layout)
  43 //!     }
  44 //! }
  45 //!
  46 //! #[global_allocator]
  47 //! static GLOBAL: MyAllocator = MyAllocator;
  48 //!
  49 //! fn main() {
  50 //!     // This `Vec` will allocate memory through `GLOBAL` above
  51 //!     let mut v = Vec::new();
  52 //!     v.push(1);
  53 //! }
  54 //! ```
  55 //!
  56 //! The attribute is used on a `static` item whose type implements the
  57 //! [`GlobalAlloc`] trait. This type can be provided by an external library:
  58 //!
  59 //! [`GlobalAlloc`]: ../../core/alloc/trait.GlobalAlloc.html
  60 //!
  61 //! ```rust,ignore (demonstrates crates.io usage)
  62 //! extern crate jemallocator;
  63 //!
  64 //! use jemallacator::Jemalloc;
  65 //!
  66 //! #[global_allocator]
  67 //! static GLOBAL: Jemalloc = Jemalloc;
  68 //!
  69 //! fn main() {}
  70 //! ```
  71 //!
  72 //! The `#[global_allocator]` can only be used once in a crate
  73 //! or its recursive dependencies.
  74
  75 #![stable(feature = "alloc_module", since = "1.28.0")]
  76
  77 use core::sync::atomic::{AtomicPtr, Ordering};
  78 use core::{mem, ptr};
  79 use sys_common::util::dumb_print;
  80
  81 #[stable(feature = "alloc_module", since = "1.28.0")]
  82 #[doc(inline)]
  83 pub use alloc_crate::alloc::*;
  84
  85 #[stable(feature = "alloc_system_type", since = "1.28.0")]
  86 #[doc(inline)]
  87 pub use alloc_system::System;
  88
  89 static HOOK: AtomicPtr<()> = AtomicPtr::new(ptr::null_mut());
  90
  91 /// Registers a custom allocation error hook, replacing any that was previously registered.
  92 ///
  93 /// The allocation error hook is invoked when an infallible memory allocation fails, before
  94 /// the runtime aborts. The default hook prints a message to standard error,
  95 /// but this behavior can be customized with the [`set_alloc_error_hook`] and
  96 /// [`take_alloc_error_hook`] functions.
  97 ///
  98 /// The hook is provided with a `Layout` struct which contains information
  99 /// about the allocation that failed.
 100 ///
 101 /// The allocation error hook is a global resource.
 102 #[unstable(feature = "alloc_error_hook", issue = "51245")]
 103 pub fn set_alloc_error_hook(hook: fn(Layout)) {
 104     HOOK.store(hook as *mut (), Ordering::SeqCst);
 105 }
 106
 107 /// Unregisters the current allocation error hook, returning it.
 108 ///
 109 /// *See also the function [`set_alloc_error_hook`].*
 110 ///
 111 /// If no custom hook is registered, the default hook will be returned.
 112 #[unstable(feature = "alloc_error_hook", issue = "51245")]
 113 pub fn take_alloc_error_hook() -> fn(Layout) {
 114     let hook = HOOK.swap(ptr::null_mut(), Ordering::SeqCst);
 115     if hook.is_null() {
 116         default_alloc_error_hook
 117     } else {
 118         unsafe { mem::transmute(hook) }
 119     }
 120 }
 121
 122 fn default_alloc_error_hook(layout: Layout) {
 123     dumb_print(format_args!("memory allocation of {} bytes failed", layout.size()));
 124 }
 125
 126 #[cfg(not(test))]
 127 #[doc(hidden)]
 128 #[lang = "oom"]
 129 #[unstable(feature = "alloc_internals", issue = "0")]
 130 pub extern fn rust_oom(layout: Layout) -> ! {
 131     let hook = HOOK.load(Ordering::SeqCst);
 132     let hook: fn(Layout) = if hook.is_null() {
 133         default_alloc_error_hook
 134     } else {
 135         unsafe { mem::transmute(hook) }
 136     };
 137     hook(layout);
 138     unsafe { ::sys::abort_internal(); }
 139 }
 140
 141 #[cfg(not(test))]
 142 #[doc(hidden)]
 143 #[allow(unused_attributes)]
 144 #[unstable(feature = "alloc_internals", issue = "0")]
 145 pub mod __default_lib_allocator {
 146     use super::{System, Layout, GlobalAlloc};
 147     // for symbol names src/librustc/middle/allocator.rs
 148     // for signatures src/librustc_allocator/lib.rs
 149
 150     // linkage directives are provided as part of the current compiler allocator
 151     // ABI
 152
 153     #[no_mangle]
 154     #[rustc_std_internal_symbol]
 155     pub unsafe extern fn __rdl_alloc(size: usize, align: usize) -> *mut u8 {
 156         let layout = Layout::from_size_align_unchecked(size, align);
 157         System.alloc(layout)
 158     }
 159
 160     #[no_mangle]
 161     #[rustc_std_internal_symbol]
 162     pub unsafe extern fn __rdl_dealloc(ptr: *mut u8,
 163                                        size: usize,
 164                                        align: usize) {
 165         System.dealloc(ptr, Layout::from_size_align_unchecked(size, align))
 166     }
 167
 168     #[no_mangle]
 169     #[rustc_std_internal_symbol]
 170     pub unsafe extern fn __rdl_realloc(ptr: *mut u8,
 171                                        old_size: usize,
 172                                        align: usize,
 173                                        new_size: usize) -> *mut u8 {
 174         let old_layout = Layout::from_size_align_unchecked(old_size, align);
 175         System.realloc(ptr, old_layout, new_size)
 176     }
 177
 178     #[no_mangle]
 179     #[rustc_std_internal_symbol]
 180     pub unsafe extern fn __rdl_alloc_zeroed(size: usize, align: usize) -> *mut u8 {
 181         let layout = Layout::from_size_align_unchecked(size, align);
 182         System.alloc_zeroed(layout)
 183     }
 184 }