diff --git a/db4-graph/src/lib.rs b/db4-graph/src/lib.rs index 38a01aa453..18592e8b55 100644 --- a/db4-graph/src/lib.rs +++ b/db4-graph/src/lib.rs @@ -456,18 +456,29 @@ where self.graph } + /// Write-lock the node and edge segments created since this graph was locked + /// (e.g. by node or edge resolution), keeping the locks already held. + pub fn lock_new_segments(&mut self) { + self.graph + .storage + .nodes() + .lock_new_segments(&mut self.nodes); + self.graph + .storage + .edges() + .lock_new_segments(&mut self.edges); + } + pub fn resize_segments_to_vid(&mut self, vid: VID) { let (segment_id, _) = self.graph.storage.nodes().resolve_pos(vid); self.graph.storage().nodes().grow(segment_id + 1); - std::mem::take(&mut self.nodes); - self.nodes = self.graph.storage.nodes().write_locked(); + self.lock_new_segments(); } pub fn resize_segments_to_eid(&mut self, eid: EID) { let (segment_id, _) = self.graph.storage.edges().resolve_pos(eid); self.graph.storage().edges().grow(segment_id + 1); - std::mem::take(&mut self.edges); - self.edges = self.graph.storage.edges().write_locked(); + self.lock_new_segments(); } pub fn edge_stats(&self) -> &Arc { diff --git a/db4-storage/src/pages/edge_store.rs b/db4-storage/src/pages/edge_store.rs index 492a19563f..7806dfd2e1 100644 --- a/db4-storage/src/pages/edge_store.rs +++ b/db4-storage/src/pages/edge_store.rs @@ -468,6 +468,21 @@ impl, EXT: PersistenceStrategy> ) } + /// Write-lock the segments created since `pages` was taken and append them, + /// keeping the locks `pages` already holds. + pub fn lock_new_segments<'a>(&'a self, pages: &mut WriteLockedEdgePages<'a, ES>) { + for segment_id in pages.len()..self.segments.count() { + let page = self.get_or_create_segment(segment_id); + pages.push(LockedEdgePage::new( + segment_id, + self.max_page_len(), + page.as_ref(), + &self.layer_counter, + page.head_mut(), + )); + } + } + /// Retrieve the segment for an edge given its EID pub fn get_edge_segment(&self, eid: EID) -> Option<&Arc> { let (segment_id, _) = resolve_pos(eid, self.max_page_len()); diff --git a/db4-storage/src/pages/locked/edges.rs b/db4-storage/src/pages/locked/edges.rs index dd1cdd2cb1..adb92cbd56 100644 --- a/db4-storage/src/pages/locked/edges.rs +++ b/db4-storage/src/pages/locked/edges.rs @@ -145,6 +145,12 @@ impl<'a, EXT: PersistenceStrategy, ES: EdgeSegmentOps> self.writers.len() } + /// Append the page for the next segment; pages are indexed by segment id. + pub fn push(&mut self, page: LockedEdgePage<'a, ES>) { + assert_eq!(page.page_id(), self.writers.len()); + self.writers.push(page); + } + pub fn is_empty(&self) -> bool { self.writers.is_empty() } diff --git a/db4-storage/src/pages/locked/nodes.rs b/db4-storage/src/pages/locked/nodes.rs index d56ea0e648..27d8a93db0 100644 --- a/db4-storage/src/pages/locked/nodes.rs +++ b/db4-storage/src/pages/locked/nodes.rs @@ -110,6 +110,12 @@ impl<'a, EXT: PersistenceStrategy, NS: NodeSegmentOps> self.writers.len() } + /// Append the page for the next segment; pages are indexed by segment id. + pub fn push(&mut self, page: LockedNodePage<'a, NS>) { + assert_eq!(page.segment_id(), self.writers.len()); + self.writers.push(page); + } + #[inline] pub fn get_mut(&mut self, segment_id: usize) -> Option<&mut LockedNodePage<'a, NS>> { self.writers.get_mut(segment_id) diff --git a/db4-storage/src/pages/node_store.rs b/db4-storage/src/pages/node_store.rs index 85749ef5bf..4d30a5282a 100644 --- a/db4-storage/src/pages/node_store.rs +++ b/db4-storage/src/pages/node_store.rs @@ -327,6 +327,21 @@ impl, EXT: PersistenceStrategy> ) } + /// Write-lock the segments created since `pages` was taken and append them, + /// keeping the locks `pages` already holds. + pub fn lock_new_segments<'a>(&'a self, pages: &mut WriteLockedNodePages<'a, NS>) { + for segment_id in pages.len()..self.segments.count() { + let page = self.get_or_create_segment(segment_id); + pages.push(LockedNodePage::new( + segment_id, + &self.stats, + self.max_segment_len(), + page.as_ref(), + page.head_mut(), + )); + } + } + pub fn reserve_vid(&self, row: usize) -> VID { let (seg, pos) = self.reserve_free_pos(row); pos.as_vid(seg, self.max_segment_len()) diff --git a/raphtory-tests/tests/io_tests/df_loaders.rs b/raphtory-tests/tests/io_tests/df_loaders.rs index 4f0e362576..7ad4947fca 100644 --- a/raphtory-tests/tests/io_tests/df_loaders.rs +++ b/raphtory-tests/tests/io_tests/df_loaders.rs @@ -987,6 +987,90 @@ mod io_tests { assert!(!g.has_node(9u64)); assert_eq!(g.node(1u64).unwrap().metadata().get("str_prop"), None); } + fn small_pages_graph() -> Graph { + Graph::new_with_config( + Args::default() + .with_max_node_page_len(2) + .with_max_edge_page_len(2), + ) + .unwrap() + } + + #[test] + fn test_lock_new_segments_covers_segments_created_while_locked() { + let g = small_pages_graph(); + g.add_edge(0, 1, 2, NO_PROPS, None).unwrap(); + + let mut wlg = g.write_lock().unwrap(); + let tg = wlg.graph; + let held_nodes = wlg.nodes.len(); + let held_edges = wlg.edges.len(); + + tg.storage().nodes().grow(held_nodes + 3); + tg.storage().edges().grow(held_edges + 2); + assert_eq!(wlg.nodes.len(), held_nodes); + assert_eq!(wlg.edges.len(), held_edges); + + wlg.lock_new_segments(); + + assert_eq!(wlg.nodes.len(), tg.storage().nodes().num_segments()); + assert_eq!(wlg.edges.len(), tg.storage().edges().num_segments()); + assert_eq!(wlg.nodes.len(), held_nodes + 3); + assert_eq!(wlg.edges.len(), held_edges + 2); + for segment_id in 0..wlg.nodes.len() { + assert_eq!( + wlg.nodes.get_mut(segment_id).unwrap().segment_id(), + segment_id + ); + } + for segment_id in 0..wlg.edges.len() { + assert_eq!(wlg.edges.get_mut(segment_id).unwrap().page_id(), segment_id); + } + + // nothing new: no-op + wlg.lock_new_segments(); + assert_eq!(wlg.nodes.len(), held_nodes + 3); + assert_eq!(wlg.edges.len(), held_edges + 2); + } + + #[test] + fn test_load_edges_creating_many_segments_in_one_chunk() { + // every edge is new and the pages are tiny, so the node pass of a single + // chunk has to create edge segments that the edge pass then writes to + let edges = (0..200u64) + .map(|i| (i, i + 1000, i as i64, i.to_string(), i as i64)) + .collect::>(); + let g = small_pages_graph(); + load_edges_from_df( + build_df(edges.len(), &edges), + ColumnNames::new("time", None, "src", "dst", None), + true, + &["str_prop", "int_prop"], + &[], + None, + None, + &g, + false, + ) + .unwrap(); + + let g2 = Graph::new(); + for (src, dst, time, str_prop, int_prop) in &edges { + g2.add_edge( + *time, + *src, + *dst, + [ + ("str_prop", str_prop.clone().into_prop()), + ("int_prop", int_prop.into_prop()), + ], + None, + ) + .unwrap(); + } + assert_eq!(g.count_edges(), edges.len()); + assert_graph_equal(&g, &g2); + } } mod parquet_tests { diff --git a/raphtory/src/arrow_loader/df_loaders/edge_props.rs b/raphtory/src/arrow_loader/df_loaders/edge_props.rs index fc7a817d82..b27c12019a 100644 --- a/raphtory/src/arrow_loader/df_loaders/edge_props.rs +++ b/raphtory/src/arrow_loader/df_loaders/edge_props.rs @@ -144,9 +144,7 @@ pub fn load_edges_metadata_from_df(()) })?; - drop(write_locked_graph); - - let mut write_locked_graph = graph.write_lock().map_err(into_graph_err)?; + write_locked_graph.lock_new_segments(); write_locked_graph.edges.par_iter_mut().for_each(|shard| { let zip = izip!( diff --git a/raphtory/src/arrow_loader/df_loaders/edges.rs b/raphtory/src/arrow_loader/df_loaders/edges.rs index 406220d343..b1af555653 100644 --- a/raphtory/src/arrow_loader/df_loaders/edges.rs +++ b/raphtory/src/arrow_loader/df_loaders/edges.rs @@ -428,8 +428,8 @@ pub fn load_edges_from_df