rs_malloc_tracker 1.0.1

#![allow(clippy::deref_addrof)] // necessary for static muts

use std::env;
use std::sync::mpsc::{Receiver, RecvTimeoutError, Sender, channel};
use std::sync::{Arc, Mutex, RwLock};
use std::thread::{self, JoinHandle};
use std::time::{Duration, Instant};

use log::*;
use prometheus::{IntGauge, IntGaugeVec, Registry};

use crate::mmap::{MMAP_SEND_FAILS, MMAP_SEND_MISSES, MMAP_ZONES};

use super::meta_instrumentation::MetaInstrumentationState;
use super::proc_maps::{Symbol, SymbolCache, read_self_smaps};
use super::{MyRegistry, SubmoduleState};

#[allow(dead_code)]
#[derive(Clone, Debug)]
pub struct Mmap {
    pub flags: i32,
    pub ip: Option<usize>,
}

struct Config {
    pub functions: bool,
    pub interval: Duration,
}

struct Metrics {
    registry: Registry,

    mmap_virt_bytes: IntGaugeVec,
    smaps_rss_bytes: IntGaugeVec,
    smaps_virt_bytes: IntGaugeVec,
    mmap_send_fails_count: IntGauge,
    mmap_send_misses_count: IntGauge,
}

pub struct MmapState {
    config: Config,

    metrics: Metrics,

    handler_mutex: Arc<Mutex<()>>,
    symbol_cache: Mutex<SymbolCache>,
    meta_instrumentation: Arc<RwLock<MetaInstrumentationState>>,

    processing_thread: Option<JoinHandle<()>>,
    processing_shutdown: Mutex<Option<Sender<()>>>,
}

impl MmapState {
    ///
    /// The operation being done here (iterating over every allocated memory zone) can be very expensive
    /// (several seconds of compute). Though not as bad as malloc or ast_malloc.
    /// Therefore we do not want to run it every time we serve a prometheus request.
    ///
    /// Instead, we recompute them in a background thread with a dynamic interval to ensure that we do
    /// not generate too much load.
    fn process(&self, started: Sender<()>, shutdown: Receiver<()>) {
        info!("Mmap process thread started");
        started.send(()).unwrap();
        loop {
            let elapsed = || -> Result<Duration, String> {
                // Hold off until request is served since we do a reset()!
                let _guard = self.handler_mutex.lock().unwrap();

                let start_time = Instant::now();

                self.metrics.mmap_virt_bytes.reset();
                self.metrics.smaps_rss_bytes.reset();
                self.metrics.smaps_virt_bytes.reset();
                self.metrics
                    .mmap_send_fails_count
                    .set(unsafe { *(*(&raw const MMAP_SEND_FAILS)).read().unwrap() as i64 });
                self.metrics
                    .mmap_send_misses_count
                    .set(unsafe { *(*(&raw const MMAP_SEND_MISSES)).read().unwrap() as i64 });

                // Clone MMAP_ZONES to avoid locking those unnecessarily long while we process
                let mmap_zones = self.meta_instrumentation.read().unwrap().meta_instrument(
                    "mmaps-clone",
                    || -> Result<_, String> {
                        Ok(MMAP_ZONES
                            .read()
                            .map_err(|_| "Failed to acquire MMAP_ZONES read lock".to_string())?
                            .clone())
                    },
                )?;

                let mut symbol_cache = self.symbol_cache.lock().unwrap();

                // For each memory zone, instrument virtual memory
                // (this is straightforward, unlike RSS)
                self.meta_instrumentation.read().unwrap().meta_instrument(
                    "mmaps-compute",
                    || -> Result<_, String> {
                        for (ival, map) in mmap_zones.iter() {
                            let sym = map
                                .ip
                                .map(|ip| symbol_cache.resolve_symbol_at(self.config.functions, ip))
                                .unwrap_or(Ok(Symbol::default()))?;
                            let map_region = sym.region.map(|s| s.as_str()).unwrap_or("<unknown>");
                            let function = sym.function.map(|s| s.as_str()).unwrap_or("<unknown>");

                            self.metrics
                                .mmap_virt_bytes
                                .with_label_values(&[map_region, function])
                                .add((ival.end() + 1 - ival.start()) as i64);
                        }

                        Ok(())
                    },
                )?;

                // Now the fun part begins.
                //
                // Instrumenting RSS is non-trivial because for each zone we only know its address
                // and length.
                // However the kernel knows the RSS and exposes it in /proc/self/smaps.
                // Therefore for each zone of virtual memory we have instrumented, we need to check
                // /proc/self/smaps to see how much RSS has been allocated for that memory region.
                let smaps = self
                    .meta_instrumentation
                    .read()
                    .unwrap()
                    .meta_instrument("read-proc-smaps", read_self_smaps)?;

                self.meta_instrumentation.read().unwrap().meta_instrument(
                    "smaps-compute",
                    || -> Result<_, String> {
                        // Iterate over smaps
                        for (ival, map) in smaps.iter() {
                            // For each region in /proc/self/smaps, this is the decision tree
                            //
                            // 1. We have an instrumented mmap() region for this address.
                            //    Great! We can count the smaps RSS towards that mmap()'d region.
                            // 2. ELSE:
                            //    a. The memory region has a name (via PR_SET_VMA), which we can report
                            //    b. The memory region does not have a name, in which case we have no
                            //       way to know who allocated it.
                            //
                            //       This last case should happen only if memory was allocated
                            //       without going through our preloaded `mmap()` function.
                            //       I have noticed one example is the PLT which is loaded by
                            //       linux-vdso which cannot be forced to use our injectged mmap.
                            let sym = if let Some(mmap) = mmap_zones.get_at_point(ival.start()) {
                                // we know the zone!
                                mmap.ip
                                    .map(|ip| {
                                        symbol_cache.resolve_symbol_at(self.config.functions, ip)
                                    })
                                    .unwrap_or(Ok(Symbol::default()))?
                            } else if let Some(owner) = map.owner.as_ref() {
                                // we don't know the zone but it has a name
                                Symbol {
                                    region: Some(owner),
                                    function: None,
                                }
                            } else {
                                // :(
                                Symbol {
                                    region: Some(&"<untracked>".to_string()),
                                    function: Some(&"<untracked>".to_string()),
                                }
                            };
                            let map_region = sym.region.map(|s| s.as_str()).unwrap_or("<unknown>");
                            let function = sym.function.map(|s| s.as_str()).unwrap_or("<unknown>");

                            self.metrics
                                .smaps_rss_bytes
                                .with_label_values(&[map_region, function])
                                .add(map.rss as i64);
                            self.metrics
                                .smaps_virt_bytes
                                .with_label_values(&[map_region, function])
                                .add(map.size as i64);
                        }

                        Ok(())
                    },
                )?;

                Ok(start_time.elapsed())
            }()
            .inspect_err(|e| error!("Error processing mmap: {}", e))
            .unwrap_or(Duration::default());

            debug!("Processed mmap in {} ms", elapsed.as_millis());
            let sleep = if elapsed > self.config.interval.div_f32(10.0) {
                let backoff = self.config.interval.mul_f32(elapsed.as_secs() as f32);
                info!(
                    "Processing mmap took {} ms. Backing off with a {} ms sleep.",
                    elapsed.as_millis(),
                    backoff.as_millis()
                );
                backoff
            } else {
                self.config.interval
            };
            // exit sleep immediately on shutdown
            let res = shutdown.recv_timeout(sleep);
            if res.is_ok() || res.is_err_and(|e| e == RecvTimeoutError::Disconnected) {
                break;
            }
        }
    }
}

