lore: speed up indexing on many-CPU hosts

src/bin/index.rs

@@ -1,9 +1,9 @@
 // SPDX-License-Identifier: MIT OR Apache-2.0
-use anyhow::Result;
+use anyhow::{Context, Result};
 use clap::Parser;
 use colored::Colorize;
 use semcode::indexer::{
-    list_shas_in_range, process_commits_pipeline, process_lore_commits_pipeline,
+    list_shas_in_range, process_commits_pipeline, process_lore_commits_pipeline, KnownLoreIds,
 };
 use semcode::{measure, process_database_path, CodeVectorizer, DatabaseManager};
 // Temporary call relationships are now embedded in function JSON columns
@@ -13,6 +13,16 @@ use std::path::{Path, PathBuf};
 use std::sync::Arc;
 use tracing::{error, info, warn};
 
+/// Upper bound on tokio worker threads and lore parser threads.  One
+/// thread per CPU (300+ on large hosts) spends more time in futex and
+/// allocator contention than it gains: lore indexing is limited by the
+/// database inserters, not by parsing.  The other modes do their CPU
+/// work on rayon, std threads or spawn_blocking, not tokio workers.
+/// The exception is lance's FTS index build, which tokenizes on tokio
+/// workers, so a full lore FTS (re)build is capped too; set
+/// TOKIO_WORKER_THREADS to lift the cap.
+const MAX_THREADS: usize = 32;
+
 #[derive(Parser, Debug)]
 #[command(name = "semcode-index")]
 #[command(about = "Index a codebase for semantic analysis", long_about = None)]
