From 3f733fa7511f0342bc63d2d53a4163e1ce67bdd4 Mon Sep 17 00:00:00 2001 From: Elliot Worth Date: Sat, 2 May 2026 18:34:07 -0500 Subject: [PATCH 1/4] fix: enforce 50k row hard cap to prevent OOM on large queries - Always apply LIMIT to SELECT/SHOW/DESCRIBE/EXPLAIN (was conditional) - Hard cap at 50,000 rows regardless of user settings - Fix rows_truncated detection to compare against capped_limit - Lower frontend virtualization threshold from 500 to 100 rows - Add truncation warning banner when results are capped Closes #25 Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../crates/mas-core/src/query/executor.rs | 46 +++++++++---------- src/components/grid/ResultsGrid.tsx | 14 +++++- 2 files changed, 35 insertions(+), 25 deletions(-) diff --git a/src-tauri/crates/mas-core/src/query/executor.rs b/src-tauri/crates/mas-core/src/query/executor.rs index 82ff193..c8cdd49 100644 --- a/src-tauri/crates/mas-core/src/query/executor.rs +++ b/src-tauri/crates/mas-core/src/query/executor.rs @@ -35,31 +35,30 @@ impl QueryExecutor { database: Option, limit: Option, ) -> Result, CoreError> { + const HARD_MAX_ROWS: u64 = 50_000; + let effective_limit = limit.unwrap_or(HARD_MAX_ROWS); + let capped_limit = effective_limit.min(HARD_MAX_ROWS); + let pool = self.connection_manager.get_pool(&connection_id)?; let statements = split_statements(&sql); - // Apply global row limit to SELECT/SHOW/DESCRIBE statements - let statements: Vec = if let Some(max_rows) = limit { - statements - .into_iter() - .map(|stmt| { - let upper = stmt.trim().to_uppercase(); - if upper.starts_with("SELECT") - || upper.starts_with("SHOW") - || upper.starts_with("DESCRIBE") - || upper.starts_with("EXPLAIN") - { - // Remove any existing LIMIT/OFFSET at the end before injecting our own - let cleaned = strip_limit(&stmt); - format!("{} LIMIT {}", cleaned, max_rows) - } else { - stmt - } - }) - .collect() - } else { - statements - }; + // Always apply row limit to SELECT/SHOW/DESCRIBE statements (hard cap prevents OOM) + let statements: Vec = statements + .into_iter() + .map(|stmt| { + let upper = stmt.trim().to_uppercase(); + if upper.starts_with("SELECT") + || upper.starts_with("SHOW") + || upper.starts_with("DESCRIBE") + || upper.starts_with("EXPLAIN") + { + let cleaned = strip_limit(&stmt); + format!("{} LIMIT {}", cleaned, capped_limit) + } else { + stmt + } + }) + .collect(); tracing::Span::current().record("statement_count", statements.len()); tracing::trace!(sql = %sql, "Full SQL input"); @@ -173,8 +172,7 @@ impl QueryExecutor { ); // Detect if rows may be truncated by our injected LIMIT - let rows_truncated = - limit.is_some() && is_select && row_count >= limit.unwrap(); + let rows_truncated = row_count >= capped_limit; results.push(QueryResult { query_id, diff --git a/src/components/grid/ResultsGrid.tsx b/src/components/grid/ResultsGrid.tsx index 989b574..4e93e76 100644 --- a/src/components/grid/ResultsGrid.tsx +++ b/src/components/grid/ResultsGrid.tsx @@ -14,6 +14,7 @@ import { ArrowDown, Loader2, AlertCircle, + AlertTriangle, Copy, FileSpreadsheet, FileJson, @@ -383,7 +384,7 @@ export function ResultsGrid() { // Row virtualization: only render visible rows to avoid DOM bloat const ROW_HEIGHT = 32; // px per row const totalRows = data.length + (editing.editMode ? editing.inserts.length : 0); - const shouldVirtualize = totalRows > 500; + const shouldVirtualize = totalRows > 100; const rowVirtualizer = useVirtualizer({ count: totalRows, getScrollElement: () => rowVirtualizerParentRef.current, @@ -465,6 +466,17 @@ export function ResultsGrid() { onDiscard={editing.discardAll} /> + {/* Truncation warning */} + {activeResult?.rows_truncated && ( +
+ + + Results truncated to {activeResult.rows.length.toLocaleString()} rows (hard limit: 50,000). + Add a LIMIT clause to fetch fewer rows. + +
+ )} + {/* Result set tabs */} {results.length > 1 && (
From 7a783f12453c8b5f9294b98d691b40f3665c8ffe Mon Sep 17 00:00:00 2001 From: Elliot Worth Date: Sat, 2 May 2026 19:56:14 -0500 Subject: [PATCH 2/4] refactor: replace hard 50k row limit with dynamic memory guard - Remove hardcoded 50k row cap - Add sysinfo crate for process memory monitoring - MemoryGuard checks RSS every 1000 rows during fetch - Triggers when process hits 75% of system RAM - Gracefully stops fetch and returns accumulated rows with warning - Add OutOfMemory error variant to CoreError - Update frontend to show backend warnings - Truncation banner no longer mentions arbitrary 50k number --- src-tauri/Cargo.lock | 107 +++++++++- src-tauri/crates/mas-core/Cargo.toml | 1 + src-tauri/crates/mas-core/src/error.rs | 3 + .../crates/mas-core/src/query/executor.rs | 194 +++++++++++++++--- src/components/grid/ResultsGrid.tsx | 10 +- 5 files changed, 280 insertions(+), 35 deletions(-) diff --git a/src-tauri/Cargo.lock b/src-tauri/Cargo.lock index 9732109..7e6de19 100644 --- a/src-tauri/Cargo.lock +++ b/src-tauri/Cargo.lock @@ -2440,6 +2440,7 @@ dependencies = [ "serde", "serde_json", "sqlx", + "sysinfo", "tempfile", "thiserror 2.0.18", "tokio", @@ -2616,6 +2617,15 @@ dependencies = [ "minimal-lexical", ] +[[package]] +name = "ntapi" +version = "0.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c3b335231dfd352ffb0f8017f3b6027a4917f7df785ea2143d8af2adc66980ae" +dependencies = [ + "winapi", +] + [[package]] name = "nu-ansi-term" version = "0.50.3" @@ -2773,6 +2783,16 @@ dependencies = [ "objc2-core-foundation", ] +[[package]] +name = "objc2-io-kit" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "33fafba39597d6dc1fb709123dfa8289d39406734be322956a69f0931c73bb15" +dependencies = [ + "libc", + "objc2-core-foundation", +] + [[package]] name = "objc2-io-surface" version = "0.3.2" @@ -4629,6 +4649,20 @@ dependencies = [ "syn 2.0.117", ] +[[package]] +name = "sysinfo" +version = "0.38.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "92ab6a2f8bfe508deb3c6406578252e491d299cbbf3bc0529ecc3313aee4a52f" +dependencies = [ + "libc", + "memchr", + "ntapi", + "objc2-core-foundation", + "objc2-io-kit", + "windows 0.62.2", +] + [[package]] name = "system-deps" version = "6.2.2" @@ -4674,7 +4708,7 @@ dependencies = [ "tao-macros", "unicode-segmentation", "url", - "windows", + "windows 0.61.3", "windows-core 0.61.2", "windows-version", "x11-dl", @@ -4745,7 +4779,7 @@ dependencies = [ "webkit2gtk", "webview2-com", "window-vibrancy", - "windows", + "windows 0.61.3", ] [[package]] @@ -4871,7 +4905,7 @@ dependencies = [ "url", "webkit2gtk", "webview2-com", - "windows", + "windows 0.61.3", ] [[package]] @@ -4896,7 +4930,7 @@ dependencies = [ "url", "webkit2gtk", "webview2-com", - "windows", + "windows 0.61.3", "wry", ] @@ -5932,7 +5966,7 @@ checksum = "7130243a7a5b33c54a444e54842e6a9e133de08b5ad7b5861cd8ed9a6a5bc96a" dependencies = [ "webview2-com-macros", "webview2-com-sys", - "windows", + "windows 0.61.3", "windows-core 0.61.2", "windows-implement", "windows-interface", @@ -5956,7 +5990,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "381336cfffd772377d291702245447a5251a2ffa5bad679c99e61bc48bacbf9c" dependencies = [ "thiserror 2.0.18", - "windows", + "windows 0.61.3", "windows-core 0.61.2", ] @@ -6034,11 +6068,23 @@ version = "0.61.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9babd3a767a4c1aef6900409f85f5d53ce2544ccdfaa86dad48c91782c6d6893" dependencies = [ - "windows-collections", + "windows-collections 0.2.0", "windows-core 0.61.2", - "windows-future", + "windows-future 0.2.1", "windows-link 0.1.3", - "windows-numerics", + "windows-numerics 0.2.0", +] + +[[package]] +name = "windows" +version = "0.62.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "527fadee13e0c05939a6a05d5bd6eec6cd2e3dbd648b9f8e447c6518133d8580" +dependencies = [ + "windows-collections 0.3.2", + "windows-core 0.62.2", + "windows-future 0.3.2", + "windows-numerics 0.3.1", ] [[package]] @@ -6050,6 +6096,15 @@ dependencies = [ "windows-core 0.61.2", ] +[[package]] +name = "windows-collections" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "23b2d95af1a8a14a3c7367e1ed4fc9c20e0a26e79551b1454d72583c97cc6610" +dependencies = [ + "windows-core 0.62.2", +] + [[package]] name = "windows-core" version = "0.61.2" @@ -6084,7 +6139,18 @@ checksum = "fc6a41e98427b19fe4b73c550f060b59fa592d7d686537eebf9385621bfbad8e" dependencies = [ "windows-core 0.61.2", "windows-link 0.1.3", - "windows-threading", + "windows-threading 0.1.0", +] + +[[package]] +name = "windows-future" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e1d6f90251fe18a279739e78025bd6ddc52a7e22f921070ccdc67dde84c605cb" +dependencies = [ + "windows-core 0.62.2", + "windows-link 0.2.1", + "windows-threading 0.2.1", ] [[package]] @@ -6131,6 +6197,16 @@ dependencies = [ "windows-link 0.1.3", ] +[[package]] +name = "windows-numerics" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6e2e40844ac143cdb44aead537bbf727de9b044e107a0f1220392177d15b0f26" +dependencies = [ + "windows-core 0.62.2", + "windows-link 0.2.1", +] + [[package]] name = "windows-result" version = "0.3.4" @@ -6293,6 +6369,15 @@ dependencies = [ "windows-link 0.1.3", ] +[[package]] +name = "windows-threading" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3949bd5b99cafdf1c7ca86b43ca564028dfe27d66958f2470940f73d86d75b37" +dependencies = [ + "windows-link 0.2.1", +] + [[package]] name = "windows-version" version = "0.1.7" @@ -6660,7 +6745,7 @@ dependencies = [ "webkit2gtk", "webkit2gtk-sys", "webview2-com", - "windows", + "windows 0.61.3", "windows-core 0.61.2", "windows-version", "x11-dl", diff --git a/src-tauri/crates/mas-core/Cargo.toml b/src-tauri/crates/mas-core/Cargo.toml index bf66a57..6eb830a 100644 --- a/src-tauri/crates/mas-core/Cargo.toml +++ b/src-tauri/crates/mas-core/Cargo.toml @@ -16,6 +16,7 @@ chrono = { workspace = true } rusqlite = { workspace = true } dashmap = "6" futures = "0.3" +sysinfo = "0.38.4" [dev-dependencies] tokio = { workspace = true, features = ["test-util"] } diff --git a/src-tauri/crates/mas-core/src/error.rs b/src-tauri/crates/mas-core/src/error.rs index c7e7862..2045b7a 100644 --- a/src-tauri/crates/mas-core/src/error.rs +++ b/src-tauri/crates/mas-core/src/error.rs @@ -32,6 +32,9 @@ pub enum CoreError { #[error("Cancelled")] Cancelled, + #[error("Out of memory: {0}")] + OutOfMemory(String), + #[error(transparent)] Sqlx(#[from] sqlx::Error), diff --git a/src-tauri/crates/mas-core/src/query/executor.rs b/src-tauri/crates/mas-core/src/query/executor.rs index c8cdd49..e841afb 100644 --- a/src-tauri/crates/mas-core/src/query/executor.rs +++ b/src-tauri/crates/mas-core/src/query/executor.rs @@ -35,30 +35,33 @@ impl QueryExecutor { database: Option, limit: Option, ) -> Result, CoreError> { - const HARD_MAX_ROWS: u64 = 50_000; - let effective_limit = limit.unwrap_or(HARD_MAX_ROWS); - let capped_limit = effective_limit.min(HARD_MAX_ROWS); - let pool = self.connection_manager.get_pool(&connection_id)?; let statements = split_statements(&sql); - // Always apply row limit to SELECT/SHOW/DESCRIBE statements (hard cap prevents OOM) - let statements: Vec = statements - .into_iter() - .map(|stmt| { - let upper = stmt.trim().to_uppercase(); - if upper.starts_with("SELECT") - || upper.starts_with("SHOW") - || upper.starts_with("DESCRIBE") - || upper.starts_with("EXPLAIN") - { - let cleaned = strip_limit(&stmt); - format!("{} LIMIT {}", cleaned, capped_limit) - } else { - stmt - } - }) - .collect(); + // Apply user-specified row limit to SELECT/SHOW/DESCRIBE statements (if provided) + let statements: Vec = if let Some(max_rows) = limit { + statements + .into_iter() + .map(|stmt| { + let upper = stmt.trim().to_uppercase(); + if upper.starts_with("SELECT") + || upper.starts_with("SHOW") + || upper.starts_with("DESCRIBE") + || upper.starts_with("EXPLAIN") + { + let cleaned = strip_limit(&stmt); + format!("{} LIMIT {}", cleaned, max_rows) + } else { + stmt + } + }) + .collect() + } else { + statements + }; + + // Memory guard: detect OOM before the OS kills us + let mut mem_guard = MemoryGuard::new(); tracing::Span::current().record("statement_count", statements.len()); tracing::trace!(sql = %sql, "Full SQL input"); @@ -98,6 +101,17 @@ impl QueryExecutor { Either::Right(row) => { // SELECT/SHOW/DESCRIBE/EXPLAIN row — accumulate until trailing Left. if stmt_idx >= 0 { + // Check memory every 1000 rows to prevent OOM + if current_rows.len() % 1000 == 0 { + if mem_guard.check().is_err() { + mem_guard.set_triggered(); + tracing::warn!( + rows_accumulated = current_rows.len(), + "Memory limit reached, stopping query fetch" + ); + break; + } + } current_rows.push(row); } } @@ -171,8 +185,10 @@ impl QueryExecutor { "Query executed" ); - // Detect if rows may be truncated by our injected LIMIT - let rows_truncated = row_count >= capped_limit; + // Detect truncation: user limit enforced, or memory guard kicked in + let rows_truncated = + limit.is_some() && is_select && row_count >= limit.unwrap() + || mem_guard.triggered(); results.push(QueryResult { query_id, @@ -228,6 +244,65 @@ impl QueryExecutor { } } + // If memory guard triggered mid-stream, process accumulated rows for current statement + if mem_guard.triggered() && stmt_idx >= 0 && !current_rows.is_empty() { + let idx = stmt_idx as usize; + let stmt = &statements[idx]; + let query_id = uuid::Uuid::new_v4().to_string(); + let execution_time = start.elapsed().as_millis() as u64; + + let upper = stmt.to_uppercase(); + let is_select = upper.starts_with("SELECT") + || upper.starts_with("SHOW") + || upper.starts_with("DESCRIBE") + || upper.starts_with("EXPLAIN"); + + if is_select { + let columns: Vec = + if let Some(first_row) = current_rows.first() { + first_row + .columns() + .iter() + .map(|col| ColumnMeta { + name: col.name().to_string(), + data_type: col.type_info().name().to_string(), + nullable: true, + is_primary_key: false, + }) + .collect() + } else { + Vec::new() + }; + + let result_rows: Vec> = current_rows + .iter() + .map(|row| { + row.columns() + .iter() + .enumerate() + .map(|(i, col)| { + extract_value(row, i, col.type_info().name()) + }) + .collect() + }) + .collect(); + + let row_count = result_rows.len() as u64; + + results.push(QueryResult { + query_id, + statement_index: idx, + columns, + rows: result_rows, + rows_affected: row_count, + execution_time_ms: execution_time, + warnings: vec!["Query truncated: memory limit reached".to_string()], + rows_truncated: true, + total_rows_available: Some(row_count), + }); + } + } + Ok(results) } } @@ -469,6 +544,79 @@ fn find_keyword_offset(s: &str, keyword: &str) -> Option { last_pos } +/// Monitors process memory usage to prevent OOM crashes. +/// Checks every N rows during query execution and triggers when +/// process RSS exceeds a threshold of total system memory. +struct MemoryGuard { + threshold_bytes: u64, + triggered: bool, +} + +impl MemoryGuard { + /// Create a new guard. Threshold defaults to 75% of system RAM. + fn new() -> Self { + let mut sys = sysinfo::System::new(); + sys.refresh_memory(); + + let total = sys.total_memory(); + let threshold = (total as f64 * 0.75) as u64; + + tracing::debug!( + total_mb = total / 1024 / 1024, + threshold_mb = threshold / 1024 / 1024, + "Memory guard initialized" + ); + + Self { + threshold_bytes: threshold, + triggered: false, + } + } + + /// Check current process memory against threshold. + /// Returns Err if memory usage exceeds threshold. + fn check(&self) -> Result<(), CoreError> { + if self.triggered { + return Err(CoreError::OutOfMemory( + "Query stopped: process memory limit reached".to_string(), + )); + } + + let mut sys = sysinfo::System::new(); + sys.refresh_memory(); + + // Find our process and check memory usage + let current_pid = std::process::id(); + if let Some(process) = sys.process(sysinfo::Pid::from_u32(current_pid)) { + let rss = process.memory(); + let rss_mb = rss / 1024 / 1024; + let threshold_mb = self.threshold_bytes / 1024 / 1024; + + if rss > self.threshold_bytes { + tracing::warn!( + rss_mb, + threshold_mb, + "Process memory exceeds threshold" + ); + return Err(CoreError::OutOfMemory(format!( + "Process memory ({rss_mb} MB) exceeds limit ({threshold_mb} MB). \ + Add a LIMIT clause to reduce result size." + ))); + } + } + + Ok(()) + } + + fn triggered(&self) -> bool { + self.triggered + } + + fn set_triggered(&mut self) { + self.triggered = true; + } +} + #[cfg(test)] mod tests { use super::*; diff --git a/src/components/grid/ResultsGrid.tsx b/src/components/grid/ResultsGrid.tsx index 4e93e76..609f432 100644 --- a/src/components/grid/ResultsGrid.tsx +++ b/src/components/grid/ResultsGrid.tsx @@ -471,12 +471,20 @@ export function ResultsGrid() {
- Results truncated to {activeResult.rows.length.toLocaleString()} rows (hard limit: 50,000). + Results truncated to {activeResult.rows.length.toLocaleString()} rows. Add a LIMIT clause to fetch fewer rows.
)} + {/* Backend warnings (e.g., memory guard) */} + {activeResult?.warnings?.map((warning, idx) => ( +
+ + {warning} +
+ ))} + {/* Result set tabs */} {results.length > 1 && (
From 67bf55c2d754bdaa92294e2ac82577c2829bc1a1 Mon Sep 17 00:00:00 2001 From: Elliot Worth Date: Sun, 3 May 2026 09:24:48 -0500 Subject: [PATCH 3/4] fix: rustfmt and clippy warnings in memory guard - Collapse nested if into single condition (clippy::collapsible_if) - Use is_multiple_of(1000) instead of % 1000 == 0 (clippy::manual_is_multiple_of) - Run cargo fmt for consistent formatting --- .../crates/mas-core/src/query/executor.rs | 55 ++++++++----------- 1 file changed, 23 insertions(+), 32 deletions(-) diff --git a/src-tauri/crates/mas-core/src/query/executor.rs b/src-tauri/crates/mas-core/src/query/executor.rs index e841afb..e3bed8e 100644 --- a/src-tauri/crates/mas-core/src/query/executor.rs +++ b/src-tauri/crates/mas-core/src/query/executor.rs @@ -102,15 +102,13 @@ impl QueryExecutor { // SELECT/SHOW/DESCRIBE/EXPLAIN row — accumulate until trailing Left. if stmt_idx >= 0 { // Check memory every 1000 rows to prevent OOM - if current_rows.len() % 1000 == 0 { - if mem_guard.check().is_err() { - mem_guard.set_triggered(); - tracing::warn!( - rows_accumulated = current_rows.len(), - "Memory limit reached, stopping query fetch" - ); - break; - } + if current_rows.len().is_multiple_of(1000) && mem_guard.check().is_err() { + mem_guard.set_triggered(); + tracing::warn!( + rows_accumulated = current_rows.len(), + "Memory limit reached, stopping query fetch" + ); + break; } current_rows.push(row); } @@ -258,21 +256,20 @@ impl QueryExecutor { || upper.starts_with("EXPLAIN"); if is_select { - let columns: Vec = - if let Some(first_row) = current_rows.first() { - first_row - .columns() - .iter() - .map(|col| ColumnMeta { - name: col.name().to_string(), - data_type: col.type_info().name().to_string(), - nullable: true, - is_primary_key: false, - }) - .collect() - } else { - Vec::new() - }; + let columns: Vec = if let Some(first_row) = current_rows.first() { + first_row + .columns() + .iter() + .map(|col| ColumnMeta { + name: col.name().to_string(), + data_type: col.type_info().name().to_string(), + nullable: true, + is_primary_key: false, + }) + .collect() + } else { + Vec::new() + }; let result_rows: Vec> = current_rows .iter() @@ -280,9 +277,7 @@ impl QueryExecutor { row.columns() .iter() .enumerate() - .map(|(i, col)| { - extract_value(row, i, col.type_info().name()) - }) + .map(|(i, col)| extract_value(row, i, col.type_info().name())) .collect() }) .collect(); @@ -593,11 +588,7 @@ impl MemoryGuard { let threshold_mb = self.threshold_bytes / 1024 / 1024; if rss > self.threshold_bytes { - tracing::warn!( - rss_mb, - threshold_mb, - "Process memory exceeds threshold" - ); + tracing::warn!(rss_mb, threshold_mb, "Process memory exceeds threshold"); return Err(CoreError::OutOfMemory(format!( "Process memory ({rss_mb} MB) exceeds limit ({threshold_mb} MB). \ Add a LIMIT clause to reduce result size." From 07b1f1f8bebc0e238dc2c27001dde53d3044cc16 Mon Sep 17 00:00:00 2001 From: Elliot Worth Date: Thu, 7 May 2026 21:50:11 -0500 Subject: [PATCH 4/4] fix: use available memory instead of total for OOM detection - Changed MemoryGuard to monitor available system memory, not total - Trigger threshold: 95% of available memory consumed (5% buffer) - Previous logic used 75% of total memory which could allow crashes when other processes consumed RAM - New logic prevents runaway queries when system memory is critically low - Removed per-process RSS tracking; now monitors system-wide availability Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../crates/mas-core/src/query/executor.rs | 56 +++++++++---------- 1 file changed, 28 insertions(+), 28 deletions(-) diff --git a/src-tauri/crates/mas-core/src/query/executor.rs b/src-tauri/crates/mas-core/src/query/executor.rs index e3bed8e..e8a20b8 100644 --- a/src-tauri/crates/mas-core/src/query/executor.rs +++ b/src-tauri/crates/mas-core/src/query/executor.rs @@ -541,59 +541,59 @@ fn find_keyword_offset(s: &str, keyword: &str) -> Option { /// Monitors process memory usage to prevent OOM crashes. /// Checks every N rows during query execution and triggers when -/// process RSS exceeds a threshold of total system memory. +/// available system memory usage exceeds 95%. struct MemoryGuard { - threshold_bytes: u64, triggered: bool, } impl MemoryGuard { - /// Create a new guard. Threshold defaults to 75% of system RAM. + /// Create a new guard. Monitors available memory. fn new() -> Self { let mut sys = sysinfo::System::new(); sys.refresh_memory(); + let available = sys.available_memory(); let total = sys.total_memory(); - let threshold = (total as f64 * 0.75) as u64; + let available_mb = available / 1024 / 1024; + let total_mb = total / 1024 / 1024; tracing::debug!( - total_mb = total / 1024 / 1024, - threshold_mb = threshold / 1024 / 1024, - "Memory guard initialized" + available_mb, + total_mb, + "Memory guard initialized, will trigger at 95% of available memory" ); - Self { - threshold_bytes: threshold, - triggered: false, - } + Self { triggered: false } } - /// Check current process memory against threshold. - /// Returns Err if memory usage exceeds threshold. + /// Check current available memory. Triggers if available memory drops below 5%. + /// Returns Err if we're at risk of OOM. fn check(&self) -> Result<(), CoreError> { if self.triggered { return Err(CoreError::OutOfMemory( - "Query stopped: process memory limit reached".to_string(), + "Query stopped: available memory critically low".to_string(), )); } let mut sys = sysinfo::System::new(); sys.refresh_memory(); - // Find our process and check memory usage - let current_pid = std::process::id(); - if let Some(process) = sys.process(sysinfo::Pid::from_u32(current_pid)) { - let rss = process.memory(); - let rss_mb = rss / 1024 / 1024; - let threshold_mb = self.threshold_bytes / 1024 / 1024; - - if rss > self.threshold_bytes { - tracing::warn!(rss_mb, threshold_mb, "Process memory exceeds threshold"); - return Err(CoreError::OutOfMemory(format!( - "Process memory ({rss_mb} MB) exceeds limit ({threshold_mb} MB). \ - Add a LIMIT clause to reduce result size." - ))); - } + let available = sys.available_memory(); + let total = sys.total_memory(); + let used = total.saturating_sub(available); + let usage_percent = (used as f64 / total as f64 * 100.0) as u32; + let available_mb = available / 1024 / 1024; + + if usage_percent >= 95 { + tracing::warn!( + usage_percent, + available_mb, + "System memory usage at {usage_percent}%, stopping query" + ); + return Err(CoreError::OutOfMemory(format!( + "System memory critically low ({available_mb} MB available, {usage_percent}% used). \ + Add a LIMIT clause to reduce result size." + ))); } Ok(())