From 1b2ac752dbff3e00824526481b695c7169322cb4 Mon Sep 17 00:00:00 2001 From: Lucas <18381968+lucaspunz@users.noreply.github.com> Date: Wed, 26 Feb 2025 00:20:12 -0800 Subject: [PATCH] weak reference holding in index --- Makefile | 8 +++ src/database.rs | 86 +++++++++++++++++++++++++++--- src/index.rs | 136 ++++++++++++++++++++++++++++++++++++++++++++---- 3 files changed, 213 insertions(+), 17 deletions(-) diff --git a/Makefile b/Makefile index bb5ad8a..c60f358 100644 --- a/Makefile +++ b/Makefile @@ -11,3 +11,11 @@ dev: test: cargo test + +pyenv: + python3 -m venv venv + source venv/bin/activate + pip install -r requirements.txt + +format: + cargo fmt diff --git a/src/database.rs b/src/database.rs index c1a1eaf..d19e853 100644 --- a/src/database.rs +++ b/src/database.rs @@ -7,6 +7,7 @@ use std::collections::HashMap; use std::fs::File; use std::io::{BufReader, BufWriter, Write}; use std::path::Path; +use std::sync::{Arc, Weak}; #[derive(Serialize, Deserialize, Debug, Default)] pub struct RDatabaseMetadata { @@ -122,15 +123,17 @@ impl RDatabase { // Push t into the tables vector so its address becomes stable. self.tables.push(t); let i = self.tables.len() - 1; - // Get a raw pointer to the table in the vector. - let table_ptr = &self.tables[i] as *const RTable; - // Set the owner pointer in the index. - self.tables[i].index.set_owner(table_ptr); - // Map a name of a table to it's index + // Map a name of a table to its index self.tables_hashmap.insert(name, i); - // Should it really be cloning here? - // I guess since it has just an Arc Mutex, the underlying data should persi + { + // Now create an Arc> with the table reference + let arc_table = std::sync::Arc::new(std::sync::RwLock::new(self.tables[i].clone())); + + // Set the owner on the actual table in our vector + self.tables[i].index.set_owner(arc_table); + } + return self.tables[i].clone(); } @@ -200,4 +203,73 @@ mod tests { assert_eq!(db.tables.len(), 2); } + + #[test] + fn test_create_table_with_set_owner() { + let mut db = RDatabase::new(); + + // Create a table + let table = db.create_table(String::from("users"), 3, 0); + + // Verify table was added to the database + assert_eq!(db.tables.len(), 1); + assert!(db.tables_hashmap.contains_key("users")); + assert_eq!(*db.tables_hashmap.get("users").unwrap(), 0); + + // Verify owner is set correctly for the index + assert!(table.index.owner.is_some()); + if let Some(owner_weak) = &table.index.owner { + // Verify owner can be upgraded (reference is valid) + assert!(owner_weak.upgrade().is_some()); + + // Verify the table referenced by the index matches what we expect + if let Some(owner_arc) = owner_weak.upgrade() { + let referenced_table = owner_arc.read().unwrap(); + assert_eq!(referenced_table.name, "users"); + assert_eq!(referenced_table.num_columns, 3); + assert_eq!(referenced_table.primary_key_column, 0); + } + } + } + #[test] + fn test_drop_table_memory_management() { + let mut db = RDatabase::new(); + + // Create a table + let table = db.create_table(String::from("users_to_drop"), 3, 0); + + // Save a weak reference to the table's owner + let weak_ref = if let Some(owner) = &table.index.owner { + owner.clone() + } else { + panic!("Owner not set properly"); + }; + + // drop the table (so we aren't holding onto it manually when checking later) + drop(table); + + // Verify the weak reference can be upgraded before dropping + assert!(weak_ref.upgrade().is_some()); + + // Verify the table's properties before dropping + if let Some(owner_arc) = weak_ref.upgrade() { + let referenced_table = owner_arc.read().unwrap(); + assert_eq!(referenced_table.name, "users_to_drop"); + assert_eq!(referenced_table.num_columns, 3); + assert_eq!(referenced_table.primary_key_column, 0); + } + + // Drop the table + db.drop_table("users_to_drop".to_string()); + + // Verify the table is removed from the database + assert_eq!(db.tables.len(), 0); + assert!(!db.tables_hashmap.contains_key("users_to_drop")); + + // Verify the reference is no longer valid (table is fully dropped) + assert!( + weak_ref.upgrade().is_none(), + "Table wasn't properly dropped - Arc reference still exists" + ); + } } diff --git a/src/index.rs b/src/index.rs index 0d8888e..74b7e16 100644 --- a/src/index.rs +++ b/src/index.rs @@ -1,13 +1,22 @@ use super::table::RTable; use pyo3::prelude::*; -use std::collections::{BTreeMap, HashMap}; +use std::{ + collections::{BTreeMap, HashMap}, + sync::{Arc, RwLock, Weak}, +}; #[pyclass] #[derive(Clone, Default)] pub struct RIndex { + #[pyo3(get, set)] pub index: BTreeMap, + #[pyo3(get, set)] pub secondary_indices: HashMap>>, - owner: Option, // RTable owner's pointer + // Using Arc> pattern which is safer than raw pointers + // these fields are not python exposed + pub owner: Option>>, + // Keep strong reference to the table to prevent it from being dropped + pub table_ref: Option>>, } #[pymethods] @@ -19,10 +28,14 @@ impl RIndex { /// When called from Python, create the secondary index for a given column pub fn create_index(&mut self, col_index: i64) { - if let Some(ptr) = self.owner { - // SAFETY: We assume the owner lives as long as the index - let table: &crate::table::RTable = unsafe { &*(ptr as *const crate::table::RTable) }; - self.create_index_internal(col_index, table); + if let Some(owner_weak) = &self.owner { + if let Some(owner_arc) = owner_weak.upgrade() { + let table = owner_arc.read().unwrap(); + self.create_index_internal(col_index, &table); + } else { + // Table was dropped, so this index should be considered invalid + panic!("Table reference no longer valid"); + } } else { panic!("Owner not set for RIndex"); } @@ -52,13 +65,16 @@ impl RIndex { index: BTreeMap::new(), secondary_indices: HashMap::new(), owner: None, + table_ref: None, } } - /// Set the owner (the table that “owns” this index) - pub fn set_owner(&mut self, owner: *const RTable) { - // Must cast the owner reference to a raw pointer - self.owner = Some(owner as usize); + // Set the owner (the table that "owns" this index) + pub fn set_owner(&mut self, table_arc: Arc>) { + // Store the Arc directly in table_ref + self.table_ref = Some(table_arc.clone()); + // Generate weak reference when needed + self.owner = Some(Arc::downgrade(&table_arc)); } /// Create a mapping from primary_key to RID @@ -230,5 +246,105 @@ mod tests { index.drop_index_internal(2); assert!(index.secondary_indices.get(&2).is_none()); } + #[test] + fn test_set_owner() { + // Create a dummy table + let table = RTable { + name: "dummy".to_string(), + primary_key_column: 0, + page_range: PageRange::new(3), + page_directory: HashMap::new(), + num_records: 0, + num_columns: 3, + index: RIndex::new(), + }; + + let mut index = RIndex::new(); + + { + // Wrap the table in an Arc> + let table_arc = Arc::new(RwLock::new(table)); + + // Create an index and set the owner + index.set_owner(table_arc); + } + + // Verify the owner is set by checking it can be upgraded + assert!(index.owner.is_some()); + let owner_weak = index.owner.as_ref().unwrap(); + assert!(owner_weak.upgrade().is_some()); + } + + #[test] + fn test_secondary_index_insert() { + let mut index = RIndex::new(); + let col_index = 1; + + // Create an empty secondary index + index.secondary_indices.insert(col_index, BTreeMap::new()); + + // Insert a value + index.secondary_index_insert(col_index, 5, 100); + + // Verify the value was inserted + let sec_index = index.secondary_indices.get(&col_index).unwrap(); + assert_eq!(sec_index.get(&100).unwrap(), &vec![5]); + + // Insert another value with the same key + index.secondary_index_insert(col_index, 10, 100); + let sec_index = index.secondary_indices.get(&col_index).unwrap(); + assert_eq!(sec_index.get(&100).unwrap(), &vec![5, 10]); + } + + #[test] + fn test_secondary_index_update() { + let mut index = RIndex::new(); + let col_index = 1; + + // Create a secondary index with initial values + let mut btree = BTreeMap::new(); + btree.insert(100, vec![5, 10]); + btree.insert(200, vec![15]); + index.secondary_indices.insert(col_index, btree); + + // Update a value (move rid 10 from value 100 to 200) + index.secondary_index_update(col_index, 10, 100, 200); + + // Verify the update + let sec_index = index.secondary_indices.get(&col_index).unwrap(); + // Check that we have the right values in each index + let vec_100 = sec_index.get(&100).unwrap(); + let vec_200 = sec_index.get(&200).unwrap(); + + assert_eq!(vec_100.len(), 1); + assert!(vec_100.contains(&5)); + + assert_eq!(vec_200.len(), 2); + assert!(vec_200.contains(&10)); + assert!(vec_200.contains(&15)); + } + + #[test] + fn test_secondary_index_delete() { + let mut index = RIndex::new(); + let col_index = 1; + + // Create a secondary index with initial values + let mut btree = BTreeMap::new(); + btree.insert(100, vec![5, 10]); + index.secondary_indices.insert(col_index, btree); + + // Delete a value + index.secondary_index_delete(col_index, 10, 100); + + // Verify the deletion + let sec_index = index.secondary_indices.get(&col_index).unwrap(); + assert_eq!(sec_index.get(&100).unwrap(), &vec![5]); + + // Delete the last value + index.secondary_index_delete(col_index, 5, 100); + let sec_index = index.secondary_indices.get(&col_index).unwrap(); + assert!(sec_index.get(&100).unwrap().is_empty()); + } } }