Bulk loads exclude regular updates (pometry-storage#425)

db4-graph/Cargo.toml

@@ -15,4 +15,5 @@ edition.workspace = true
 storage.workspace = true
 raphtory-api.workspace = true
 raphtory-core.workspace = true
+parking_lot.workspace = true
 rayon.workspace = true

db4-graph/src/lib.rs

@@ -1,3 +1,4 @@
+use parking_lot::RwLockWriteGuard;
 use raphtory_api::core::{
     entities::{
         self,
@@ -434,6 +435,8 @@ where
     pub edges: WriteLockedEdgePages<'a, ES<EXT>>,
     pub graph_props: WriteLockedGraphPropPages<'a, GS<EXT>>,
     pub graph: &'a TemporalGraph<EXT>,
+    // Declared last so it is released after the segment locks.
+    _bulk_load: RwLockWriteGuard<'a, ()>,
 }
 
 impl<'a, EXT> WriteLockedGraph<'a, EXT>
@@ -443,12 +446,15 @@ where
     ES<EXT>: EdgeSegmentOps<Extension = EXT>,
     GS<EXT>: GraphPropSegmentOps<Extension = EXT>,
 {
+    /// Takes the bulk load write guard first, waiting for regular updates in flight to finish.
     pub fn new(graph: &'a TemporalGraph<EXT>) -> Self {
+        let bulk_load = graph.storage.bulk_load_write();
         WriteLockedGraph {
             nodes: graph.storage.nodes().write_locked(),
             edges: graph.storage.edges().write_locked(),
             graph_props: graph.storage.graph_props().write_locked(),
             graph,
+            _bulk_load: bulk_load,
         }
     }
 

db4-storage/src/pages/mod.rs

@@ -18,7 +18,7 @@ use edge_store::EdgeStorageInner;
 use graph_prop_store::GraphPropStorageInner;
 use node_page::writer::NodeWriter;
 use node_store::NodeStorageInner;
-use parking_lot::RwLockWriteGuard;
+use parking_lot::{RwLock, RwLockReadGuard, RwLockWriteGuard};
 use raphtory_api::core::{
     entities::properties::meta::Meta, storage::graph_folder::InnerGraphFolder,
 };
@@ -59,6 +59,8 @@ pub struct GraphStore<
     graph_dir: Option<PathBuf>,
     event_id: AtomicUsize,
     ext: EXT,
+    /// Write-held for the whole of a bulk load; regular updates and vacuum hold a read guard.
+    bulk_loader_lock: RwLock<()>,
 }
 
 impl<
@@ -164,6 +166,7 @@ impl<
             event_id: AtomicUsize::new(0),
             graph_dir: graph_dir.map(|p| p.to_path_buf()),
             ext,
+            bulk_loader_lock: RwLock::new(()),
         }
     }
 
@@ -202,6 +205,7 @@ impl<
             event_id: AtomicUsize::new(t_len),
             graph_dir: Some(graph_dir.as_ref().to_path_buf()),
             ext,
+            bulk_loader_lock: RwLock::new(()),
         })
     }
 
@@ -301,7 +305,19 @@ impl<
         self.edges().get_free_writer()
     }
 
+    /// Read guard held by regular updates and vacuum, excluding them from bulk loads.
+    /// Recursive, so a regular update nested in another does not wait behind a queued bulk load.
+    pub fn bulk_load_read(&self) -> RwLockReadGuard<'_, ()> {
+        self.bulk_loader_lock.read_recursive()
+    }
+
+    /// Write guard held for the whole of a bulk load or WAL replay.
+    pub fn bulk_load_write(&self) -> RwLockWriteGuard<'_, ()> {
+        self.bulk_loader_lock.write()
+    }
+
     pub fn vacuum(self: &Arc<Self>) -> Result<(), StorageError> {
+        let _bulk_load = self.bulk_load_read();
         let mut locked_nodes = self.nodes.write_locked();
         let mut locked_edges = self.edges.write_locked();
 

db4-storage/src/transaction/mod.rs

@@ -1,4 +1,8 @@
-use std::sync::atomic::{self, AtomicU64};
+use parking_lot::RwLockReadGuard;
+use std::{
+    ops::Deref,
+    sync::atomic::{self, AtomicU64},
+};
 
 use crate::wal::TransactionID;
 
@@ -38,3 +42,27 @@ impl Default for TransactionManager {
         Self::new()
     }
 }
+
+/// The transaction manager together with a bulk load read guard, so a regular update
+/// cannot run while a bulk load holds the graph. Keep it until the update ends.
+pub struct TransactionGuard<'a> {
+    manager: &'a TransactionManager,
+    _bulk_load: RwLockReadGuard<'a, ()>,
+}
+
+impl<'a> TransactionGuard<'a> {
+    pub fn new(manager: &'a TransactionManager, bulk_load: RwLockReadGuard<'a, ()>) -> Self {
+        Self {
+            manager,
+            _bulk_load: bulk_load,
+        }
+    }
+}
+
+impl Deref for TransactionGuard<'_> {
+    type Target = TransactionManager;
+
+    fn deref(&self) -> &Self::Target {
+        self.manager
+    }
+}

