diff --git a/src/database/symbol_filename.rs b/src/database/symbol_filename.rs index fc1b74c..25d4b4e 100644 --- a/src/database/symbol_filename.rs +++ b/src/database/symbol_filename.rs @@ -4,10 +4,22 @@ use arrow::array::{ArrayRef, RecordBatch, RecordBatchIterator, StringBuilder}; use futures::TryStreamExt; use lancedb::connection::Connection; use lancedb::query::{ExecutableQuery, QueryBase}; +use once_cell::sync::Lazy; use std::sync::Arc; +use tokio::sync::Mutex; use crate::database::connection::OPTIMAL_BATCH_SIZE; +/// Process-global lock serializing symbol_filename upserts. +/// +/// One batch's functions and types inserts both upsert into this table +/// concurrently, and several inserters run at once. Two merge_inserts of the +/// same new (symbol, filename) pair both open the table at the same version, +/// both miss, and both insert, so the pair keeps two rows. The lock is +/// process-global rather than per instance so that it also serializes stores +/// belonging to different DatabaseManagers opened on one database. +static SYMBOL_FILENAME_LOCK: Lazy> = Lazy::new(|| Mutex::new(())); + /// Stores (symbol_name, filename) pairs for fast lookups of functions, macros, types, and typedefs pub struct SymbolFilenameStore { connection: Connection, @@ -24,6 +36,10 @@ impl SymbolFilenameStore { return Ok(()); } + // Taken before open_table and held across every chunk: a handle + // opened earlier reads a version without the other writer's rows, + // and its merge_insert inserts them again. + let _guard = SYMBOL_FILENAME_LOCK.lock().await; let table = self .connection .open_table("symbol_filename") @@ -184,3 +200,61 @@ impl SymbolFilenameStore { Ok(pairs) } } + +#[cfg(test)] +mod tests { + use super::*; + use lancedb::connection::LanceFileVersion; + use lancedb::database::listing::ListingDatabaseOptions; + + async fn fresh_store() -> (tempfile::TempDir, SymbolFilenameStore) { + let dir = tempfile::tempdir().unwrap(); + let db_path = dir.path().join("symbol-filename-race.lancedb"); + let options = ListingDatabaseOptions::builder() + .data_storage_version(LanceFileVersion::V2_2) + .build(); + let connection = lancedb::connect(db_path.to_str().unwrap()) + .database_options(&options) + .execute() + .await + .unwrap(); + crate::database::schema::SchemaManager::new(connection.clone()) + .create_all_tables() + .await + .unwrap(); + (dir, SymbolFilenameStore::new(connection)) + } + + /// Two concurrent insert_batch calls with the same new pairs must not + /// leave duplicate rows: without the lock both merge_inserts miss and + /// both insert. Fresh pairs every round so one lucky round cannot hide + /// the race. + #[tokio::test] + async fn concurrent_upserts_of_same_pairs_leave_no_duplicates() { + const ROUNDS: u32 = 40; + const PER_ROUND: u32 = 64; + + let (_dir, store) = fresh_store().await; + let store = Arc::new(store); + let mut unique_pairs = std::collections::HashSet::new(); + + for round in 0..ROUNDS { + let pairs: Vec<(String, String)> = (0..PER_ROUND) + .map(|item| (format!("sym_{round}_{item}"), format!("dir/file_{item}.c"))) + .collect(); + unique_pairs.extend(pairs.iter().cloned()); + + let (first, second) = + tokio::join!(store.insert_batch(pairs.clone()), store.insert_batch(pairs)); + first.unwrap(); + second.unwrap(); + } + + let rows = store.get_all().await.unwrap(); + assert_eq!( + rows.len(), + unique_pairs.len(), + "concurrent upserts of the same new pairs left duplicate rows" + ); + } +}