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::consumer::{MEMORY_ZONES, SEND_FAILS, ZoneKind};

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

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

struct Metrics {
    registry: Registry,

    allocated_bytes: IntGaugeVec,
    allocated_regions: IntGaugeVec,
    allocated_regions_count: IntGauge,
    send_fails_count: IntGaugeVec,
    usable_allocated_bytes: IntGaugeVec,
}

pub struct MallocState {
    config: Config,
    metrics: Metrics,

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

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

impl MallocState {
    ///
    /// The operation being done here (iterating over every allocated memory zone) is very expensive
    /// (several seconds of compute).
    /// 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!("Malloc 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();

                // Clone the zones so we don't keep the read lock for longer than necessary
                // (otherwise sent zones will pile up and performance will degrade!)
                let zones = self.meta_instrumentation.read().unwrap().meta_instrument(
                    "malloc-clone",
                    || -> Result<_, String> {
                        Ok(MEMORY_ZONES
                            .read()
                            .map_err(|_| "Failed to acquire MEMORY_ZONES read lock".to_string())?
                            .clone())
                    },
                )?;

                self.metrics.allocated_bytes.reset();
                self.metrics.allocated_regions.reset();
                self.metrics.usable_allocated_bytes.reset();
                self.metrics.allocated_regions_count.set(zones.len() as i64);

                let send_fails = unsafe { &*&raw const SEND_FAILS }.read().unwrap();
                for (error, count) in send_fails.iter() {
                    self.metrics
                        .send_fails_count
                        .with_label_values(&[error])
                        .add(*count as i64);
                }

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

                self.meta_instrumentation.read().unwrap().meta_instrument(
                    "malloc-compute",
                    || -> Result<_, String> {
                        for zone in zones.values() {
                            let sym = zone
                                .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>");
                            let kind = match zone.kind {
                                ZoneKind::Malloc => "malloc",
                                ZoneKind::Calloc => "calloc",
                                ZoneKind::Memalign => "memalign",
                                ZoneKind::Realloc => "realloc",
                            };

                            self.metrics
                                .allocated_regions
                                .with_label_values(&[kind, map_region, function])
                                .inc();
                            self.metrics
                                .allocated_bytes
                                .with_label_values(&[kind, map_region, function])
                                .add(zone.size as i64);

                            self.metrics
                                .usable_allocated_bytes
                                .with_label_values(&[kind, map_region, function])
                                .add(zone.usable_size as i64);
                        }

                        Ok(())
                    },
                )?;

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

            debug!("Processed malloc 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 malloc 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 MallocState {
    fn new(module: &super::PrometheusModule) -> Result<Arc<RwLock<Self>>, String>
    where
        Self: Sized,
    {
        let registry = Registry::new();

        let state = Arc::new(RwLock::new(MallocState {
            config: Config {
                functions: env::var("RS_MALLOC_TRACKER_MALLOC_FUNCTIONS")
                    .unwrap_or("no".to_string())
                    == "yes",
                interval: Duration::from_secs(
                    env::var("RS_MALLOC_TRACKER_MALLOC_PROCESSING_INTERVAL_SECONDS")
                        .unwrap_or("30".to_string())
                        .parse::<u64>()
                        .map_err(|e| format!("Invalid interval string: {}", e))?,
                ),
            },
            meta_instrumentation: Arc::clone(module.meta_instrumentation.as_ref().unwrap()),
            handler_mutex: Arc::clone(&module.handler_mutex),
            symbol_cache: Mutex::new(SymbolCache::new(module)?),
            metrics: Metrics {
                allocated_bytes: registry.register_int_gauge_vec(
                    "rs_malloc_tracker_malloc_allocated_bytes_count",
                    "Amount of allocated memory via malloc",
                    &["kind", "region", "function"],
                )?,
                allocated_regions: registry.register_int_gauge_vec(
                    "rs_malloc_tracker_malloc_allocated_regions_count",
                    "Number of allocated memory regions via malloc",
                    &["kind", "region", "function"],
                )?,
                allocated_regions_count: registry.register_int_gauge(
                    "rs_malloc_tracker_malloc_tracked_regions_count",
                    "Number of tracked memory regions",
                )?,
                send_fails_count: registry.register_int_gauge_vec(
                    "rs_malloc_tracker_malloc_send_fails",
                    "Regions failed to be registered",
                    &["error"],
                )?,
                usable_allocated_bytes: registry.register_int_gauge_vec(
                    "rs_malloc_tracker_malloc_usable_allocated_bytes_count",
                    "Amount of allocated memory via malloc (usable as reported by jemalloc)",
                    &["kind", "region", "function"],
                )?,
                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!("malloc sub-module stopped");

        Ok(())
    }

    fn destroy(&mut self) -> Result<(), String> {
        self.processing_thread
            .take()
            .unwrap()
            .join()
            .map_err(|_| "Failed to join processing thread")?;

        Ok(())
    }
}