From b50a32a038b6299e1c2cb87c3fa6f71c3f733f83 Mon Sep 17 00:00:00 2001 From: Alex-Gilbert Date: Wed, 18 Feb 2026 11:20:16 -0600 Subject: [PATCH] fix(solver): add max_results cap to prevent OOM in full-name solver MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The full-name solver collected all results into a Vec with no limit, causing OOM crashes for realistic targets where millions of valid first×last combinations exist. Thread max_results through solve_subtree, dfs_first, and the Cartesian product loop so the search stops early once enough results are found. Co-Authored-By: Claude Opus 4.6 --- solver_rs/src/dfs.rs | 16 +++ solver_rs/src/full_name.rs | 20 +++- solver_rs/src/main.rs | 2 + tests/test_full_name_stress.py | 193 +++++++++++++++++++++++++++++++++ 4 files changed, 230 insertions(+), 1 deletion(-) create mode 100644 tests/test_full_name_stress.py diff --git a/solver_rs/src/dfs.rs b/solver_rs/src/dfs.rs index ba80f7b..fea9d2a 100644 --- a/solver_rs/src/dfs.rs +++ b/solver_rs/src/dfs.rs @@ -77,6 +77,7 @@ fn compute_state_max(wt: &WidthTable, constraint: &Constraint) -> Vec { } /// Solve a single subtree: find all strings matching the target width. +/// Stops early once `max_results` matches are found (0 = unlimited). pub fn solve_subtree( wt: &WidthTable, target: f64, @@ -90,6 +91,7 @@ pub fn solve_subtree( start_state: usize, equiv: &EquivClasses, deduped_allowed: &[Vec], + max_results: usize, ) -> Vec { let smax = compute_state_max(wt, constraint); let byte_to_idx = build_byte_to_idx(&wt.charset); @@ -98,6 +100,9 @@ pub fn solve_subtree( if prefix.is_empty() { // Start from scratch — iterate first chars (deduped: one representative per class) for &first_idx in &deduped_allowed[start_state] { + if max_results > 0 && results.len() >= max_results { + break; + } let start_width = wt.left_edge[first_idx]; if start_width > target + tolerance { continue; @@ -130,6 +135,7 @@ pub fn solve_subtree( &mut path, ns as usize, 0, // prefix_len + max_results, &mut results, ); } @@ -152,6 +158,7 @@ pub fn solve_subtree( &mut path, start_state, prefix.len(), + max_results, &mut results, ); } @@ -177,8 +184,13 @@ fn dfs( path: &mut Vec, cstate: usize, prefix_len: usize, + max_results: usize, results: &mut Vec, ) { + if max_results > 0 && results.len() >= max_results { + return; + } + let current_length = prefix_len + depth; if current_length >= min_length && constraint.accept_states[cstate] { @@ -242,6 +254,7 @@ fn dfs( path, ns as usize, prefix_len, + max_results, results, ); path.pop(); @@ -607,6 +620,7 @@ mod tests { "", 0.0, -1, &constraint, 0, &equiv, &deduped, + 0, ); let mut texts: Vec = results.iter().map(|r| r.text.clone()).collect(); texts.sort(); @@ -635,6 +649,7 @@ mod tests { "", 0.0, -1, &constraint, 0, &equiv, &deduped, + 0, ); for r in &results { assert!((r.width - 15.0).abs() <= 1.0, "result {} has error {} > tolerance", r.text, r.error); @@ -667,6 +682,7 @@ mod tests { "", 0.0, -1, &constraint, 0, &equiv, &deduped, + 0, ); let mut texts: Vec = results.iter().map(|r| r.text.clone()).collect(); texts.sort(); diff --git a/solver_rs/src/full_name.rs b/solver_rs/src/full_name.rs index 2cbf98c..51b827d 100644 --- a/solver_rs/src/full_name.rs +++ b/solver_rs/src/full_name.rs @@ -59,6 +59,7 @@ pub fn solve_full_name( target: f64, tolerance: f64, uppercase_only: bool, + max_results: usize, ) -> Vec { let n = word_charset.len(); @@ -170,6 +171,7 @@ pub fn solve_full_name( &equiv1, &byte_to_idx, first_deduped, body_deduped, &equiv2, &deduped2, 0, 0.0, 0, &mut Vec::new(), true, + max_results, &mut results, &mut seen, ); @@ -209,9 +211,14 @@ fn dfs_first( last_idx: usize, path: &mut Vec, is_start: bool, + max_results: usize, results: &mut Vec, seen: &mut HashSet, ) { + if max_results > 0 && results.len() >= max_results { + return; + } + // At valid first-word leaf, solve second word if depth >= min_word_len && !is_start { let first_width = acc_width + wt1.right_edge[last_idx]; @@ -223,19 +230,29 @@ fn dfs_first( let auto_min2 = auto_min2.max(min_word_len); if auto_min2 <= auto_max2 { + // Cap inner solve to remaining budget + let inner_limit = if max_results > 0 { + max_results.saturating_sub(results.len()) + } else { + 0 + }; let second_results = solve_subtree( wt2, remaining, tolerance, auto_min2, auto_max2, "", 0.0, -1, constraint, 0, equiv2, deduped2, + inner_limit, ); // Expand first word (all class-member variants) let first_expanded = expand_match(path, equiv1, charset, byte_to_idx); - for first_bytes in &first_expanded { + 'outer: for first_bytes in &first_expanded { let first_text = unsafe { String::from_utf8_unchecked(first_bytes.clone()) }; for r2 in &second_results { + if max_results > 0 && results.len() >= max_results { + break 'outer; + } let full_text = format!("{} {}", first_text, r2.text); if !seen.contains(&full_text) { let full_width = first_width + sp + r2.width; @@ -293,6 +310,7 @@ fn dfs_first( equiv1, byte_to_idx, first_deduped, body_deduped, equiv2, deduped2, depth + 1, new_width, next_idx, path, false, + max_results, results, seen, ); path.pop(); diff --git a/solver_rs/src/main.rs b/solver_rs/src/main.rs index bbe37ab..935fd06 100644 --- a/solver_rs/src/main.rs +++ b/solver_rs/src/main.rs @@ -158,6 +158,7 @@ async fn handle_solve( prefix, *prefix_width, *last_idx, &constraint, *pfx_state, &equiv, &deduped_allowed, + result_limit, ); results.sort_by(|a, b| a.error.partial_cmp(&b.error).unwrap()); @@ -205,6 +206,7 @@ async fn handle_full_name( &req.wt2_table, &req.wt2_right_edge, &req.space_advance, &req.left_after_space, req.target, req.tolerance, req.uppercase_only, + result_limit, ); let mut sent = 0; diff --git a/tests/test_full_name_stress.py b/tests/test_full_name_stress.py new file mode 100644 index 0000000..bd114fd --- /dev/null +++ b/tests/test_full_name_stress.py @@ -0,0 +1,193 @@ +"""Stress test for the Rust full-name solver. + +Starts the solver server, sends a full-name request with a realistic +font and wide target, and verifies it completes without crashing. +""" + +import json +import subprocess +import sys +import time + +import requests +from PIL import ImageFont + +sys.path.insert(0, ".") +from unredact.pipeline.width_table import CHARSETS, build_width_table + +SOLVER_PORT = 3199 # Use a non-default port to avoid conflicts +SOLVER_URL = f"http://127.0.0.1:{SOLVER_PORT}" + + +def build_full_name_payload( + font: ImageFont.FreeTypeFont, + target_width: float, + tolerance: float, + left_context: str = "", + right_context: str = "", + uppercase_only: bool = False, + max_results: int = 100, + filter_mode: str = "names", +) -> dict: + word_charset = CHARSETS["uppercase"] if uppercase_only else CHARSETS["alpha"] + wt1 = build_width_table(font, word_charset, left_context, "") + wt2 = build_width_table(font, word_charset, "", right_context) + + space_advance = [] + for c in word_charset: + space_advance.append(font.getlength(c + " ") - font.getlength(c)) + space_base = font.getlength(" ") + left_after_space = [] + for c in word_charset: + left_after_space.append(font.getlength(" " + c) - space_base) + + return { + "word_charset": word_charset, + "wt1_table": wt1.width_table.flatten().tolist(), + "wt1_left_edge": wt1.left_edge.tolist(), + "wt1_right_edge": wt1.right_edge.tolist(), + "wt2_table": wt2.width_table.flatten().tolist(), + "wt2_right_edge": wt2.right_edge.tolist(), + "space_advance": space_advance, + "left_after_space": left_after_space, + "target": float(target_width), + "tolerance": float(tolerance), + "uppercase_only": uppercase_only, + "max_results": max_results, + "filter": filter_mode, + "filter_prefix": "", + "filter_suffix": "", + } + + +def start_solver(): + """Start the Rust solver server and wait for it to be ready.""" + import os + env = os.environ.copy() + env["SOLVER_PORT"] = str(SOLVER_PORT) + proc = subprocess.Popen( + ["cargo", "run", "--release"], + cwd="solver_rs", + env=env, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + ) + # Wait for health check + for _ in range(30): + try: + r = requests.get(f"{SOLVER_URL}/health", timeout=1) + if r.status_code == 200: + return proc + except requests.ConnectionError: + pass + time.sleep(0.5) + proc.kill() + raise RuntimeError("Solver failed to start") + + +def test_full_name(font, label, target, tolerance, uppercase_only=False, timeout=30): + """Send a full-name solve request and check it completes.""" + payload = build_full_name_payload( + font, target, tolerance, + uppercase_only=uppercase_only, + max_results=200, + filter_mode="names", + ) + + print(f"\n{'='*60}") + print(f"TEST: {label}") + print(f" target={target}px, tolerance={tolerance}px") + print(f" charset={'uppercase' if uppercase_only else 'alpha'} ({len(payload['word_charset'])} chars)") + print(f" uppercase_only={uppercase_only}") + + start = time.time() + try: + resp = requests.post( + f"{SOLVER_URL}/solve/full-name", + json=payload, + stream=True, + timeout=timeout, + ) + results = [] + for line in resp.iter_lines(): + if not line: + continue + text = line.decode() + if text.startswith("data: "): + data = json.loads(text[6:]) + if "done" in data: + elapsed = time.time() - start + print(f" DONE in {elapsed:.2f}s — {data['total']} results") + break + results.append(data) + if len(results) <= 5: + print(f" match: {data['text']} (width={data['width']:.2f}, err={data['error']:.2f})") + elif len(results) == 6: + print(f" ... (more results)") + else: + elapsed = time.time() - start + print(f" Stream ended without done event in {elapsed:.2f}s") + + print(f" Total matches received: {len(results)}") + return True + + except requests.Timeout: + elapsed = time.time() - start + print(f" TIMEOUT after {elapsed:.2f}s") + return False + except requests.ConnectionError: + print(f" CONNECTION ERROR — solver may have crashed!") + return False + except Exception as e: + print(f" ERROR: {e}") + return False + + +def main(): + font_path = "/usr/share/fonts/TTF/Arial.TTF" + font_size = 16 + + print(f"Loading font: {font_path} at {font_size}pt") + font = ImageFont.truetype(font_path, font_size) + + print("Starting Rust solver...") + proc = start_solver() + print(f"Solver running (PID {proc.pid})") + + all_passed = True + try: + # Test 1: Small target (short names, should be fast) + ok = test_full_name(font, "Small capitalized name (~80px)", 80.0, 2.0, timeout=15) + all_passed &= ok + + # Test 2: Medium target (typical name) + ok = test_full_name(font, "Medium capitalized name (~120px)", 120.0, 2.0, timeout=30) + all_passed &= ok + + # Test 3: Large target (long names — the stress test) + ok = test_full_name(font, "Large capitalized name (~180px)", 180.0, 3.0, timeout=60) + all_passed &= ok + + # Test 4: All-caps small + ok = test_full_name(font, "Small CAPS name (~80px)", 80.0, 2.0, uppercase_only=True, timeout=15) + all_passed &= ok + + # Test 5: All-caps medium + ok = test_full_name(font, "Medium CAPS name (~120px)", 120.0, 2.0, uppercase_only=True, timeout=30) + all_passed &= ok + + finally: + proc.terminate() + proc.wait(timeout=5) + print(f"\nSolver stopped (exit code: {proc.returncode})") + + print(f"\n{'='*60}") + if all_passed: + print("ALL TESTS PASSED") + else: + print("SOME TESTS FAILED") + return 0 if all_passed else 1 + + +if __name__ == "__main__": + sys.exit(main())