symbol_filename: serialize upserts

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<Mutex<()>> = 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"
+        );
+    }
+}