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::{MemoryRefreshKind, RefreshKind, System};
13
14use crate::errors::AppError;
15
16/// Retorna a memória disponível atual em MiB.
17pub fn available_memory_mb() -> u64 {
18    let sys =
19        System::new_with_specifics(RefreshKind::new().with_memory(MemoryRefreshKind::everything()));
20    let available_bytes = sys.available_memory();
21    available_bytes / (1024 * 1024)
22}
23
24/// Calcula o teto seguro de concorrência para cargas pesadas de embedding.
25///
26/// Fórmula canônica:
27/// `permits = min(cpus, available_memory_mb / ram_por_task_mb) * 0.5`
28///
29/// O resultado é clampado entre `1` e `max_concurrency`.
30pub fn calculate_safe_concurrency(
31    available_mb: u64,
32    cpu_count: usize,
33    ram_per_task_mb: u64,
34    max_concurrency: usize,
35) -> usize {
36    let cpu_count = cpu_count.max(1);
37    let max_concurrency = max_concurrency.max(1);
38    let ram_per_task_mb = ram_per_task_mb.max(1);
39
40    let memory_bound = (available_mb / ram_per_task_mb) as usize;
41    let resource_bound = cpu_count.min(memory_bound).max(1);
42    let safe_with_margin = (resource_bound / 2).max(1);
43
44    safe_with_margin.min(max_concurrency)
45}
46
47/// Verifica se há memória disponível suficiente para iniciar o carregamento do modelo.
48///
49/// # Parâmetros
50/// - `min_mb`: piso mínimo em MiB de memória disponível (tipicamente
51///   [`crate::constants::MIN_AVAILABLE_MEMORY_MB`]).
52///
53/// # Erros
54/// Retorna [`AppError::LowMemory`] quando `available_mb < min_mb`.
55///
56/// # Retorno
57/// Retorna `Ok(available_mb)` com o valor real de memória disponível em MiB.
58pub fn check_available_memory(min_mb: u64) -> Result<u64, AppError> {
59    let available_mb = available_memory_mb();
60
61    if available_mb < min_mb {
62        return Err(AppError::LowMemory {
63            available_mb,
64            required_mb: min_mb,
65        });
66    }
67
68    Ok(available_mb)
69}
70
71#[cfg(test)]
72mod testes {
73    use super::*;
74
75    #[test]
76    fn check_available_memory_com_zero_sempre_passa() {
77        let resultado = check_available_memory(0);
78        assert!(
79            resultado.is_ok(),
80            "min_mb=0 deve sempre passar, got: {resultado:?}"
81        );
82        let mb = resultado.unwrap();
83        assert!(mb > 0, "sistema deve reportar memória positiva");
84    }
85
86    #[test]
87    fn check_available_memory_com_valor_gigante_falha() {
88        let resultado = check_available_memory(u64::MAX);
89        assert!(
90            matches!(resultado, Err(AppError::LowMemory { .. })),
91            "u64::MAX MiB deve falhar com LowMemory, got: {resultado:?}"
92        );
93    }
94
95    #[test]
96    fn low_memory_error_contem_valores_corretos() {
97        match check_available_memory(u64::MAX) {
98            Err(AppError::LowMemory {
99                available_mb,
100                required_mb,
101            }) => {
102                assert_eq!(required_mb, u64::MAX);
103                assert!(available_mb < u64::MAX);
104            }
105            outro => panic!("esperado LowMemory, got: {outro:?}"),
106        }
107    }
108
109    #[test]
110    fn calculate_safe_concurrency_respeita_metade_da_margem() {
111        let permits = calculate_safe_concurrency(8_000, 8, 1_000, 4);
112        assert_eq!(permits, 4);
113    }
114
115    #[test]
116    fn calculate_safe_concurrency_nunca_retorna_zero() {
117        let permits = calculate_safe_concurrency(100, 1, 10_000, 4);
118        assert_eq!(permits, 1);
119    }
120
121    #[test]
122    fn calculate_safe_concurrency_respeita_teto_maximo() {
123        let permits = calculate_safe_concurrency(128_000, 64, 500, 4);
124        assert_eq!(permits, 4);
125    }
126}
sqlite_graphrag/memory_guard.rs

sqlite_graphrag/
memory_guard.rs
