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

//! A Rust allocator which makes sound when active, like a Geiger counter.
//!
//! Currently this just writes an ASCII [`BEL`] to `/dev/tty`.
//!
//! Inspired by [Malloc Geiger].
//!
//! [`BEL`]: https://en.wikipedia.org/wiki/Bell_character
//! [Malloc Geiger]: https://github.com/laserallan/malloc_geiger
//!
//! ## Usage
//!
//! To use `alloc_geiger` add it as a dependency:
//!
//! ```toml
//! # Cargo.toml
//! [dependencies]
//! alloc_geiger = "0.1.0"
//! ```
//!
//! To set `alloc_geiger::Geiger` as the global allocator, it must be initialized
//! with an underlying allocator. The `type System` alias and its `const SYSTEM`
//! make it easy to use the default system allocator:
//!
//! ```rust
//! #[global_allocator]
//! static ALLOC: alloc_geiger::System = alloc_geiger::SYSTEM;
//!
//! fn main() {
//!     // ...
//! }
//! ```
//!
//! Alternatives like [`jemallocator`] may also be used:
//!
//! ```rust
//! use alloc_geiger::Geiger;
//! use jemallocator::Jemalloc;
//!
//! #[global_allocator]
//! static ALLOC: Geiger<Jemalloc> = Geiger::new(Jemalloc);
//!
//! fn main() {
//!     // ...
//! }
//! ```
//!
//! [`jemallocator`]: https://crates.io/crates/jemallocator

use once_cell::sync::OnceCell;
use std::alloc::{self, GlobalAlloc, Layout};
use std::cell::Cell;
use std::fs::{File, OpenOptions};
use std::io::Write;

/// Geiger counter allocator.
#[derive(Debug, Default)]
pub struct Geiger<Alloc> {
    inner: Alloc,
    tty: OnceCell<Option<File>>,
}

/// `Geiger` allocator based on `std::alloc::System`.
pub type System = Geiger<alloc::System>;

/// `Geiger` allocator based on `std::alloc::System`.
pub const SYSTEM: System = Geiger::new(alloc::System);

fn open_tty() -> Option<File> {
    OpenOptions::new().append(true).open("/dev/tty").ok()
}

impl<Alloc> Geiger<Alloc> {
    pub const fn new(inner: Alloc) -> Self {
        Geiger {
            inner,
            tty: OnceCell::new(),
        }
    }

    fn bell(&self) {
        const BEL: u8 = 0x07;

        thread_local! {
            // Guard against recursion
            static BUSY: Cell<bool> = Cell::new(false);
        }

        BUSY.with(|busy| {
            if !busy.replace(true) {
                let tty = self.tty.get_or_init(open_tty);
                if let Some(ref mut file) = tty.as_ref() {
                    file.write_all(&[BEL]).ok();
                }
                busy.set(false);
            }
        });
    }
}

unsafe impl<Alloc: GlobalAlloc> GlobalAlloc for Geiger<Alloc> {
    #[inline]
    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
        self.bell();
        self.inner.alloc(layout)
    }

    #[inline]
    unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
        self.bell();
        self.inner.alloc_zeroed(layout)
    }

    #[inline]
    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
        self.bell();
        self.inner.dealloc(ptr, layout)
    }

    #[inline]
    unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
        self.bell();
        self.inner.realloc(ptr, layout, new_size)
    }
}