Skip to main content

hypercall_admin/monitoring/
system.rs

1//! Memory stats + heap profiling admin endpoints.
2
3use axum::{http::StatusCode, response::IntoResponse};
4use serde::Serialize;
5
6use hypercall_runtime_api::sonic_json::SonicJson;
7
8/// Response for memory stats endpoint.
9#[derive(Debug, Serialize, utoipa::ToSchema)]
10pub struct MemoryStatsResponse {
11    /// Physical memory used by the process (bytes)
12    pub physical_mem_bytes: Option<usize>,
13    /// Virtual memory used by the process (bytes)
14    pub virtual_mem_bytes: Option<usize>,
15    /// Heap profiling enabled at compile time
16    pub heap_profiling_enabled: bool,
17    /// jemalloc stats (only if heap-profiling feature enabled)
18    #[serde(skip_serializing_if = "Option::is_none")]
19    pub jemalloc: Option<JemallocStats>,
20}
21
22/// jemalloc memory statistics
23#[derive(Debug, Serialize, utoipa::ToSchema)]
24pub struct JemallocStats {
25    /// Total bytes allocated by the application
26    pub allocated: usize,
27    /// Total bytes in active pages allocated by the application
28    pub active: usize,
29    /// Total bytes dedicated to metadata
30    pub metadata: usize,
31    /// Total bytes in physically mapped data pages
32    pub mapped: usize,
33    /// Total bytes retained (not released to OS)
34    pub retained: usize,
35    /// Total bytes in resident data pages
36    pub resident: usize,
37}
38
39/// Get current memory usage statistics.
40///
41/// Returns process memory stats and jemalloc stats (if heap-profiling enabled).
42#[utoipa::path(
43    get,
44    path = "/monitoring/memory",
45    responses(
46        (status = 200, description = "Memory statistics", body = MemoryStatsResponse)
47    ),
48    tag = "Monitoring",
49    security(("admin_key" = []))
50)]
51pub async fn memory_stats() -> impl IntoResponse {
52    let mem = memory_stats::memory_stats();
53
54    let (physical_mem_bytes, virtual_mem_bytes) = match mem {
55        Some(stats) => (Some(stats.physical_mem), Some(stats.virtual_mem)),
56        None => (None, None),
57    };
58
59    #[cfg(feature = "heap-profiling")]
60    let jemalloc = {
61        use tikv_jemalloc_ctl::{epoch, stats};
62        // Advance the epoch to get fresh stats
63        epoch::advance().ok();
64
65        Some(JemallocStats {
66            allocated: stats::allocated::read().unwrap_or(0),
67            active: stats::active::read().unwrap_or(0),
68            metadata: stats::metadata::read().unwrap_or(0),
69            mapped: stats::mapped::read().unwrap_or(0),
70            retained: stats::retained::read().unwrap_or(0),
71            resident: stats::resident::read().unwrap_or(0),
72        })
73    };
74
75    #[cfg(not(feature = "heap-profiling"))]
76    let jemalloc: Option<JemallocStats> = None;
77
78    SonicJson(MemoryStatsResponse {
79        physical_mem_bytes,
80        virtual_mem_bytes,
81        heap_profiling_enabled: cfg!(feature = "heap-profiling"),
82        jemalloc,
83    })
84}
85
86/// Response for heap dump endpoint.
87#[derive(Debug, Serialize, utoipa::ToSchema)]
88pub struct HeapDumpResponse {
89    pub success: bool,
90    pub message: String,
91    #[serde(skip_serializing_if = "Option::is_none")]
92    pub dump_path: Option<String>,
93}
94
95/// Trigger a jemalloc heap profile dump.
96///
97/// Requires the binary to be built with --features heap-profiling and
98/// run with MALLOC_CONF=prof:true environment variable.
99///
100/// The dump file will be written to /tmp/heap_dump_<timestamp>.prof
101#[utoipa::path(
102    post,
103    path = "/monitoring/heap-dump",
104    responses(
105        (status = 200, description = "Heap dump triggered", body = HeapDumpResponse),
106        (status = 501, description = "Heap profiling not enabled")
107    ),
108    tag = "Monitoring",
109    security(("admin_key" = []))
110)]
111pub async fn heap_dump() -> impl IntoResponse {
112    #[cfg(feature = "heap-profiling")]
113    {
114        use std::ffi::CString;
115
116        let timestamp = chrono::Utc::now().format("%Y%m%d_%H%M%S");
117        let dump_path = format!("/tmp/heap_dump_{}.prof", timestamp);
118
119        // jemalloc requires a null-terminated C string for the dump path
120        match CString::new(dump_path.clone()) {
121            Ok(c_path) => {
122                // Use raw mallctl interface to write to "prof.dump"
123                // This triggers jemalloc to write a heap profile to the specified path
124                // The value must be a pointer to the C string (char*)
125                let path_ptr = c_path.as_ptr();
126                let result: Result<(), tikv_jemalloc_ctl::Error> =
127                    unsafe { tikv_jemalloc_ctl::raw::write(b"prof.dump\0", path_ptr) };
128
129                match result {
130                    Ok(()) => {
131                        tracing::info!("Heap dump written to {}", dump_path);
132                        (
133                            StatusCode::OK,
134                            SonicJson(HeapDumpResponse {
135                                success: true,
136                                message: format!("Heap dump written to {}", dump_path),
137                                dump_path: Some(dump_path),
138                            }),
139                        )
140                            .into_response()
141                    }
142                    Err(e) => {
143                        let msg = format!(
144                            "Failed to write heap dump: {}. Make sure MALLOC_CONF=prof:true is set.",
145                            e
146                        );
147                        tracing::error!("{}", msg);
148                        (
149                            StatusCode::INTERNAL_SERVER_ERROR,
150                            SonicJson(HeapDumpResponse {
151                                success: false,
152                                message: msg,
153                                dump_path: None,
154                            }),
155                        )
156                            .into_response()
157                    }
158                }
159            }
160            Err(e) => {
161                let msg = format!("Invalid dump path: {}", e);
162                tracing::error!("{}", msg);
163                (
164                    StatusCode::INTERNAL_SERVER_ERROR,
165                    SonicJson(HeapDumpResponse {
166                        success: false,
167                        message: msg,
168                        dump_path: None,
169                    }),
170                )
171                    .into_response()
172            }
173        }
174    }
175
176    #[cfg(not(feature = "heap-profiling"))]
177    {
178        (
179            StatusCode::NOT_IMPLEMENTED,
180            SonicJson(HeapDumpResponse {
181                success: false,
182                message: "Heap profiling not enabled. Rebuild with --features heap-profiling"
183                    .to_string(),
184                dump_path: None,
185            }),
186        )
187            .into_response()
188    }
189}