Hold one WriteLockedGraph for the whole bulk load (pometry-storage#424)

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

@@ -11,7 +11,7 @@ use crate::{
     segments::node::segment::MemNodeSegment,
 };
 use parking_lot::RwLockWriteGuard;
-use raphtory_api::core::entities::LayerId;
+use raphtory_api::core::entities::{LayerId, properties::prop::Prop};
 use raphtory_core::entities::VID;
 use rayon::prelude::*;
 use std::ops::DerefMut;
@@ -84,6 +84,12 @@ impl<'a, NS: NodeSegmentOps> LockedNodePage<'a, NS> {
         self.lock.get_or_create_layer(layer_id);
         self.layer_counter.get(layer_id);
     }
+
+    pub fn get_metadata(&self, pos: LocalPOS, layer_id: LayerId, prop_id: usize) -> Option<Prop> {
+        self.lock
+            .get_metadata(pos, layer_id, prop_id)
+            .or_else(|| self.page.get_metadata_immut(pos, layer_id, prop_id))
+    }
 }
 
 pub struct WriteLockedNodePages<'a, NS> {
@@ -110,6 +116,18 @@ impl<'a, EXT: PersistenceStrategy<NS = NS>, NS: NodeSegmentOps<Extension = EXT>>
         self.writers.len()
     }
 
+    /// Read node metadata through the held write locks. Use this instead of the graph's
+    /// read path while the pages are held, which would deadlock on the segment lock.
+    ///
+    /// Returns `None` if the node's segment is not held.
+    pub fn node_metadata(&self, vid: VID, layer_id: LayerId, prop_id: usize) -> Option<Prop> {
+        let max_page_len = self.writers.first()?.max_page_len;
+        let (segment_id, pos) = resolve_pos(vid, max_page_len);
+        self.writers
+            .get(segment_id)?
+            .get_metadata(pos, layer_id, prop_id)
+    }
+
     /// 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());

raphtory-tests/tests/io_tests/df_loaders.rs

@@ -987,6 +987,129 @@ mod io_tests {
         assert!(!g.has_node(9u64));
         assert_eq!(g.node(1u64).unwrap().metadata().get("str_prop"), None);
     }