impl SubmoduleState for MmapState {
    fn new(module: &super::PrometheusModule) -> Result<Arc<RwLock<Self>>, String>
    where
        Self: Sized,
    {
        let registry = Registry::new();

        let state = Arc::new(RwLock::new(MmapState {
            meta_instrumentation: Arc::clone(module.meta_instrumentation.as_ref().unwrap()),
            handler_mutex: Arc::clone(&module.handler_mutex),
            symbol_cache: Mutex::new(SymbolCache::new(module)?),
            config: Config {
                functions: env::var("RS_MALLOC_TRACKER_MMAP_FUNCTIONS")
                    .unwrap_or("yes".to_string())
                    == "yes",
                interval: Duration::from_secs(
                    env::var("RS_MALLOC_TRACKER_MMAP_PROCESSING_INTERVAL_SECONDS")
                        .unwrap_or("30".to_string())
                        .parse::<u64>()
                        .map_err(|e| format!("Invalid interval string: {}", e))?,
                ),
            },
            metrics: Metrics {
                mmap_virt_bytes: registry.register_int_gauge_vec(
                    "rs_malloc_tracker_mmap_virt_bytes_count",
                    "Amount of allocated virtal memory via mmap",
                    &["region", "function"],
                )?,
                smaps_rss_bytes: registry.register_int_gauge_vec(
                    "rs_malloc_tracker_smaps_rss_bytes_count",
                    "Amount of allocated resident memory via mmap (checked via /proc/self/smaps)",
                    &["region", "function"],
                )?,
                smaps_virt_bytes: registry.register_int_gauge_vec(
                    "rs_malloc_tracker_smaps_virt_bytes_count",
                    "Amount of virtual memory via mmap (checked via /proc/self/smaps)",
                    &["region", "function"],
                )?,
                mmap_send_fails_count: registry.register_int_gauge(
                    "rs_malloc_tracker_mmap_send_fails",
                    "Mmap Regions failed to be registered",
                )?,
                mmap_send_misses_count: registry.register_int_gauge(
                    "rs_malloc_tracker_mmap_send_misses",
                    "Mmap Regions could not be registered at the beginning of program execution",
                )?,
                registry,
            },

            processing_shutdown: Mutex::default(),
            processing_thread: None,
        }));

        let (processing_shutdown_sender, processing_shutdown_receiver) = channel();
        let processing_state = Arc::clone(&state);
        state
            .write()
            .unwrap()
            .processing_shutdown
            .lock()
            .map_err(|_| "Failed to acquire processing_shutdown mutex")?
            .replace(processing_shutdown_sender);

        let (processing_started_sender, processing_started_receiver) = channel();
        state
            .write()
            .unwrap()
            .processing_thread
            .replace(thread::spawn(move || {
                processing_state
                    .read()
                    .unwrap()
                    .process(processing_started_sender, processing_shutdown_receiver)
            }));

        // Necessary to avoid deadlock if shutdown was initiated before thread could acquire read
        // lock.
        processing_started_receiver.recv().unwrap();

        Ok(state)
    }

    fn get_registry(&self) -> &Registry {
        &self.metrics.registry
    }

    fn stop(&self) -> Result<(), String> {
        self.processing_shutdown
            .lock()
            .map_err(|_| "Failed to acquire processing_shutdown mutex")?
            .as_ref()
            .unwrap()
            .send(())
            .unwrap();

        debug!("mmap sub-module stopped");

        Ok(())
    }

    fn destroy(&mut self) -> Result<(), String> {
        debug!("joining mmap sub-module");

        self.processing_thread
            .take()
            .unwrap()
            .join()
            .map_err(|_| "Failed to join processing thread")?;

        debug!("mmap sub-module joined");

        Ok(())
    }
}