From 048c33454a7e91709362b55f286e47b1a01d6459 Mon Sep 17 00:00:00 2001 From: Fabian Murariu Date: Thu, 15 Oct 2026 15:17:14 +0100 Subject: [PATCH] Hold one WriteLockedGraph for the whole bulk load Each bulk loader (nodes, node properties, edges/deletions, edge metadata) now takes its WriteLockedGraph once before the chunk loop and picks up segments created by node/edge resolution with lock_new_segments, instead of re-taking every segment lock per chunk. The existing-node type check during node resolution now reads through the held node pages (WriteLockedNodePages::node_metadata): going through the graph's read path deadlocked against the held write lock. Part of Pometry/pometry-storage#423. Fixes Pometry/pometry-storage#424 Co-Authored-By: Claude Opus 5.5 (1M context) --- db4-storage/src/pages/locked/nodes.rs | 20 ++- raphtory-tests/tests/io_tests/df_loaders.rs | 123 ++++++++++++++++++ .../src/arrow_loader/df_loaders/edge_props.rs | 5 +- raphtory/src/arrow_loader/df_loaders/edges.rs | 6 +- raphtory/src/arrow_loader/df_loaders/nodes.rs | 71 +++++----- 5 files changed, 192 insertions(+), 33 deletions(-) diff --git a/db4-storage/src/pages/locked/nodes.rs b/db4-storage/src/pages/locked/nodes.rs index 27d8a93db0..bd88aabd05 100644 --- a/db4-storage/src/pages/locked/nodes.rs +++ b/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 { + 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: NodeSegmentOps> 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 { + 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()); diff --git a/raphtory-tests/tests/io_tests/df_loaders.rs b/raphtory-tests/tests/io_tests/df_loaders.rs index 7ad4947fca..8a88f61720 100644 --- a/raphtory-tests/tests/io_tests/df_loaders.rs +++ b/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(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)> { + let mut result = g + .nodes() + .into_iter() + .map(|node| (node.id(), node.node_type())) + .collect::>(); + 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::>(); + 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::>(); + 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::>(); + 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::>(); + 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() diff --git a/raphtory/src/arrow_loader/df_loaders/edge_props.rs b/raphtory/src/arrow_loader/df_loaders/edge_props.rs index b27c12019a..f733680d1d 100644 --- a/raphtory/src/arrow_loader/df_loaders/edge_props.rs +++ b/raphtory/src/arrow_loader/df_loaders/edge_props.rs @@ -88,6 +88,9 @@ pub fn load_edges_metadata_from_df = vec![]; let mut eid_col_resolved: Vec = 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 = vec![]; // exists or needs to be created let mut node_resolve_cache: Option> = 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(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() diff --git a/raphtory/src/arrow_loader/df_loaders/nodes.rs b/raphtory/src/arrow_loader/df_loaders/nodes.rs index a6017965df..e0f3dd4ba6 100644 --- a/raphtory/src/arrow_loader/df_loaders/nodes.rs +++ b/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> = 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::( 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>, src_index: usize, node_resolve_cache: &mut Option>, src_col_resolved: &'a mut Vec, @@ -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>, cache: &mut NodeResolveCache<(VID, usize)>, node_col_resolved: &'a mut Vec, 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( - 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>, vid: VID, ) -> Option { - 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) }