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/// Verifica se há memória disponível suficiente para iniciar o carregamento do modelo.
17///
18/// # Parâmetros
19/// - `min_mb`: piso mínimo em MiB de memória disponível (tipicamente
20///   [`crate::constants::MIN_AVAILABLE_MEMORY_MB`]).
21///
22/// # Erros
23/// Retorna [`AppError::LowMemory`] quando `available_mb < min_mb`.
24///
25/// # Retorno
26/// Retorna `Ok(available_mb)` com o valor real de memória disponível em MiB.
27pub fn check_available_memory(min_mb: u64) -> Result<u64, AppError> {
28    let sys =
29        System::new_with_specifics(RefreshKind::new().with_memory(MemoryRefreshKind::everything()));
30    let available_bytes = sys.available_memory();
31    let available_mb = available_bytes / (1024 * 1024);
32    drop(sys);
33
34    if available_mb < min_mb {
35        return Err(AppError::LowMemory {
36            available_mb,
37            required_mb: min_mb,
38        });
39    }
40
41    Ok(available_mb)
42}
43
44#[cfg(test)]
45mod testes {
46    use super::*;
47
48    #[test]
49    fn check_available_memory_com_zero_sempre_passa() {
50        let resultado = check_available_memory(0);
51        assert!(
52            resultado.is_ok(),
53            "min_mb=0 deve sempre passar, got: {resultado:?}"
54        );
55        let mb = resultado.unwrap();
56        assert!(mb > 0, "sistema deve reportar memória positiva");
57    }
58
59    #[test]
60    fn check_available_memory_com_valor_gigante_falha() {
61        let resultado = check_available_memory(u64::MAX);
62        assert!(
63            matches!(resultado, Err(AppError::LowMemory { .. })),
64            "u64::MAX MiB deve falhar com LowMemory, got: {resultado:?}"
65        );
66    }
67
68    #[test]
69    fn low_memory_error_contem_valores_corretos() {
70        match check_available_memory(u64::MAX) {
71            Err(AppError::LowMemory {
72                available_mb,
73                required_mb,
74            }) => {
75                assert_eq!(required_mb, u64::MAX);
76                assert!(available_mb < u64::MAX);
77            }
78            outro => panic!("esperado LowMemory, got: {outro:?}"),
79        }
80    }
81}
sqlite_graphrag/memory_guard.rs

sqlite_graphrag/
memory_guard.rs
