1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124

use std::{
    alloc::{GlobalAlloc, Layout, System},
    ptr,
    sync::atomic::{AtomicUsize, Ordering},
};

/// Wraps an allocator, adding memory use limits and tracking for peak
/// memory usage.
pub struct Alloc<A = System> {
    parent: A,
    used: AtomicUsize,
    max: AtomicUsize,
    limit: AtomicUsize,
}

impl Alloc<System> {
    /// Creates a new wrapper allocator using the system allocator.
    pub const fn new(limit: usize) -> Alloc<System> {
        Alloc {
            parent: System,
            used: AtomicUsize::new(0),
            max: AtomicUsize::new(0),
            limit: AtomicUsize::new(limit),
        }
    }
}

impl<A> Alloc<A>
where
    A: GlobalAlloc,
{
    /// Creates a new wrapper allocator using any GlobalAllocator.
    ///
    /// Unfortunately, this can't be a const fn, because trait bounds on const fns are unstable.
    /// <https://github.com/rust-lang/rust/issues/57563>
    pub fn new_with(parent: A, limit: usize) -> Alloc<A> {
        Alloc {
            parent,
            used: AtomicUsize::new(0),
            max: AtomicUsize::new(0),
            limit: AtomicUsize::new(limit),
        }
    }

    /// Clears the current peak value.
    pub fn reset_max(&self) {
        self.max
            .store(self.used.load(Ordering::Acquire), Ordering::Release);
    }

    /// Returns the peak memory that's been used since startup or since
    /// `reset_max()` was called.
    pub fn get_max(&self) -> usize {
        self.max.load(Ordering::Acquire)
    }

    /// Sets the maximum amount of memory that can be used. This should
    /// be called early in the application lifecycle.
    pub fn set_limit(&self, limit: usize) {
        self.limit.store(limit, Ordering::Release);
    }
}

unsafe impl GlobalAlloc for Alloc {
    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
        let size = layout.size();
        let limit = self.limit.load(Ordering::Acquire);
        let new_size = self.used.fetch_add(size, Ordering::Acquire) + size;
        if new_size <= limit {
            self.max.fetch_max(new_size, Ordering::Relaxed);
            let result = self.parent.alloc(layout);
            if result.is_null() {
                self.used.fetch_sub(size, Ordering::Release);
            }
            result
        } else {
            self.used.fetch_sub(size, Ordering::Release);
            ptr::null_mut()
        }
    }

    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
        let size = layout.size();
        self.parent.dealloc(ptr, layout);
        self.used.fetch_sub(size, Ordering::Release);
    }

    unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
        let size = layout.size();
        let limit = self.limit.load(Ordering::Acquire);
        let new_size = self.used.fetch_add(size, Ordering::Acquire) + size;
        if new_size <= limit {
            self.max.fetch_max(new_size, Ordering::Relaxed);
            let result = self.parent.alloc_zeroed(layout);
            if result.is_null() {
                self.used.fetch_sub(size, Ordering::Release);
            }
            result
        } else {
            self.used.fetch_sub(size, Ordering::Release);
            ptr::null_mut()
        }
    }

    unsafe fn realloc(&self, ptr: *mut u8, old_layout: Layout, realloc_size: usize) -> *mut u8 {
        let new_layout = Layout::from_size_align_unchecked(realloc_size, old_layout.align());
        let (old_size, new_size) = (old_layout.size(), new_layout.size());

        let limit = self.limit.load(Ordering::Acquire);
        let new_used = self.used.fetch_add(new_size, Ordering::Acquire) + new_size;
        if new_used <= limit {
            let result = self.parent.realloc(ptr, old_layout, realloc_size);
            if result.is_null() {
                self.used.fetch_sub(new_size, Ordering::Release);
            } else {
                self.used.fetch_sub(old_size, Ordering::Release);
            }
            result
        } else {
            self.used.fetch_sub(new_size, Ordering::Release);
            ptr::null_mut()
        }
    }
}