rs_malloc_tracker 1.0.1

use std::{
    env,
    fs::{create_dir, remove_file},
    io::{self, Write},
    os::unix::net::UnixListener,
    path::PathBuf,
    sync::{
        Arc, LazyLock, Mutex, RwLock,
        mpsc::{RecvTimeoutError, Sender, channel},
    },
    thread::{JoinHandle, spawn},
    time::Duration,
};

use ctor::{ctor, dtor};
use log::*;
use prometheus::{
    Encoder, IntCounter, IntCounterVec, IntGauge, IntGaugeVec, Opts, Registry, TextEncoder,
};

use crate::prometheus::{malloc::MallocState, meta_instrumentation::MetaInstrumentationState};

#[cfg(feature = "mmap")]
use crate::prometheus::mmap::MmapState;

mod malloc;
mod meta_instrumentation;
#[cfg(feature = "mmap")]
mod mmap;
mod proc_maps;

#[allow(dead_code)]
trait MyRegistry {
    fn register_int_counter(&self, name: &str, help: &str) -> Result<IntCounter, String>;
    fn register_int_counter_vec(
        &self,
        name: &str,
        help: &str,
        labels: &[&str],
    ) -> Result<IntCounterVec, String>;
    fn register_int_gauge(&self, name: &str, help: &str) -> Result<IntGauge, String>;
    fn register_int_gauge_vec(
        &self,
        name: &str,
        help: &str,
        labels: &[&str],
    ) -> Result<IntGaugeVec, String>;
}

impl MyRegistry for Registry {
    fn register_int_counter(&self, name: &str, help: &str) -> Result<IntCounter, String> {
        let metric = IntCounter::new(name, help).map_err(|e| format!("Invalid metric: {}", e))?;
        self.register(Box::new(metric.clone()))
            .map_err(|e| format!("Could not register metric: {}", e))?;

        Ok(metric)
    }

    fn register_int_counter_vec(
        &self,
        name: &str,
        help: &str,
        labels: &[&str],
    ) -> Result<IntCounterVec, String> {
        let metric = IntCounterVec::new(Opts::new(name, help), labels)
            .map_err(|e| format!("Invalid metric: {}", e))?;
        self.register(Box::new(metric.clone()))
            .map_err(|e| format!("Could not register metric: {}", e))?;

        Ok(metric)
    }

    fn register_int_gauge(&self, name: &str, help: &str) -> Result<IntGauge, String> {
        let metric = IntGauge::new(name, help).map_err(|e| format!("Invalid metric: {}", e))?;
        self.register(Box::new(metric.clone()))
            .map_err(|e| format!("Could not register metric: {}", e))?;

        Ok(metric)
    }

    fn register_int_gauge_vec(
        &self,
        name: &str,
        help: &str,
        labels: &[&str],
    ) -> Result<IntGaugeVec, String> {
        let metric = IntGaugeVec::new(Opts::new(name, help), labels)
            .map_err(|e| format!("Invalid metric: {}", e))?;
        self.register(Box::new(metric.clone()))
            .map_err(|e| format!("Could not register metric: {}", e))?;

        Ok(metric)
    }
}

trait SubmoduleState {
    /// Attempt to instanciate the submodule.
    fn new(module: &PrometheusModule) -> Result<Arc<RwLock<Self>>, String>
    where
        Self: Sized;

    fn get_registry(&self) -> &Registry;

    /// If the submodule uses threads that borrow `self`, we FIRST need to use the read mutex to
    /// tell the threads to shut down. Only when those threads are shut down is the last read lock
    /// released which allows us to use the write lock to join the threads.
    #[allow(dead_code)] // In the future we could support unloading at runtime
    fn stop(&self) -> Result<(), String> {
        Ok(())
    }

    #[allow(dead_code)] // In the future we could support unloading at runtime
    fn destroy(&mut self) -> Result<(), String> {
        Ok(())
    }
}

struct PrometheusModule {
    server_thread: Option<(JoinHandle<()>, Sender<()>)>,

    /// Used to ensure we are not doing a gauge reset while serving results, which could cause
    /// data artifacts.
    /// This also makes sure not all those expensive threads run concurrently.
    pub handler_mutex: LazyLock<Arc<Mutex<()>>>,

    pub meta_instrumentation: Option<Arc<RwLock<MetaInstrumentationState>>>,
    pub malloc: Option<Arc<RwLock<MallocState>>>,
    #[cfg(feature = "mmap")]
    pub mmap: Option<Arc<RwLock<MmapState>>>,
}

impl PrometheusModule {
    fn encode_metrics(&self) -> Result<String, String> {
        let _guard = self.handler_mutex.lock().unwrap();

        let mut buffer = vec![];
        let encoder = TextEncoder::new();

        encoder
            .encode(&prometheus::gather(), &mut buffer)
            .map_err(|e| e.to_string())?;

        encoder
            .encode(
                &self
                    .malloc
                    .as_ref()
                    .unwrap()
                    .read()
                    .unwrap()
                    .get_registry()
                    .gather(),
                &mut buffer,
            )
            .map_err(|e| e.to_string())?;

        #[cfg(feature = "mmap")]
        encoder
            .encode(
                &self
                    .mmap
                    .as_ref()
                    .unwrap()
                    .read()
                    .unwrap()
                    .get_registry()
                    .gather(),
                &mut buffer,
            )
            .map_err(|e| e.to_string())?;

        Ok(String::from_utf8_lossy(buffer.as_slice()).to_string())
    }