@@ -682,6 +692,8 @@ struct LoreIndexResult {
 
 /// Index a lore archive repository, filtering out already-indexed commits.
 /// This is the common logic shared by both --lore and --refresh-lore handlers.
+/// Every archive of one run shares `known_ids`.
+#[allow(clippy::too_many_arguments)]
 async fn index_lore_archive(
     lore_repo: gix::Repository,
     archive_path: &Path,
@@ -690,6 +702,7 @@ async fn index_lore_archive(
     batch_size: usize,
     num_workers: usize,
     db_threads: usize,
+    known_ids: &KnownLoreIds,
 ) -> Result<LoreIndexResult> {
     // Get the remote tracking branch to include all fetched commits
     // Try refs/remotes/origin/master first, then refs/remotes/origin/main as fallback
@@ -764,6 +777,7 @@ async fn index_lore_archive(
         batch_size,
         num_workers,
         db_threads,
+        known_ids,
     )
     .await?;
 
@@ -1008,8 +1022,21 @@ async fn run_branch_indexing(args: Args, branches: Vec<semcode::git::BranchRef>)
 
 // ==================== End Branch Indexing Support ====================
 
-#[tokio::main]
-async fn main() -> Result<()> {
+fn main() -> Result<()> {
+    // #[tokio::main] starts one worker per CPU.  Cap it, but let
+    // TOKIO_WORKER_THREADS override as it does for the default runtime.
+    let mut runtime = tokio::runtime::Builder::new_multi_thread();
+    runtime.enable_all();
+    if std::env::var_os("TOKIO_WORKER_THREADS").is_none() {
+        runtime.worker_threads(num_cpus::get().min(MAX_THREADS));
+    }
+    runtime
+        .build()
+        .context("failed to build the tokio runtime")?
+        .block_on(async_main())
+}
+
+async fn async_main() -> Result<()> {
     // Suppress ORT verbose logging
     std::env::set_var("ORT_LOG_LEVEL", "ERROR");
 
@@ -1112,11 +1139,12 @@ async fn main() -> Result<()> {
 
             let start_time = std::time::Instant::now();
             let batch_size = 1024;
-            let num_workers = num_cpus::get();
+            let num_workers = num_cpus::get().min(MAX_THREADS);
 
             let mut total_new_emails = 0usize;
             let mut total_emails_all_archives = 0usize;
             let mut failed_archives: Vec<(String, String)> = Vec::new();
+            let known_ids = KnownLoreIds::new();
 
             // Process each lore archive
             for lore_url in lore_args {
@@ -1153,6 +1181,7 @@ async fn main() -> Result<()> {
                     batch_size,
                     num_workers,
                     args.db_threads,
+                    &known_ids,
                 )
                 .await;
 
@@ -1170,6 +1199,8 @@ async fn main() -> Result<()> {
                     }
                 }
             }
+            // One string per stored email: free it before compaction.
+            drop(known_ids);
 
             let total_time = start_time.elapsed();
 
@@ -1290,7 +1321,7 @@ async fn main() -> Result<()> {
 
             let start_time = std::time::Instant::now();
             let batch_size = 1024;
-            let num_workers = num_cpus::get();
+            let num_workers = num_cpus::get().min(MAX_THREADS);
             let db_threads = args.db_threads;
             let total_archives = archives_with_names.len();
 
@@ -1300,6 +1331,7 @@ async fn main() -> Result<()> {
             // causing both to stall indefinitely.
             let mut results: Vec<Result<(String, LoreIndexResult), (String, anyhow::Error)>> =
                 Vec::with_capacity(total_archives);
+            let known_ids = KnownLoreIds::new();
 
             for (archive_path, display_name) in archives_with_names {
                 println!("\n=== Refreshing lore archive: {} ===", display_name);
@@ -1321,6 +1353,7 @@ async fn main() -> Result<()> {
                     batch_size,
                     num_workers,
                     db_threads,
+                    &known_ids,
                 )
                 .await;
 
@@ -1331,6 +1364,8 @@ async fn main() -> Result<()> {
                     Err(e) => results.push(Err((display_name, e))),
                 }
             }
+            // One string per stored email: free it before compaction.
+            drop(known_ids);
 
             // Aggregate results
             let mut total_new_emails = 0usize;

src/database/connection.rs

@@ -150,6 +150,16 @@ fn log_absence(kind: &str, name: &str, git_sha: &str, why: Absent, elsewhere: &s
 /// moves the generation and drops every entry).
 type CallerMemoKey = (String, String, u64);
 
+/// Where a message_id stands in the duplicate check that the lore
+/// inserters share (see DatabaseManager::insert_lore_emails()).
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum LoreIdState {
+    /// The email is in the lore table.
+    Stored,
+    /// A batch has claimed the id and its append has not finished.
+    Pending,
+}
+
 pub struct DatabaseManager {
     connection: Connection,
     git_repo_path: String,
@@ -6574,11 +6584,22 @@ impl DatabaseManager {
     /// any emails that could not be stored.  Callers must not record
     /// commit SHAs for failed indices in the lore_indexed_commits
     /// table, otherwise those emails are permanently lost.
+    ///
+    /// `known_ids` holds every message_id already in the lore table
+    /// (see get_lore_message_ids()) and is shared by all inserters.
+    /// Emails whose message_id is stored are skipped.  New ids are
+    /// claimed as pending before the append, then marked stored, or
+    /// released if their email failed.  An email whose id another
+    /// inserter has pending is deferred: it is returned with the failed
+    /// indices, so its SHA is not recorded and the next run retries it
+    /// against the table.
     pub async fn insert_lore_emails(
         &self,
         emails: &[crate::types::LoreEmailInfo],
+        known_ids: &std::sync::Mutex<HashMap<String, LoreIdState>>,
     ) -> Result<Vec<usize>> {
         use arrow::datatypes::{DataType, Field, Schema};
+        use std::collections::hash_map::Entry;
         use std::sync::Arc;
 
         if emails.is_empty() {
@@ -6629,35 +6650,66 @@ impl DatabaseManager {
         let mut failed_indices: Vec<usize> = Vec::new();
 
         // Filter out emails whose message_id already exists in the
-        // table.  Normally every email here is new because
-        // index_lore_archive skips already-indexed commits, but a
-        // partial failure on a previous run can leave rows in the
-        // table whose commit SHA was never recorded in
-        // lore_indexed_commits.  Filtering avoids duplicates without
-        // relying on merge_insert, which hits a spurious null-column
-        // error in lance-file 3.0.1.
-        let new_indices = Self::filter_existing_lore_ids(&table, emails, &dedup_indices).await?;
-
-        if new_indices.len() < dedup_indices.len() {
+        // table.  A crossposted email has the same message_id in
+        // every list's archive, and a partial failure on a previous
+        // run can leave rows in the table whose commit SHA was never
+        // recorded in lore_indexed_commits.  Filtering avoids
+        // duplicates without relying on merge_insert, which hits a
+        // spurious null-column error in lance-file 3.0.1.
+        //
+        // The check runs against an in-memory set rather than a
+        // `message_id IN (...)` query: lance plans that predicate once
+        // per fragment, and a filtered query costs a scan of every
+        // fragment appended since the last index update.  Crossposts
+        // are caught because one set, loaded from the table once per
+        // run, is shared by the inserters of every archive in the run.
+        // An id is claimed under the lock so that two batches holding
+        // the same message_id cannot both store it.
+        let mut new_indices: Vec<usize> = Vec::new();
+        let mut deferred: Vec<usize> = Vec::new();
+        {
+            let mut known = known_ids.lock().unwrap_or_else(|e| e.into_inner());
+            for &i in &dedup_indices {
+                match known.entry(emails[i].message_id.clone()) {
+                    Entry::Vacant(slot) => {
+                        slot.insert(LoreIdState::Pending);
+                        new_indices.push(i);
+                    }
+                    Entry::Occupied(slot) if *slot.get() == LoreIdState::Pending => {
+                        deferred.push(i);
+                    }
+                    Entry::Occupied(_) => {}
+                }
+            }
+        }
+
+        let skipped = dedup_indices.len() - new_indices.len() - deferred.len();
+        if skipped > 0 {
             tracing::info!(
                 "insert_lore_emails: {} of {} already in table, inserting {}",
-                dedup_indices.len() - new_indices.len(),
+                skipped,
                 dedup_indices.len(),
                 new_indices.len(),
             );
         }
+        if !deferred.is_empty() {
+            tracing::info!(
+                "insert_lore_emails: deferring {} emails another batch is storing",
+                deferred.len()
+            );
+        }
 
         if new_indices.is_empty() {
             tracing::info!("insert_lore_emails: no new emails to insert");
-            return Ok(failed_indices);
+            return Ok(deferred);
         }
 
         // Use table.add() instead of merge_insert.  merge_insert in
         // lance-file 3.0.1 introduces spurious nulls during its
         // internal join/write phase, causing every insert to fail
         // with "subject contained null values".  Plain append avoids
         // the merge codepath entirely.  Duplicates are prevented by
-        // the filter_existing_lore_ids check above.
+        // the known_ids check above.
         if let Err(e) = Self::add_lore_chunk(&table, emails, &new_indices, &schema).await {
             tracing::warn!(
                 "insert_lore_emails: full batch of {} failed ({}), \
@@ -6700,60 +6752,59 @@ impl DatabaseManager {
             );
         }
 
+        // Mark the claimed ids stored, and release those whose email
+        // failed so that another copy of it can still be stored.
+        {
+            let mut known = known_ids.lock().unwrap_or_else(|e| e.into_inner());
+            for &idx in &new_indices {
+                if let Some(state) = known.get_mut(&emails[idx].message_id) {
+                    *state = LoreIdState::Stored;
+                }
+            }
+            for &idx in &failed_indices {
+                known.remove(&emails[idx].message_id);
+            }
+        }
+
         tracing::info!("insert_lore_emails: Batch insertion complete");
 
+        failed_indices.extend(deferred);
         Ok(failed_indices)
     }
 
-    /// Return the subset of `indices` whose message_id does not
-    /// already exist in the lore table.
-    async fn filter_existing_lore_ids(
-        table: &lancedb::Table,
-        emails: &[crate::types::LoreEmailInfo],
-        indices: &[usize],
-    ) -> Result<Vec<usize>> {
-        use futures::TryStreamExt as _;
-
-        // Build an IN-list predicate.  message_ids are already
-        // validated to be non-empty, but escape single quotes for
-        // the SQL literal.
-        let id_list: Vec<String> = indices
-            .iter()
-            .map(|&i| format!("'{}'", emails[i].message_id.replace('\'', "''")))
-            .collect();
-
-        let predicate = format!("message_id IN ({})", id_list.join(", "));
+    /// Return every message_id in the lore table, marked stored, for the
+    /// duplicate check in insert_lore_emails().  One unfiltered column
+    /// scan is far cheaper than a filtered query per batch.
+    pub async fn get_lore_message_ids(&self) -> Result<HashMap<String, LoreIdState>> {
+        let table = self.connection.open_table("lore").execute().await?;
 
-        let stream = table
+        let mut stream = table
             .query()
             .select(lancedb::query::Select::Columns(vec![
                 "message_id".to_string()
             ]))
-            .only_if(&predicate)
             .execute()
             .await?;
-        let batches: Vec<_> = stream.try_collect().await?;
 
-        let mut existing = std::collections::HashSet::new();
-        for batch in &batches {
+        // Consume batch by batch, so the Arrow copy of the column is
+        // never held in memory next to the map.
+        let mut existing = HashMap::new();
+        while let Some(batch) = stream.try_next().await? {
             if let Some(col) = batch.column_by_name("message_id") {
                 let arr = col
                     .as_any()
-                    .downcast_ref::<arrow::array::StringArray>()
+                    .downcast_ref::<StringArray>()
                     .expect("message_id column must be StringArray");
+                existing.reserve(arr.len());
                 for i in 0..arr.len() {
                     if !arr.is_null(i) {
-                        existing.insert(arr.value(i).to_string());
+                        existing.insert(arr.value(i).to_string(), LoreIdState::Stored);
                     }
                 }
             }
         }
 
-        Ok(indices
-            .iter()
-            .copied()
-            .filter(|&i| !existing.contains(&emails[i].message_id))
-            .collect())
+        Ok(existing)
     }
 
     /// Build a [`RecordBatch`] from the given email indices and
@@ -8543,6 +8594,12 @@ impl DatabaseManager {
     }
 
     /// Record git commit SHAs that have been processed for lore indexing.
+    ///
+    /// This is a plain append, not a merge_insert.  The callers only pass
+    /// SHAs that were filtered against get_indexed_lore_commits(), and
+    /// readers treat the column as a set, so a rare duplicate row is
+    /// harmless.  merge_insert on this unindexed table hash-joins
+    /// against every existing row, so its cost grows with the table.
     pub async fn insert_lore_indexed_commits(&self, commit_shas: &[String]) -> Result<()> {
         use arrow::array::{ArrayRef, StringArray};
         use arrow::datatypes::{DataType, Field, Schema};
@@ -8561,21 +8618,14 @@ impl DatabaseManager {
 
         let columns: Vec<ArrayRef> = vec![Arc::new(StringArray::from(commit_shas.to_vec()))];
 
-        let batch = RecordBatch::try_new(schema.clone(), columns)?;
-        let batches = vec![Ok(batch)];
-        let batch_iterator =
-            arrow::record_batch::RecordBatchIterator::new(batches.into_iter(), schema);
+        let batch = RecordBatch::try_new(schema, columns)?;
 
         let table = self
             .connection
             .open_table("lore_indexed_commits")
             .execute()
             .await?;
-        let mut merge_insert = table.merge_insert(&["git_commit_sha"]);
-        merge_insert
-            .when_matched_update_all(None)
-            .when_not_matched_insert_all();
-        merge_insert.execute(Box::new(batch_iterator)).await?;
+        table.add(vec![batch]).execute().await?;
 
         Ok(())
     }

src/database/mod.rs

@@ -21,7 +21,7 @@ mod symbol_filename;
 mod types;
 mod vectors;
 
-pub use connection::DatabaseManager;
+pub use connection::{DatabaseManager, LoreIdState};
 
 use anyhow::Result;
 use arrow::array::RecordBatch;