fix tabledb race condition

veilid-core/src/table_store/mod.rs

@@ -15,6 +15,7 @@ mod native;
 use native::*;
 
 use keyvaluedb::*;
+use weak_table::WeakValueHashMap;
 
 impl_veilid_log_facility!("tstore");
 
@@ -72,7 +73,7 @@ pub struct TableInfo {
 
 #[must_use]
 struct TableStoreInner {
-    opened: BTreeMap<String, Weak<TableDBUnlockedInner>>,
+    opened: WeakValueHashMap<String, Weak<TableDBUnlockedInner>>,
     encryption_key: Option<TypedSharedSecret>,
     all_table_names: HashMap<String, String>,
     all_tables_db: Option<Database>,
@@ -115,7 +116,7 @@ impl_veilid_component!(TableStore);
 impl TableStore {
     fn new_inner() -> TableStoreInner {
         TableStoreInner {
-            opened: BTreeMap::new(),
+            opened: WeakValueHashMap::new(),
             encryption_key: None,
             all_table_names: HashMap::new(),
             all_tables_db: None,
@@ -528,6 +529,7 @@ impl TableStore {
         self.flush().await;
 
         let mut inner = self.inner.lock();
+        inner.opened.shrink_to_fit();
         if !inner.opened.is_empty() {
             panic!(
                 "all open databases should have been closed: {:?}",
@@ -539,15 +541,6 @@ impl TableStore {
         inner.encryption_key = None;
     }
 
-    #[instrument(level = "trace", target = "tstore", skip_all)]
-    pub(crate) fn on_table_db_drop(&self, table: String) {
-        veilid_log!(self trace "dropping table db: {}", table);
-        let mut inner = self.inner.lock();
-        if inner.opened.remove(&table).is_none() {
-            unreachable!("should have removed an item");
-        }
-    }
-
     /// Get or create a TableDB database table. If the column count is greater than an
     /// existing TableDB's column count, the database will be upgraded to add the missing columns.
     #[instrument(level = "trace", target = "tstore", skip_all)]
@@ -566,10 +559,10 @@ impl TableStore {
 
         // See if this table is already opened, if so the column count must be the same
         {
-            let mut inner = self.inner.lock();
-            if let Some(table_db_weak_inner) = inner.opened.get(&table_name) {
-                match TableDB::try_new_from_weak_inner(table_db_weak_inner.clone(), column_count) {
-                    Some(tdb) => {
+            let inner = self.inner.lock();
+            if let Some(table_db_unlocked_inner) = inner.opened.get(&table_name) {
+                let tdb = TableDB::new_from_unlocked_inner(table_db_unlocked_inner, column_count);
+
                 // Ensure column count isnt bigger
                 let existing_col_count = tdb.get_column_count()?;
                 if column_count > existing_col_count {
@@ -581,11 +574,6 @@ impl TableStore {
 
                 return Ok(tdb);
             }
-                    None => {
-                        inner.opened.remove(&table_name);
-                    }
-                };
-            }
         }
 
         // Open table db using platform-specific driver
@@ -637,7 +625,7 @@ impl TableStore {
         // Keep track of opened DBs
         inner
             .opened
-            .insert(table_name.clone(), table_db.weak_unlocked_inner());
+            .insert(table_name.clone(), table_db.unlocked_inner());
 
         Ok(table_db)
     }

veilid-core/src/table_store/table_db.rs

@@ -38,13 +38,6 @@ impl fmt::Debug for TableDBUnlockedInner {
     }
 }
 
-impl Drop for TableDBUnlockedInner {
-    fn drop(&mut self) {
-        let table_store = self.registry.table_store();
-        table_store.on_table_db_drop(self.table.clone());
-    }
-}
-
 #[derive(Debug, Clone)]
 #[must_use]
 pub struct TableDB {
@@ -87,12 +80,11 @@ impl TableDB {
         }
     }
 
-    pub(super) fn try_new_from_weak_inner(
-        weak_inner: Weak<TableDBUnlockedInner>,
+    pub(super) fn new_from_unlocked_inner(
+        unlocked_inner: Arc<TableDBUnlockedInner>,
         opened_column_count: u32,
-    ) -> Option<Self> {
-        weak_inner.upgrade().map(|table_db_unlocked_inner| {
-            let db = &table_db_unlocked_inner.database;
+    ) -> Self {
+        let db = &unlocked_inner.database;
         let total_columns = db.num_columns().unwrap();
         Self {
             opened_column_count: if opened_column_count == 0 {
@@ -100,13 +92,12 @@ impl TableDB {
             } else {
                 opened_column_count
             },
-                unlocked_inner: table_db_unlocked_inner,
+            unlocked_inner,
         }
-        })
     }
 
-    pub(super) fn weak_unlocked_inner(&self) -> Weak<TableDBUnlockedInner> {
-        Arc::downgrade(&self.unlocked_inner)
+    pub(super) fn unlocked_inner(&self) -> Arc<TableDBUnlockedInner> {
+        self.unlocked_inner.clone()
     }
 
     /// Get the internal name of the table