diff --git a/Cargo.lock b/Cargo.lock index 3a5c35a..d33c70f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -313,6 +313,7 @@ dependencies = [ "ignore", "indicatif", "predicates", + "rayon", "regex", "reqwest", "rusqlite", @@ -1400,6 +1401,26 @@ dependencies = [ "rand_core", ] +[[package]] +name = "rayon" +version = "1.12.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fb39b166781f92d482534ef4b4b1b2568f42613b53e5b6c160e24cfbfa30926d" +dependencies = [ + "either", + "rayon-core", +] + +[[package]] +name = "rayon-core" +version = "1.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "22e18b0f0062d30d4230b2e85ff77fdfe4326feb054b9783a3460d8435c8ab91" +dependencies = [ + "crossbeam-deque", + "crossbeam-utils", +] + [[package]] name = "redox_users" version = "0.5.2" diff --git a/Cargo.toml b/Cargo.toml index 6308393..fbdd74a 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -45,6 +45,9 @@ dirs = "6" glob = "0.3" ignore = "0.4" +# Parallelism (cold-rebuild file parsing) +rayon = "1" + # Error handling anyhow = "1" which = "7" diff --git a/src/index/brain.rs b/src/index/brain.rs index e1bbfc5..db9d3b2 100644 --- a/src/index/brain.rs +++ b/src/index/brain.rs @@ -7,6 +7,7 @@ use crate::embed::tokens::embed_code; use crate::index::symbols::SymbolQuery; use crate::index::vector::{CodeVectorIndex, DEFAULT_DIMS, cosine_distance_to_similarity}; use anyhow::{Context, Result}; +use rayon::prelude::*; use rusqlite::Connection; use std::collections::{HashMap, HashSet}; use std::sync::{LazyLock, RwLock}; @@ -78,23 +79,41 @@ pub fn embed_project(conn: &Connection, project_id: i64) -> Result { .filter_map(|r| r.ok()) .collect(); + // ── Parallel embedding computation (Rayon) ───────────────────────── + // embed_code is pure + CPU-bound. usearch insert is serial. + let t_compute = std::time::Instant::now(); + let embedded: Vec<(i64, Vec)> = rows + .par_iter() + .map(|(sym_id, name, _kind, signature)| { + let text = if signature.is_empty() || signature == name { + name.clone() + } else { + format!("{name} {signature}") + }; + let embedding = embed_code(&text); + let vec: Vec = embedding.as_slice().iter().map(|&v| v as f32).collect(); + (*sym_id, vec) + }) + .collect(); + let compute_ms = t_compute.elapsed().as_millis(); + + // ── Serial usearch insert ──────────────────────────────────────── + let t_insert = std::time::Instant::now(); let mut count = 0; - let mut new_ids: HashSet = HashSet::with_capacity(rows.len()); - for (sym_id, name, _kind, signature) in &rows { - let text = if signature.is_empty() || signature == name { - name.clone() - } else { - format!("{name} {signature}") - }; - - let embedding = embed_code(&text); - let vec: Vec = embedding.as_slice().iter().map(|&v| v as f32).collect(); - - vi.insert(*sym_id, &vec) - .context("insert symbol embedding")?; + let mut new_ids: HashSet = HashSet::with_capacity(embedded.len()); + for (sym_id, vec) in &embedded { + vi.insert(*sym_id, vec).context("insert symbol embedding")?; new_ids.insert(*sym_id); count += 1; } + let insert_ms = t_insert.elapsed().as_millis(); + + tracing::debug!( + "embed_compute={}ms, usearch_insert={}ms, symbols={}", + compute_ms, + insert_ms, + count + ); if vi.is_dirty() { vi.save().context("save vector index")?; diff --git a/src/index/mod.rs b/src/index/mod.rs index 31801a2..4d9f1aa 100644 --- a/src/index/mod.rs +++ b/src/index/mod.rs @@ -16,6 +16,7 @@ pub mod vector; use std::collections::HashMap; use std::path::Path; +use rayon::prelude::*; use rusqlite::Connection; #[cfg(test)] use sha2::{Digest, Sha256}; @@ -322,21 +323,37 @@ fn index_project_with_id( } if !files_to_index.is_empty() { - // Single transaction for ALL files — eliminates per-file fsync + // ── Phase 1: Parallel extraction (CPU-bound tree-sitter + regex) ─── + // Rayon par_iter distributes file parsing across all CPU cores. + // extract_all is CPU-bound (~4ms/file) and fully thread-safe. + let t_extract = std::time::Instant::now(); + let extracted_files: Vec<(String, extract::ExtractedAll, String, String)> = files_to_index + .par_iter() + .map(|(rel_str, content, language, cheap_fp)| { + let extracted = extract::extract_all(content, language, rel_str); + ( + rel_str.clone(), + extracted, + language.clone(), + cheap_fp.clone(), + ) + }) + .collect(); + let extract_ms = t_extract.elapsed().as_millis(); + + // ── Phase 2: Serial SQLite writes (single transaction) ───────── + let t_db = std::time::Instant::now(); let tx = conn.unchecked_transaction()?; // Disable FTS5 triggers during bulk insert to avoid 3x write amplification. - // For a content-sync FTS5 table (content='symbols'), the correct bulk-load - // sequence is: drop triggers → insert rows → 'rebuild' command → recreate triggers. tx.execute_batch( "DROP TRIGGER IF EXISTS symbols_fts_insert;\ DROP TRIGGER IF EXISTS symbols_fts_delete;\ DROP TRIGGER IF EXISTS symbols_fts_update;", )?; - for (rel_str, content, language, cheap_fp) in &files_to_index { - let extracted = extract::extract_all(content, language, rel_str); - match index_file_in_tx(&tx, project_id, rel_str, cheap_fp, language, &extracted) { + for (rel_str, extracted, language, cheap_fp) in &extracted_files { + match index_file_in_tx(&tx, project_id, rel_str, cheap_fp, language, extracted) { Ok(n) => { stats.files_indexed += 1; stats.symbols_indexed += n; @@ -382,6 +399,12 @@ fn index_project_with_id( )?; tx.commit()?; + tracing::debug!( + "extract={}ms (rayon), db={}ms, files={}", + extract_ms, + t_db.elapsed().as_millis(), + files_to_index.len() + ); } // Update project's last_indexed timestamp