raphtory-storage/src/durability_ops.rs

@@ -1,17 +1,19 @@
 use crate::{graph::graph::GraphStorage, mutation::MutationError};
-use storage::{transaction::TransactionManager, ControlFile, Wal};
+use storage::{transaction::TransactionGuard, ControlFile, Wal};
 
 /// Accessor methods for supporting durability.
 pub trait DurabilityOps {
-    fn transaction_manager(&self) -> Result<&TransactionManager, MutationError>;
+    /// Start a regular update. The guard excludes bulk loads until it is dropped, so keep it
+    /// for the whole update. Never call this from inside a bulk load: it would deadlock.
+    fn transaction_manager(&self) -> Result<TransactionGuard<'_>, MutationError>;
 
     fn wal(&self) -> Result<&Wal, MutationError>;
 
     fn control_file(&self) -> Result<&ControlFile, MutationError>;
 }
 
 impl DurabilityOps for GraphStorage {
-    fn transaction_manager(&self) -> Result<&TransactionManager, MutationError> {
+    fn transaction_manager(&self) -> Result<TransactionGuard<'_>, MutationError> {
         self.mutable()?.transaction_manager()
     }
 

raphtory-storage/src/mutation/addition_ops_ext.rs

@@ -35,7 +35,7 @@ use storage::{
     persist::{config::ConfigOps, strategy::PersistenceStrategy},
     properties::props_meta_writer::PropsMetaWriter,
     resolver::{GIDResolverOps, Initialiser, MaybeInit},
-    transaction::TransactionManager,
+    transaction::TransactionGuard,
     wal::LSN,
     ControlFile, Extension, LocalPOS, Wal, ES, GS, NS,
 };
@@ -743,8 +743,11 @@ impl InternalAdditionOps for TemporalGraph {
 }
 
 impl DurabilityOps for TemporalGraph {
-    fn transaction_manager(&self) -> Result<&TransactionManager, MutationError> {
-        Ok(&self.transaction_manager)
+    fn transaction_manager(&self) -> Result<TransactionGuard<'_>, MutationError> {
+        Ok(TransactionGuard::new(
+            &self.transaction_manager,
+            self.storage().bulk_load_read(),
+        ))
     }
 
     fn wal(&self) -> Result<&Wal, MutationError> {

raphtory-tests/tests/io_tests/df_loaders.rs

@@ -1194,6 +1194,123 @@ mod io_tests {
         assert_eq!(g.count_edges(), edges.len());
         assert_graph_equal(&g, &g2);
     }
+
+    /// Runs `f` on another thread and asserts it is still blocked after a short wait.
+    /// Returns a receiver that yields once `f` finishes.
+    fn assert_blocked(f: impl FnOnce() + Send + 'static) -> std::sync::mpsc::Receiver<()> {
+        let (tx, rx) = std::sync::mpsc::channel();
+        std::thread::spawn(move || {
+            f();
+            tx.send(()).unwrap();
+        });
+        assert!(
+            rx.recv_timeout(std::time::Duration::from_millis(200))
+                .is_err(),
+            "ran while the graph was write locked for a bulk load"
+        );
+        rx
+    }
+
+    #[test]
+    fn test_write_locked_graph_excludes_regular_updates_and_vacuum() {
+        let g = small_pages_graph();
+        g.add_edge(0, 1, 2, NO_PROPS, None).unwrap();
+
+        let wlg = g.write_lock().unwrap();
+
+        // gids far beyond the held segments, so only the bulk load lock can block these
+        let g_update = g.clone();
+        let update_done = assert_blocked(move || {
+            g_update.add_edge(1, 100, 101, NO_PROPS, None).unwrap();
+        });
+        let g_node = g.clone();
+        let node_done = assert_blocked(move || {
+            g_node.add_node(1, 200, NO_PROPS, None, None).unwrap();
+        });
+        let g_vacuum = g.clone();
+        let vacuum_done = assert_blocked(move || {
+            g_vacuum.vacuum().unwrap();
+        });
+
+        drop(wlg);
+        let timeout = std::time::Duration::from_secs(60);
+        update_done.recv_timeout(timeout).unwrap();
+        node_done.recv_timeout(timeout).unwrap();
+        vacuum_done.recv_timeout(timeout).unwrap();
+        assert!(g.has_edge(100, 101));
+        assert!(g.has_node(200));
+    }
+
+    #[test]
+    fn test_regular_updates_during_edge_bulk_load() {
+        const BULK_EDGES: u64 = 2000;
+        const THREADS: u64 = 4;
+        const UPDATES: u64 = 200;
+
+        let g = small_pages_graph();
+        let edges = (0..BULK_EDGES)
+            .map(|i| (i, i + 1, i as i64, i.to_string(), i as i64))
+            .collect::<Vec<_>>();
+
+        let g_load = g.clone();
+        let g_updates = g.clone();
+        finishes_within(move || {
+            std::thread::scope(|s| {
+                s.spawn(|| {
+                    load_edges_from_df(
+                        build_df(50, &edges),
+                        ColumnNames::new("time", None, "src", "dst", None),
+                        true,
+                        &["str_prop", "int_prop"],
+                        &[],
+                        None,
+                        None,
+                        &g_load,
+                        false,
+                    )
+                    .unwrap();
+                });
+                for thread in 0..THREADS {
+                    let g = &g_updates;
+                    s.spawn(move || {
+                        for j in 0..UPDATES {
+                            let t = (thread * UPDATES + j) as i64;
+                            // new nodes and a new edge, disjoint from the bulk load
+                            let base = 1_000_000 + 2 * (thread * UPDATES + j);
+                            g.add_edge(t, base, base + 1, NO_PROPS, None).unwrap();
+                            g.add_node(t, base, NO_PROPS, Some("regular"), None)
+                                .unwrap();
+                            g.node(base + 1).unwrap().set_node_type("regular").unwrap();
+                            // a new edge between nodes the bulk load also writes
+                            let k = thread * UPDATES + j;
+                            g.add_edge(t, k, k + 3, NO_PROPS, None).unwrap();
+                        }
+                    });
+                }
+            });
+        });
+
+        let regular_nodes = 2 * THREADS * UPDATES;
+        let regular_edges = 2 * THREADS * UPDATES;
+        assert_eq!(g.count_nodes() as u64, BULK_EDGES + 1 + regular_nodes);
+        assert_eq!(g.count_edges() as u64, BULK_EDGES + regular_edges);
+        for i in 0..BULK_EDGES {
+            let edge = g.edge(i, i + 1).unwrap();
+            assert_eq!(edge.properties().get("int_prop"), Some(Prop::I64(i as i64)));
+        }
+        for k in 0..THREADS * UPDATES {
+            assert!(g.has_edge(k, k + 3));
+            let base = 1_000_000 + 2 * k;
+            assert_eq!(
+                g.node(base).unwrap().node_type(),
+                Some(ArcStr::from("regular"))
+            );
+            assert_eq!(
+                g.node(base + 1).unwrap().node_type(),
+                Some(ArcStr::from("regular"))
+            );
+        }
+    }
 }
 
 mod parquet_tests {

raphtory/src/arrow_loader/df_loaders/mod.rs

@@ -17,7 +17,10 @@ use raphtory_api::core::{
     storage::{dict_mapper::MaybeNew, timeindex::EventTime, FxDashMap},
 };
 use raphtory_core::entities::{GidRef, VID};
-use raphtory_storage::mutation::addition_ops::{InternalAdditionOps, SessionAdditionOps};
+use raphtory_storage::{
+    durability_ops::DurabilityOps,
+    mutation::addition_ops::{InternalAdditionOps, SessionAdditionOps},
+};
 use rayon::prelude::*;
 use std::{
     collections::HashMap,
@@ -179,6 +182,8 @@ pub(crate) fn load_graph_props_from_df<
 
     #[cfg(feature = "progress")]
     let mut pb = build_progress_bar("Loading graph properties".to_string(), df_view.num_rows)?;
+    // Graph properties are written as regular updates, not bulk loaded.
+    let _transaction_manager = graph.core_graph().transaction_manager()?;
     let session = graph.write_session().map_err(into_graph_err)?;
 
     for chunk in df_view.chunks {

raphtory/src/db/api/mutation/import_ops.rs

@@ -20,8 +20,11 @@ use raphtory_api::core::{
     entities::GID,
     storage::{arc_str::OptionAsStr, timeindex::AsTime},
 };
-use raphtory_storage::mutation::{
-    addition_ops::InternalAdditionOps, property_addition_ops::InternalPropertyAdditionOps,
+use raphtory_storage::{
+    durability_ops::DurabilityOps,
+    mutation::{
+        addition_ops::InternalAdditionOps, property_addition_ops::InternalPropertyAdditionOps,
+    },
 };
 use std::{borrow::Borrow, fmt::Debug};
 
@@ -285,6 +288,8 @@ fn import_node_internal<
     id: V,
     merge: bool,
 ) -> Result<NodeView<'static, G>, GraphError> {
+    // Writes through the internal addition ops, so take the regular update guard here.
+    let _transaction_manager = graph.core_graph().transaction_manager()?;
     let id = id.as_node_ref();
     let gid_ref = id.as_gid_ref();
     graph.validate_gids(gid_ref).map_err(into_graph_err)?;