+    /// Runs `f` on another thread and fails the test if it does not finish in time,
+    /// so a deadlocked load fails instead of hanging the test run.
+    fn finishes_within<T: Send + 'static>(f: impl FnOnce() -> T + Send + 'static) -> T {
+        let (tx, rx) = std::sync::mpsc::channel();
+        std::thread::spawn(move || tx.send(f()).unwrap());
+        rx.recv_timeout(std::time::Duration::from_secs(60))
+            .expect("bulk load did not finish (deadlock?)")
+    }
+
+    fn load_typed_nodes(
+        g: &Graph,
+        chunk_size: usize,
+        nodes: &[(u64, i64, &str)],
+    ) -> Result<(), GraphError> {
+        load_nodes_from_df(
+            build_nodes_df(chunk_size, nodes),
+            "time",
+            None,
+            "id",
+            &[],
+            &[],
+            None,
+            None,
+            Some("node_type"),
+            g,
+            true,
+            None,
+            None,
+            None,
+        )
+    }
+
+    fn sorted_node_types(g: &Graph) -> Vec<(GID, Option<ArcStr>)> {
+        let mut result = g
+            .nodes()
+            .into_iter()
+            .map(|node| (node.id(), node.node_type()))
+            .collect::<Vec<_>>();
+        result.sort();
+        result
+    }
+
+    #[test]
+    fn test_reload_existing_typed_nodes() {
+        let nodes: Vec<(u64, i64, &str)> = (0..20u64)
+            .map(|i| (i, i as i64, if i % 2 == 0 { "even" } else { "odd" }))
+            .collect();
+        let g = small_pages_graph();
+        let g_load = g.clone();
+        let reloaded = nodes
+            .iter()
+            .map(|(id, t, ty)| (*id, t + 100, *ty))
+            .collect::<Vec<_>>();
+        finishes_within(move || {
+            load_typed_nodes(&g_load, 7, &nodes).unwrap();
+            // every gid already exists and already has the same type
+            load_typed_nodes(&g_load, 7, &reloaded).unwrap();
+        });
+
+        let expected = (0..20u64)
+            .map(|i| {
+                (
+                    GID::U64(i),
+                    Some(ArcStr::from(if i % 2 == 0 { "even" } else { "odd" })),
+                )
+            })
+            .collect::<Vec<_>>();
+        assert_eq!(sorted_node_types(&g), expected);
+        for node in g.nodes() {
+            assert_eq!(node.history().len(), 2);
+        }
+    }
+
+    #[test]
+    fn test_load_typed_nodes_after_edges() {
+        let edges = (0..20u64)
+            .map(|i| (i, (i + 1) % 20, i as i64, i.to_string(), i as i64))
+            .collect::<Vec<_>>();
+        let nodes: Vec<(u64, i64, &str)> = (0..20u64).map(|i| (i, 50, "typed")).collect();
+        let g = small_pages_graph();
+        let g_load = g.clone();
+        finishes_within(move || {
+            load_edges_from_df(
+                build_df(7, &edges),
+                ColumnNames::new("time", None, "src", "dst", None),
+                true,
+                &[],
+                &[],
+                None,
+                None,
+                &g_load,
+                false,
+            )
+            .unwrap();
+            // every gid already exists without a stored type
+            load_typed_nodes(&g_load, 7, &nodes).unwrap();
+        });
+
+        let expected = (0..20u64)
+            .map(|i| (GID::U64(i), Some(ArcStr::from("typed"))))
+            .collect::<Vec<_>>();
+        assert_eq!(sorted_node_types(&g), expected);
+        assert_eq!(g.count_edges(), 20);
+    }
+
+    #[test]
+    fn test_reload_existing_node_with_conflicting_type_fails() {
+        let g = small_pages_graph();
+        let g_load = g.clone();
+        let result = finishes_within(move || {
+            load_typed_nodes(&g_load, 7, &[(1, 1, "a"), (2, 2, "b")]).unwrap();
+            load_typed_nodes(&g_load, 7, &[(1, 3, "b")])
+        });
+        assert!(result.is_err());
+        assert_eq!(
+            sorted_node_types(&g),
+            vec![
+                (GID::U64(1), Some(ArcStr::from("a"))),
+                (GID::U64(2), Some(ArcStr::from("b"))),
+            ]
+        );
+    }
+
     fn small_pages_graph() -> Graph {
         Graph::new_with_config(
             Args::default()

raphtory/src/arrow_loader/df_loaders/edge_props.rs

@@ -88,6 +88,9 @@ pub fn load_edges_metadata_from_df<G: StaticGraphViewOps + PropertyAdditionOps +
     let mut dst_col_resolved: Vec<VID> = vec![];
     let mut eid_col_resolved: Vec<EID> = vec![];
 
+    // held for the whole load
+    let mut write_locked_graph = graph.write_lock().map_err(into_graph_err)?;
+
     for chunk in df_view.chunks {
         let df = chunk?;
         let metadata_cols =
@@ -129,8 +132,6 @@ pub fn load_edges_metadata_from_df<G: StaticGraphViewOps + PropertyAdditionOps +
             &dst_col,
         )?;
 
-        let mut write_locked_graph = graph.write_lock().map_err(into_graph_err)?;
-
         eid_col_resolved.resize_with(df.len(), Default::default);
         let eid_col_shared = atomic_usize_from_mut_slice(cast_slice_mut(&mut eid_col_resolved));
 

raphtory/src/arrow_loader/df_loaders/edges.rs

@@ -224,6 +224,9 @@ pub fn load_edges_from_df<G: StaticGraphViewOps + PropertyAdditionOps + Addition
     let mut layer_eids_exist: Vec<AtomicBool> = vec![]; // exists or needs to be created
     let mut node_resolve_cache: Option<NodeResolveCache<VID>> = None;
 
+    // held for the whole load
+    let mut write_locked_graph = graph.write_lock().map_err(into_graph_err)?;
+
     for chunk in df_view.chunks {
         let df = chunk?;
         let prop_cols =
@@ -288,7 +291,8 @@ pub fn load_edges_from_df<G: StaticGraphViewOps + PropertyAdditionOps + Addition
         let secondary_index_col =
             extract_secondary_index_col::<G>(secondary_index_index, &session, &df)?;
 
-        let mut write_locked_graph = graph.write_lock().map_err(into_graph_err)?;
+        // node resolution may have created new node segments
+        write_locked_graph.lock_new_segments();
         let max_node_segment_len = write_locked_graph
             .graph()
             .storage()

raphtory/src/arrow_loader/df_loaders/nodes.rs

@@ -31,9 +31,9 @@ use std::{
 };
 use storage::{
     api::{node_type_index::NodeTypeIndexOps, nodes::NodeSegmentOps},
-    pages::locked::nodes::LockedNodePage,
+    pages::locked::nodes::{LockedNodePage, WriteLockedNodePages},
     segments::node_type_index::index::MemNodeTypeIndex,
-    Extension,
+    Extension, NS,
 };
 
 #[cfg(feature = "progress")]
@@ -179,6 +179,15 @@ pub fn load_nodes_from_df<
         let mut node_col_resolved = vec![];
         let mut node_resolve_cache: Option<NodeResolveCache<(VID, usize)>> = None;
 
+        // held for the whole load
+        let mut write_locked_graph = graph.write_lock().map_err(into_graph_err)?;
+        let node_stats = write_locked_graph.node_stats().clone();
+        let max_node_segment_len = write_locked_graph
+            .graph()
+            .storage()
+            .nodes()
+            .max_segment_len();
+
         for chunk in df_view.chunks {
             let df = chunk?;
 
@@ -233,6 +242,7 @@ pub fn load_nodes_from_df<
 
             let (src_vids, gid_str_cache) = get_or_resolve_node_vids::<G>(
                 graph,
+                &write_locked_graph.nodes,
                 node_id_index,
                 &mut node_resolve_cache,
                 &mut node_col_resolved,
@@ -248,18 +258,12 @@ pub fn load_nodes_from_df<
                 index.notify_write();
             }
 
-            let mut write_locked_graph = graph.write_lock().map_err(into_graph_err)?;
-            let node_stats = write_locked_graph.node_stats().clone();
-
-            let max_node_segment_len = write_locked_graph
-                .graph()
-                .storage()
-                .nodes()
-                .max_segment_len() as usize;
+            // node resolution may have created new node segments
+            write_locked_graph.lock_new_segments();
 
             let rows_by_segment = group_rows_by_vid_segment(
                 src_vids,
-                max_node_segment_len as u32,
+                max_node_segment_len,
                 write_locked_graph.nodes.len(),
             );
 
@@ -387,6 +391,9 @@ pub fn load_node_props_from_df<
     let mut node_col_resolved = vec![];
     let mut node_type_resolved = vec![];
 
+    // held for the whole load
+    let mut write_locked_graph = graph.write_lock().map_err(into_graph_err)?;
+
     for chunk in df_view.chunks {
         let df = chunk?;
         if df.is_empty() {
@@ -432,7 +439,6 @@ pub fn load_node_props_from_df<
             .max()
             .map(VID)
             .unwrap_or(VID(0));
-        let mut write_locked_graph = graph.write_lock().map_err(into_graph_err)?;
         write_locked_graph.resize_segments_to_vid(max_vid);
 
         write_locked_graph
@@ -491,6 +497,7 @@ fn get_or_resolve_node_vids<
     G: StaticGraphViewOps + PropertyAdditionOps + AdditionOps,
 >(
     graph: &G,
+    locked_nodes: &WriteLockedNodePages<'_, NS<Extension>>,
     src_index: usize,
     node_resolve_cache: &mut Option<NodeResolveCache<(VID, usize)>>,
     src_col_resolved: &'a mut Vec<VID>,
@@ -505,6 +512,7 @@ fn get_or_resolve_node_vids<
 
         resolve_node_vids_and_types_with_cache(
             graph,
+            locked_nodes,
             cache,
             src_col_resolved,
             src_col,
@@ -536,6 +544,7 @@ fn resolve_node_vids_and_types_with_cache<
     G: StaticGraphViewOps + PropertyAdditionOps + AdditionOps,
 >(
     graph: &G,
+    locked_nodes: &WriteLockedNodePages<'_, NS<Extension>>,
     cache: &mut NodeResolveCache<(VID, usize)>,
     node_col_resolved: &'a mut Vec<VID>,
     node_col: &'a NodeCol,
@@ -566,20 +575,22 @@ fn resolve_node_vids_and_types_with_cache<
                             should_store = true;
                             Ok((vid, node_type))
                         }
-                        MaybeNew::Existing(vid) => match existing_stored_node_type(graph, vid) {
-                            None => {
-                                should_store = true;
-                                Ok((vid, node_type))
+                        MaybeNew::Existing(vid) => {
+                            match existing_stored_node_type(locked_nodes, vid) {
+                                None => {
+                                    should_store = true;
+                                    Ok((vid, node_type))
+                                }
+                                Some(existing) if existing == node_type => Ok((vid, existing)),
+                                Some(existing) => Err(GraphError::LoadError {
+                                    source: LoadError::ConflictingNodeType {
+                                        gid: gid.into(),
+                                        existing: node_type_name(graph, existing),
+                                        new: node_type_name(graph, node_type),
+                                    },
+                                }),
                             }
-                            Some(existing) if existing == node_type => Ok((vid, existing)),
-                            Some(existing) => Err(GraphError::LoadError {
-                                source: LoadError::ConflictingNodeType {
-                                    gid: gid.into(),
-                                    existing: node_type_name(graph, existing),
-                                    new: node_type_name(graph, node_type),
-                                },
-                            }),
-                        },
+                        }
                     }
                 })?;
 
@@ -610,12 +621,14 @@ fn resolve_node_vids_and_types_with_cache<
     Ok((node_col_resolved.as_slice(), new_gids_by_shard.concat()))
 }
 
-fn existing_stored_node_type<G: StaticGraphViewOps + PropertyAdditionOps + AdditionOps>(
-    graph: &G,
+/// Reads through the held write locks: the graph's read path would deadlock on them.
+/// An existing node always lives in a segment that is already held.
+fn existing_stored_node_type(
+    locked_nodes: &WriteLockedNodePages<'_, NS<Extension>>,
     vid: VID,
 ) -> Option<usize> {
-    graph
-        .node_metadata(vid, NODE_TYPE_PROP_ID)
+    locked_nodes
+        .node_metadata(vid, STATIC_GRAPH_LAYER_ID, NODE_TYPE_PROP_ID)
         .and_then(|prop| prop.into_u64())
         .map(|id| id as usize)
 }