sqlite_graphrag/
memory_guard.rs

1//! Guarda de memória: verifica disponibilidade de RAM antes de carregar o modelo ONNX.
2//!
3//! O carregamento do modelo via `fastembed` consome aproximadamente
4//! [`crate::constants::EMBEDDING_LOAD_EXPECTED_RSS_MB`] MiB de memória residente.
5//! Se o sistema não tiver memória suficiente disponível, múltiplas invocações
6//! paralelas podem esgotar a RAM e causar OOM (Out-Of-Memory), travando o sistema.
7//!
8//! Esta guard interroga o SO via `sysinfo` antes de qualquer inicialização pesada,
9//! abortando com [`crate::errors::AppError::LowMemory`] (exit 77) quando o piso
10//! configurado não é atingido.
11
12use sysinfo::{
13    get_current_pid, MemoryRefreshKind, ProcessRefreshKind, ProcessesToUpdate, RefreshKind, System,
14    UpdateKind,
15};
16
17use crate::errors::AppError;
18
19/// Retorna a memória disponível atual em MiB.
20pub fn available_memory_mb() -> u64 {
21    let sys =
22        System::new_with_specifics(RefreshKind::new().with_memory(MemoryRefreshKind::everything()));
23    let available_bytes = sys.available_memory();
24    available_bytes / (1024 * 1024)
25}
26
27/// Retorna o RSS atual do processo em MiB quando disponível.
28pub fn current_process_memory_mb() -> Option<u64> {
29    let pid = get_current_pid().ok()?;
30    let mut sys =
31        System::new_with_specifics(RefreshKind::new().with_memory(MemoryRefreshKind::everything()));
32    sys.refresh_processes_specifics(
33        ProcessesToUpdate::Some(&[pid]),
34        true,
35        ProcessRefreshKind::new()
36            .with_memory()
37            .with_exe(UpdateKind::OnlyIfNotSet),
38    );
39    sys.process(pid).map(|p| p.memory() / (1024 * 1024))
40}
41
42/// Calcula o teto seguro de concorrência para cargas pesadas de embedding.
43///
44/// Fórmula canônica:
45/// `permits = min(cpus, available_memory_mb / ram_por_task_mb) * 0.5`
46///
47/// O resultado é clampado entre `1` e `max_concurrency`.
48pub fn calculate_safe_concurrency(
49    available_mb: u64,
50    cpu_count: usize,
51    ram_per_task_mb: u64,
52    max_concurrency: usize,
53) -> usize {
54    let cpu_count = cpu_count.max(1);
55    let max_concurrency = max_concurrency.max(1);
56    let ram_per_task_mb = ram_per_task_mb.max(1);
57
58    let memory_bound = (available_mb / ram_per_task_mb) as usize;
59    let resource_bound = cpu_count.min(memory_bound).max(1);
60    let safe_with_margin = (resource_bound / 2).max(1);
61
62    safe_with_margin.min(max_concurrency)
63}
64
65/// Verifica se há memória disponível suficiente para iniciar o carregamento do modelo.
66///
67/// # Parâmetros
68/// - `min_mb`: piso mínimo em MiB de memória disponível (tipicamente
69///   [`crate::constants::MIN_AVAILABLE_MEMORY_MB`]).
70///
71/// # Erros
72/// Retorna [`AppError::LowMemory`] quando `available_mb < min_mb`.
73///
74/// # Retorno
75/// Retorna `Ok(available_mb)` com o valor real de memória disponível em MiB.
76pub fn check_available_memory(min_mb: u64) -> Result<u64, AppError> {
77    let available_mb = available_memory_mb();
78
79    if available_mb < min_mb {
80        return Err(AppError::LowMemory {
81            available_mb,
82            required_mb: min_mb,
83        });
84    }
85
86    Ok(available_mb)
87}
88
89#[cfg(test)]
90mod testes {
91    use super::*;
92
93    #[test]
94    fn check_available_memory_com_zero_sempre_passa() {
95        let resultado = check_available_memory(0);
96        assert!(
97            resultado.is_ok(),
98            "min_mb=0 deve sempre passar, got: {resultado:?}"
99        );
100        let mb = resultado.unwrap();
101        assert!(mb > 0, "sistema deve reportar memória positiva");
102    }
103
104    #[test]
105    fn check_available_memory_com_valor_gigante_falha() {
106        let resultado = check_available_memory(u64::MAX);
107        assert!(
108            matches!(resultado, Err(AppError::LowMemory { .. })),
109            "u64::MAX MiB deve falhar com LowMemory, got: {resultado:?}"
110        );
111    }
112
113    #[test]
114    fn low_memory_error_contem_valores_corretos() {
115        match check_available_memory(u64::MAX) {
116            Err(AppError::LowMemory {
117                available_mb,
118                required_mb,
119            }) => {
120                assert_eq!(required_mb, u64::MAX);
121                assert!(available_mb < u64::MAX);
122            }
123            outro => panic!("esperado LowMemory, got: {outro:?}"),
124        }
125    }
126
127    #[test]
128    fn calculate_safe_concurrency_respeita_metade_da_margem() {
129        let permits = calculate_safe_concurrency(8_000, 8, 1_000, 4);
130        assert_eq!(permits, 4);
131    }
132
133    #[test]
134    fn calculate_safe_concurrency_nunca_retorna_zero() {
135        let permits = calculate_safe_concurrency(100, 1, 10_000, 4);
136        assert_eq!(permits, 1);
137    }
138
139    #[test]
140    fn calculate_safe_concurrency_respeita_teto_maximo() {
141        let permits = calculate_safe_concurrency(128_000, 64, 500, 4);
142        assert_eq!(permits, 4);
143    }
144
145    #[test]
146    fn current_process_memory_mb_retorna_algum_valor() {
147        let rss = current_process_memory_mb();
148        assert!(rss.is_some());
149    }
150}
sqlite_graphrag/memory_guard.rs

sqlite_graphrag/
memory_guard.rs
