perf(diffctx): Arc<str> for Fragment::content + concurrent stdout drain in git

nikolay-e · nikolay-e · commit 92f70448dd2b · 2026-04-26T07:19:02.000+02:00
Reduces per-instance peak RSS ~140x on huggingface/transformers (22 GB → 153 MB)
by sharing fragment content via Arc&lt;str&gt; instead of String clones across pipeline
stages. Also fixes a pipe-buffer deadlock in wait_with_timeout: stdout/stderr are
now drained on background threads concurrently with try_wait, instead of after the
child exits — git ls-files on large repos exceeds the 64 KB pipe buffer and would
otherwise block until the global timeout.

Also: env-overridable BENCH_WORKERS / BENCH_BATCH_SIZE, batched ProcessPoolExecutor
that survives single-worker OOMs, Docker bench harness with sweep orchestrator and
aggregation report.
diff --git a/.dockerignore b/.dockerignore
@@ -37,6 +37,15 @@ private/
 .zsh_history
 .sh_history
 
+# Rust build artifacts (3.5G+)
+diffctx/target/
+target/
+
+# Benchmark caches and results
+.cache/
+results/
+results_archive/
+
 # Python artifacts
 __pycache__/
 *.py[cod]
diff --git a/Dockerfile.bench b/Dockerfile.bench
@@ -0,0 +1,52 @@
+FROM python:3.12-slim-bookworm AS rust-builder
+
+RUN apt-get update && apt-get install -y --no-install-recommends \
+    git build-essential pkg-config libssl-dev curl ca-certificates \
+    && rm -rf /var/lib/apt/lists/*
+
+ENV RUSTUP_HOME=/usr/local/rustup CARGO_HOME=/usr/local/cargo PATH=/usr/local/cargo/bin:$PATH
+RUN curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | \
+    sh -s -- -y --default-toolchain 1.92.0 --profile minimal
+
+ENV VIRTUAL_ENV=/opt/venv
+RUN python3 -m venv $VIRTUAL_ENV
+ENV PATH="$VIRTUAL_ENV/bin:$PATH"
+
+RUN pip install --no-cache-dir "maturin>=1,<2"
+
+WORKDIR /build/diffctx
+COPY diffctx/Cargo.toml diffctx/Cargo.lock diffctx/pyproject.toml ./
+COPY diffctx/src ./src
+COPY diffctx/python ./python
+
+RUN maturin build --release --out /wheels
+
+
+FROM python:3.12-slim-bookworm AS runtime
+
+RUN apt-get update && apt-get install -y --no-install-recommends \
+    git ca-certificates \
+    && rm -rf /var/lib/apt/lists/*
+
+ENV VIRTUAL_ENV=/opt/venv
+RUN python3 -m venv $VIRTUAL_ENV
+ENV PATH="$VIRTUAL_ENV/bin:$PATH"
+
+COPY --from=rust-builder /wheels/*.whl /tmp/wheels/
+RUN pip install --no-cache-dir /tmp/wheels/*.whl && rm -rf /tmp/wheels
+
+WORKDIR /app
+COPY pyproject.toml README.md ./
+COPY src ./src
+COPY benchmarks ./benchmarks
+
+RUN pip install --no-cache-dir -e ".[full]" && \
+    pip install --no-cache-dir datasets rank-bm25
+
+ENV CB_REPOS_DIR=/cache/contextbench_repos
+ENV PYTHONUNBUFFERED=1
+
+VOLUME ["/cache/contextbench_repos", "/app/results"]
+
+ENTRYPOINT ["python", "-m", "benchmarks"]
+CMD ["cb", "--help"]
diff --git a/benchmarks/common.py b/benchmarks/common.py
@@ -10,7 +10,7 @@
 import uuid
 from pathlib import Path
 
-WORKERS = 11
+WORKERS = int(os.environ.get("BENCH_WORKERS", "11"))
 RESULTS_DIR = Path("results")
 
 LINES_RE = re.compile(r"^(\d+)-(\d+)$")
@@ -240,8 +240,10 @@ def parse_lines_field(lines_str: str) -> tuple[int, int] | None:
 def load_results(path: Path) -> list[dict]:
     data = json.loads(path.read_text())
     if isinstance(data, dict) and "results" in data:
-        return data["results"]
-    return data
+        return list(data["results"])
+    if isinstance(data, list):
+        return data
+    raise ValueError(f"unexpected results shape in {path}: {type(data).__name__}")
 
 
 def _git_commit_sha() -> str:
@@ -314,15 +316,24 @@ def _run_serial(worker_fn, run_args: list, collect: str) -> list:
 
 def _run_pool(worker_fn, run_args: list, workers: int, collect: str) -> list:
     from concurrent.futures import ProcessPoolExecutor, as_completed
+    from concurrent.futures.process import BrokenProcessPool
 
+    batch_size = int(os.environ.get("BENCH_BATCH_SIZE", str(max(workers * 4, 20))))
     results: list = []
-    with ProcessPoolExecutor(max_workers=workers, initializer=_init_worker) as pool:
-        futures = {pool.submit(worker_fn, a): a[0] for a in run_args}
-        for future in as_completed(futures):
-            try:
-                _collect_result(results, future.result(), collect)
-            except Exception as e:
-                print(f"  WORKER CRASH [{futures[future]}]: {type(e).__name__}: {e}", flush=True)
+    for batch_start in range(0, len(run_args), batch_size):
+        batch = run_args[batch_start : batch_start + batch_size]
+        try:
+            with ProcessPoolExecutor(max_workers=workers, initializer=_init_worker) as pool:
+                futures = {pool.submit(worker_fn, a): a[0] for a in batch}
+                for future in as_completed(futures):
+                    try:
+                        _collect_result(results, future.result(), collect)
+                    except BrokenProcessPool as e:
+                        print(f"  WORKER CRASH [{futures[future]}]: {type(e).__name__}", flush=True)
+                    except Exception as e:
+                        print(f"  WORKER CRASH [{futures[future]}]: {type(e).__name__}: {e}", flush=True)
+        except BrokenProcessPool as e:
+            print(f"  POOL CRASH batch {batch_start}-{batch_start+len(batch)}: {e}", flush=True)
     return results
 
 
diff --git a/diffctx/src/fragmentation.rs b/diffctx/src/fragmentation.rs
@@ -148,11 +148,12 @@ fn truncate_generated_fragments(file_frags: Vec<Fragment>) -> Vec<Fragment> {
                 remaining
             );
             let new_end = frag.start_line() + max_lines - 1;
+            let identifiers = extract_identifiers(&truncated_content, 2);
             Fragment {
                 id: FragmentId::new(frag.id.path.clone(), frag.start_line(), new_end),
                 kind: frag.kind,
-                content: truncated_content.clone(),
-                identifiers: extract_identifiers(&truncated_content, 2),
+                content: Arc::from(truncated_content),
+                identifiers,
                 token_count: 0,
                 symbol_name: frag.symbol_name,
             }
@@ -311,12 +312,13 @@ pub fn create_whole_file_fragment(
     let line_count = lines.len() as u32;
     let path_arc: Arc<str> = Arc::from(path.to_string_lossy().as_ref());
     let token_count = count_tokens(&content) + LIMITS.overhead_per_fragment;
+    let identifiers = extract_identifiers(&content, 2);
 
     Some(Fragment {
         id: FragmentId::new(path_arc, 1, line_count),
         kind: FragmentKind::Chunk,
-        content: content.clone(),
-        identifiers: extract_identifiers(&content, 2),
+        content: Arc::from(content),
+        identifiers,
         token_count,
         symbol_name: None,
     })
diff --git a/diffctx/src/git.rs b/diffctx/src/git.rs
@@ -89,25 +89,25 @@ fn wait_with_timeout(
     _args: &[&str],
 ) -> Result<std::process::Output> {
     let mut child = child;
-    let start = std::time::Instant::now();
+    let stdout_handle = child.stdout.take().map(|mut s| {
+        std::thread::spawn(move || -> std::io::Result<Vec<u8>> {
+            let mut buf = Vec::new();
+            s.read_to_end(&mut buf)?;
+            Ok(buf)
+        })
+    });
+    let stderr_handle = child.stderr.take().map(|mut s| {
+        std::thread::spawn(move || -> std::io::Result<Vec<u8>> {
+            let mut buf = Vec::new();
+            s.read_to_end(&mut buf)?;
+            Ok(buf)
+        })
+    });
 
-    loop {
+    let start = std::time::Instant::now();
+    let status = loop {
         match child.try_wait() {
-            Ok(Some(status)) => {
-                let mut stdout = Vec::new();
-                let mut stderr = Vec::new();
-                if let Some(mut out) = child.stdout.take() {
-                    let _ = out.read_to_end(&mut stdout);
-                }
-                if let Some(mut err) = child.stderr.take() {
-                    let _ = err.read_to_end(&mut stderr);
-                }
-                return Ok(std::process::Output {
-                    status,
-                    stdout,
-                    stderr,
-                });
-            }
+            Ok(Some(status)) => break status,
             Ok(None) => {
                 if start.elapsed() >= timeout {
                     let _ = child.kill();
@@ -118,7 +118,22 @@ fn wait_with_timeout(
             }
             Err(e) => return Err(GitError::Io(e)),
         }
-    }
+    };
+
+    let stdout = stdout_handle
+        .and_then(|h| h.join().ok())
+        .and_then(|r| r.ok())
+        .unwrap_or_default();
+    let stderr = stderr_handle
+        .and_then(|h| h.join().ok())
+        .and_then(|r| r.ok())
+        .unwrap_or_default();
+
+    Ok(std::process::Output {
+        status,
+        stdout,
+        stderr,
+    })
 }
 
 pub fn is_git_repo(path: &Path) -> bool {
diff --git a/diffctx/src/parsers/config_parser.rs b/diffctx/src/parsers/config_parser.rs
@@ -88,11 +88,12 @@ fn fragment_json(path: Arc<str>, content: &str) -> Vec<Fragment> {
             snippet.push('\n');
         }
 
+        let identifiers = extract_identifiers(&snippet, 2);
         fragments.push(Fragment {
             id: FragmentId::new(Arc::clone(&path), start_line, end_line),
             kind: FragmentKind::Chunk,
-            content: snippet.clone(),
-            identifiers: extract_identifiers(&snippet, 2),
+            content: Arc::from(snippet),
+            identifiers,
             token_count: 0,
             symbol_name: None,
         });
@@ -141,11 +142,12 @@ fn split_at_top_level_pattern(path: Arc<str>, content: &str, pattern: &Regex) ->
             snippet.push('\n');
         }
 
+        let identifiers = extract_identifiers(&snippet, 2);
         fragments.push(Fragment {
             id: FragmentId::new(Arc::clone(&path), start_line, end_line),
             kind: FragmentKind::Chunk,
-            content: snippet.clone(),
-            identifiers: extract_identifiers(&snippet, 2),
+            content: Arc::from(snippet),
+            identifiers,
             token_count: 0,
             symbol_name: None,
         });
@@ -176,11 +178,12 @@ fn make_single_fragment(path: Arc<str>, lines: &[&str]) -> Vec<Fragment> {
         snippet.push('\n');
     }
     let end_line = lines.len() as u32;
+    let identifiers = extract_identifiers(&snippet, 2);
     vec![Fragment {
         id: FragmentId::new(path, 1, end_line),
         kind: FragmentKind::Chunk,
-        content: snippet.clone(),
-        identifiers: extract_identifiers(&snippet, 2),
+        content: Arc::from(snippet),
+        identifiers,
         token_count: 0,
         symbol_name: None,
     }]
diff --git a/diffctx/src/parsers/generic.rs b/diffctx/src/parsers/generic.rs
@@ -33,12 +33,13 @@ impl FragmentationStrategy for GenericStrategy {
             if !snippet.ends_with('\n') {
                 snippet.push('\n');
             }
+            let identifiers = extract_identifiers(&snippet, 2);
 
             fragments.push(Fragment {
                 id: FragmentId::new(Arc::clone(&path), start_line, end_line),
                 kind: FragmentKind::Chunk,
-                content: snippet.clone(),
-                identifiers: extract_identifiers(&snippet, 2),
+                content: Arc::from(snippet),
+                identifiers,
                 token_count: 0,
                 symbol_name: None,
             });
diff --git a/diffctx/src/parsers/markdown.rs b/diffctx/src/parsers/markdown.rs
@@ -73,11 +73,12 @@ impl FragmentationStrategy for MarkdownStrategy {
                 snippet.push('\n');
             }
 
+            let identifiers = extract_identifiers(&snippet, 2);
             fragments.push(Fragment {
                 id: FragmentId::new(Arc::clone(&path), start_line, end_line),
                 kind: FragmentKind::Section,
-                content: snippet.clone(),
-                identifiers: extract_identifiers(&snippet, 2),
+                content: Arc::from(snippet),
+                identifiers,
                 token_count: 0,
                 symbol_name: None,
             });
diff --git a/diffctx/src/parsers/mod.rs b/diffctx/src/parsers/mod.rs
@@ -119,11 +119,12 @@ fn create_code_gap_fragments(
             continue;
         }
         if let Some(snippet) = create_snippet(lines, start, end) {
+            let identifiers = crate::types::extract_identifiers(&snippet, 2);
             fragments.push(Fragment {
                 id: crate::types::FragmentId::new(Arc::clone(&path), start, end),
                 kind: crate::types::FragmentKind::Chunk,
-                content: snippet.clone(),
-                identifiers: crate::types::extract_identifiers(&snippet, 2),
+                content: Arc::from(snippet),
+                identifiers,
                 token_count: 0,
                 symbol_name: None,
             });
diff --git a/diffctx/src/parsers/tree_sitter_strategy.rs b/diffctx/src/parsers/tree_sitter_strategy.rs
@@ -1178,11 +1178,12 @@ fn create_and_append_fragment(
         Some(s) => s,
         None => return false,
     };
+    let identifiers = extract_identifiers(&snippet, 2);
     fragments.push(Fragment {
         id: FragmentId::new(Arc::clone(path), start, end),
         kind: FragmentKind::from_str(kind),
-        content: snippet.clone(),
-        identifiers: extract_identifiers(&snippet, 2),
+        content: Arc::from(snippet),
+        identifiers,
         token_count: 0,
         symbol_name: sym_name.map(|s| s.to_string()),
     });
@@ -1318,11 +1319,12 @@ fn try_container_split(
     }
     let header_end = first_child_start - 1;
     if let Some(snippet) = create_snippet(lines, start, header_end) {
+        let identifiers = extract_identifiers(&snippet, 2);
         fragments.push(Fragment {
             id: FragmentId::new(Arc::clone(path), start, header_end),
             kind: FragmentKind::from_str(kind),
-            content: snippet.clone(),
-            identifiers: extract_identifiers(&snippet, 2),
+            content: Arc::from(snippet),
+            identifiers,
             token_count: 0,
             symbol_name: sym_name.map(|s| s.to_string()),
         });
diff --git a/diffctx/src/render.rs b/diffctx/src/render.rs
@@ -141,7 +141,7 @@ fn create_fragment_entry(frag: &Fragment, path_str: &str) -> FragmentEntry {
     let content = if frag.content.is_empty() {
         None
     } else {
-        Some(frag.content.clone())
+        Some(frag.content.to_string())
     };
 
     FragmentEntry {
diff --git a/diffctx/src/signatures.rs b/diffctx/src/signatures.rs
@@ -1,3 +1,5 @@
+use std::sync::Arc;
+
 use rustc_hash::FxHashSet;
 
 use crate::types::{Fragment, FragmentId, FragmentKind};
@@ -107,7 +109,7 @@ pub fn generate_signature_variants(fragments: &[Fragment]) -> Vec<Fragment> {
         signatures.push(Fragment {
             id: sig_id,
             kind: signature_kind(frag.kind),
-            content: sig_content,
+            content: Arc::from(sig_content),
             identifiers: frag.identifiers.clone(),
             token_count: 0,
             symbol_name: frag.symbol_name.clone(),
diff --git a/diffctx/src/types.rs b/diffctx/src/types.rs
@@ -208,7 +208,7 @@ impl FragmentId {
 pub struct Fragment {
     pub id: FragmentId,
     pub kind: FragmentKind,
-    pub content: String,
+    pub content: Arc<str>,
     pub identifiers: FxHashSet<String>,
     pub token_count: u32,
     pub symbol_name: Option<String>,
diff --git a/scripts/aggregate_sweep_report.py b/scripts/aggregate_sweep_report.py
diff --git a/scripts/run_bench_docker.sh b/scripts/run_bench_docker.sh
diff --git a/scripts/sweep_orchestrator.sh b/scripts/sweep_orchestrator.sh