    fn new() -> Result<Arc<RwLock<Self>>, String> {
        let module = Arc::new(RwLock::new(Self {
            server_thread: None,
            handler_mutex: LazyLock::new(|| Arc::new(Mutex::new(()))),
            meta_instrumentation: None,
            malloc: None,
            #[cfg(feature = "mmap")]
            mmap: None,
        }));

        let meta_instrumentation = MetaInstrumentationState::new(module.read().as_ref().unwrap())?;
        module
            .write()
            .unwrap()
            .meta_instrumentation
            .replace(meta_instrumentation.clone());

        let malloc = MallocState::new(module.read().as_ref().unwrap())?;
        module.write().unwrap().malloc.replace(malloc);
        debug!("Initialized malloc sub-module");

        #[cfg(feature = "mmap")]
        {
            let mmap = MmapState::new(module.read().as_ref().unwrap())?;
            module.write().unwrap().mmap.replace(mmap);
            debug!("Initialized mmap sub-module");
        }

        let (tx, rx) = channel();
        let module2 = module.clone();
        let h = spawn(move || {
            let socket = UnixListener::bind(&*SOCKET_PATH).unwrap();
            socket.set_nonblocking(true).unwrap();
            for stream in socket.incoming() {
                match stream {
                    Ok(mut stream) => {
                        stream
                            .write(
                                module2
                                    .read()
                                    .unwrap()
                                    .encode_metrics()
                                    .unwrap_or_else(|e| format!("# Error computing metrics: {}", e))
                                    .as_bytes(),
                            )
                            .inspect_err(|e| {
                                error!("Failed to write to {:?}: {}", stream.peer_addr(), e)
                            })
                            .unwrap_or(0);
                    }
                    Err(e) => match e.kind() {
                        io::ErrorKind::WouldBlock => {
                            let res = rx.recv_timeout(Duration::from_millis(50));
                            if res.is_ok()
                                || res.is_err_and(|e| e == RecvTimeoutError::Disconnected)
                            {
                                break;
                            }
                        }
                        _ => error!("Socket connection failed: {}", e),
                    },
                }
            }
            remove_file(&*SOCKET_PATH).unwrap();
        });
        module.write().unwrap().server_thread.replace((h, tx));

        Ok(module)
    }
}

impl Drop for PrometheusModule {
    fn drop(&mut self) {
        debug!("Unloading...");

        self.malloc
            .as_ref()
            .unwrap()
            .read()
            .unwrap()
            .stop()
            .unwrap();
        self.malloc
            .as_mut()
            .unwrap()
            .write()
            .unwrap()
            .destroy()
            .unwrap();
        debug!("malloc sub-module unloaded");

        #[cfg(feature = "mmap")]
        {
            self.mmap.as_ref().unwrap().read().unwrap().stop().unwrap();
            self.mmap
                .as_mut()
                .unwrap()
                .write()
                .unwrap()
                .destroy()
                .unwrap();
            debug!("mmap sub-module unloaded");
        }

        if let Some((h, s)) = self.server_thread.take() {
            s.send(()).unwrap();
            h.join().unwrap();
            debug!("HTTP Server thread joined");
        }

        debug!("Unloaded prometheus");
    }
}

static SOCKET_PATH: LazyLock<PathBuf> = LazyLock::new(|| {
    let mut socket_path = PathBuf::from(env::var("XDG_RUNTIME_DIR").unwrap_or("/run".to_string()));
    socket_path.push(env!("CARGO_PKG_NAME"));
    create_dir(&socket_path).unwrap_or(()); // ignore error because path may already
    // exist (will fail on bind otherwise)
    socket_path.push(format!("{}.sock", std::process::id()));
    socket_path
});
static mut PROMETHEUS_MODULE: *mut Arc<RwLock<PrometheusModule>> = std::ptr::null_mut();

#[ctor]
fn prometheus_init() {
    let pid = std::process::id();
    env_logger::builder()
        .format(move |buf, record| {
            writeln!(
                buf,
                "[{} {} {} {}]: {}",
                pid,
                buf.timestamp(),
                record.level(),
                record.module_path().unwrap_or("unknown"),
                record.args()
            )
        })
        .init();

    if unsafe { PROMETHEUS_MODULE.is_null() } {
        let m = Box::new(PrometheusModule::new().unwrap());
        unsafe {
            PROMETHEUS_MODULE = Box::into_raw(m);
        }
    }
}

#[dtor]
fn prometheus_deinit() {
    if let Some(ptr) = unsafe { PROMETHEUS_MODULE.as_mut() } {
        let m = unsafe { Box::from_raw(ptr) };
        drop(m);
    }

    // It is possible that the remove_file call has not happened in case the server thread exited
    // early
    remove_file(&*SOCKET_PATH)
        .inspect_err(|e| {
            error!(
                "Could not detelete {}: {}",
                SOCKET_PATH.to_str().unwrap(),
                e
            )
        })
        .unwrap_or(());
}