Add WriteLockedGraph::lock_new_segments (pometry-storage#427)

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<GraphStats> {

db4-storage/src/pages/edge_store.rs

@@ -468,6 +468,21 @@ impl<ES: EdgeSegmentOps<Extension = EXT>, EXT: PersistenceStrategy<ES = ES>>
         )
     }
 
+    /// 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<ES>> {
         let (segment_id, _) = resolve_pos(eid, self.max_page_len());

db4-storage/src/pages/locked/edges.rs

@@ -145,6 +145,12 @@ impl<'a, EXT: PersistenceStrategy<ES = ES>, ES: EdgeSegmentOps<Extension = EXT>>
         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()
     }

db4-storage/src/pages/locked/nodes.rs

@@ -110,6 +110,12 @@ impl<'a, EXT: PersistenceStrategy<NS = NS>, NS: NodeSegmentOps<Extension = EXT>>
         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)

db4-storage/src/pages/node_store.rs

@@ -327,6 +327,21 @@ impl<NS: NodeSegmentOps<Extension = EXT>, EXT: PersistenceStrategy<NS = NS>>
         )
     }
 
+    /// 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())

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::<Vec<_>>();
+        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 {

raphtory/src/arrow_loader/df_loaders/edge_props.rs

@@ -144,9 +144,7 @@ pub fn load_edges_metadata_from_df<G: StaticGraphViewOps + PropertyAdditionOps +
             Ok::<_, GraphError>(())
         })?;
 
-        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!(

raphtory/src/arrow_loader/df_loaders/edges.rs

@@ -428,8 +428,8 @@ pub fn load_edges_from_df<G: StaticGraphViewOps + PropertyAdditionOps + Addition
                 update_inbound_edges(shard, zip, delete);
             });
 
-        drop(write_locked_graph);
-        let mut write_locked_graph = graph.write_lock().map_err(into_graph_err)?;
+        // the node pass may have created new edge segments
+        write_locked_graph.lock_new_segments();
         let edge_rows_by_segment = group_rows_by_eid_segment(
             &eid_col_resolved,
             max_edge_page_len,