From 58b96fb601eda5f02caff0dc537c1ffa0f18bd15 Mon Sep 17 00:00:00 2001 From: David Hughes Date: Sun, 25 Oct 2026 09:54:33 -0700 Subject: [PATCH] fix(graph): prevent stale concurrent refresh publication Retry refresh when another replacement has been published since its starting snapshot. Add deterministic publication coverage and preserve failed-refresh and active-reader behavior. Resolves #12. Implemented with assistance from Codex; runtime and benchmark validation pending. --- crates/lance-graph-python/src/graph.rs | 123 ++++++++++++++++++++--- docs/api-reference.md | 3 +- docs/lance.md | 3 + python/python/tests/test_engine_lance.py | 23 ++++- 4 files changed, 134 insertions(+), 18 deletions(-) diff --git a/crates/lance-graph-python/src/graph.rs b/crates/lance-graph-python/src/graph.rs index 778d688..595b445 100644 --- a/crates/lance-graph-python/src/graph.rs +++ b/crates/lance-graph-python/src/graph.rs @@ -1921,6 +1921,16 @@ impl CypherEngine { .clone() } + /// Publish only if no other refresh replaced the snapshot used to build this state. + fn try_publish_refresh(&self, expected: &Arc, replacement: EngineState) -> bool { + let mut current = self.state.write().expect("engine state lock poisoned"); + if !Arc::ptr_eq(¤t, expected) { + return false; + } + *current = Arc::new(replacement); + true + } + fn parse( &self, query: &str, @@ -2085,25 +2095,31 @@ impl CypherEngine { /// The replacement is validated in full before it is swapped in; on error nothing /// changes and the previous snapshot stays usable. Queries already running finish /// on the snapshot they started with. PyArrow inputs are unchanged. + /// If another refresh publishes first, retry from its snapshot instead of replacing + /// it with a candidate built from an older state. fn refresh(&self, py: Python) -> PyResult<()> { - let current = self.snapshot(); - let mut sources = current.sources.clone(); - for source in sources.values_mut() { - if let EngineSource::Lance { - uri, - storage_options, - dataset, - } = source - { - let latest = open_lance_dataset(py, uri, None, storage_options)?; - if !same_lance_snapshot(&latest, dataset) { - *dataset = Arc::new(latest); + loop { + let current = self.snapshot(); + let mut sources = current.sources.clone(); + for source in sources.values_mut() { + if let EngineSource::Lance { + uri, + storage_options, + dataset, + } = source + { + let latest = open_lance_dataset(py, uri, None, storage_options)?; + if !same_lance_snapshot(&latest, dataset) { + *dataset = Arc::new(latest); + } } } + let state = + build_engine_state(py, &self.config, sources, self.use_csr, Some(¤t))?; + if self.try_publish_refresh(¤t, state) { + return Ok(()); + } } - let state = build_engine_state(py, &self.config, sources, self.use_csr, Some(¤t))?; - *self.state.write().expect("engine state lock poisoned") = Arc::new(state); - Ok(()) } /// Node id -> row index statistics per node label (experimental, ``use_csr=True``). @@ -2453,6 +2469,83 @@ pub fn register_graph_module(py: Python, parent_module: &Bound<'_, PyModule>) -> mod tests { use super::*; + #[test] + fn test_refresh_rejects_stale_snapshot() { + use arrow_array::Int64Array; + use std::sync::mpsc; + use std::time::Duration; + + fn state(id: i64) -> EngineState { + let batch = RecordBatch::try_from_iter([( + "id", + Arc::new(Int64Array::from(vec![id])) as arrow_array::ArrayRef, + )]) + .unwrap(); + let provider = + Arc::new(MemTable::try_new(batch.schema(), vec![vec![batch.clone()]]).unwrap()); + let context = SessionContext::new(); + context.register_table("person", provider.clone()).unwrap(); + EngineState { + sources: HashMap::from([("Person".to_string(), EngineSource::Arrow(batch))]), + catalog: Arc::new( + InMemoryCatalog::new() + .with_node_source("Person", Arc::new(DefaultTableSource::new(provider))), + ), + context, + csr: HashMap::new(), + node_index: HashMap::new(), + key_stats: HashMap::new(), + } + } + + let engine = CypherEngine { + config: RustGraphConfig::builder() + .node_table("Person", "id") + .build() + .unwrap(), + state: RwLock::new(Arc::new(state(1))), + use_csr: false, + }; + let reader_snapshot = engine.snapshot(); + let (started_tx, started_rx) = mpsc::channel(); + let (published_tx, published_rx) = mpsc::channel(); + std::thread::scope(|threads| { + let engine = &engine; + let older = threads.spawn(move || { + let starting_snapshot = engine.snapshot(); + started_tx.send(()).unwrap(); + published_rx.recv_timeout(Duration::from_secs(5)).unwrap(); + engine.try_publish_refresh(&starting_snapshot, state(2)) + }); + started_rx.recv_timeout(Duration::from_secs(5)).unwrap(); + assert!(engine.try_publish_refresh(&engine.snapshot(), state(3))); + published_tx.send(()).unwrap(); + assert!(!older.join().unwrap()); + }); + + for (snapshot, expected) in [(reader_snapshot, 1), (engine.snapshot(), 3)] { + let query = RustCypherQuery::new("MATCH (p:Person) RETURN p.id") + .unwrap() + .with_config(engine.config.clone()); + let result = RT + .runtime + .block_on(query.execute_with_catalog_and_context( + snapshot.catalog.clone(), + snapshot.context.clone(), + )) + .unwrap(); + assert_eq!( + result + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .value(0), + expected + ); + } + } + #[test] fn test_split_vector_column_with_alias() { let (alias, column) = split_vector_column("d.embedding"); diff --git a/docs/api-reference.md b/docs/api-reference.md index 81ffd03..2a4292b 100644 --- a/docs/api-reference.md +++ b/docs/api-reference.md @@ -112,7 +112,8 @@ be opened. `#!python refresh() -> None` : Moves every Lance dataset to its latest version and updates the CSR, node indexes and statistics. - All-or-nothing; running queries keep their snapshot. See [Lance datasets](lance.md#refresh). + All-or-nothing; running queries keep their snapshot. If another refresh publishes first, retries from + that snapshot so an older candidate cannot replace it. See [Lance datasets](lance.md#refresh). `#!python dataset_versions() -> dict[str, int]` : The pinned version of each Lance dataset. PyArrow inputs are not listed. diff --git a/docs/lance.md b/docs/lance.md index d18a1cd..23037af 100644 --- a/docs/lance.md +++ b/docs/lance.md @@ -71,6 +71,9 @@ What `refresh()` guarantees: - **Schema changes are picked up.** New columns are queryable after the refresh. - **All or nothing.** The replacement is validated in full before it's swapped in. On error, nothing changes and the previous snapshot keeps working. +- **Concurrent refreshes cannot overwrite a newer published snapshot with an older candidate.** If another + refresh publishes first, the overlapping refresh retries from that snapshot and reopens the latest datasets. + The retry can add work when refreshes overlap; it does not block queries while the replacement is built. - **Running queries are unaffected.** Queries and `execute_batches()` readers that started before the refresh finish on their original snapshot. - **The CSR updates incrementally**, as described in [Keeping the CSR current](csr.md#keeping-the-csr-current). diff --git a/python/python/tests/test_engine_lance.py b/python/python/tests/test_engine_lance.py index 1b815d3..0c81356 100644 --- a/python/python/tests/test_engine_lance.py +++ b/python/python/tests/test_engine_lance.py @@ -153,9 +153,10 @@ def test_refresh_sees_table_replaced_at_same_location(graph): assert count_people(engine) == 3 -def test_failed_refresh_keeps_previous_snapshot(graph): +@pytest.mark.parametrize("use_csr", [False, True]) +def test_failed_refresh_keeps_previous_snapshot(graph, use_csr): path, person, knows = graph - engine = CypherEngine(CONFIG, {"Person": person, "KNOWS": knows}) + engine = CypherEngine(CONFIG, {"Person": person, "KNOWS": knows}, use_csr=use_csr) # Dropping the configured id column must not be accepted. lance.write_dataset( @@ -255,6 +256,24 @@ def test_execute_batches_partial_consumption(graph): assert count_people(engine) == N +@pytest.mark.parametrize("use_csr", [False, True]) +def test_execute_batches_keeps_snapshot_across_refresh(graph, use_csr): + path, person, knows = graph + engine = CypherEngine(CONFIG, {"Person": person, "KNOWS": knows}, use_csr=use_csr) + reader = engine.execute_batches("MATCH (p:Person) RETURN p.id") + first = reader.read_next_batch() + + lance.write_dataset(people(N, N + 5), path / "person.lance", mode="append") + engine.refresh() + + remaining = reader.read_all() + assert first.num_rows + remaining.num_rows == N + assert pc.sum(first.column(0)).as_py() + ( + pc.sum(remaining.column(0)).as_py() or 0 + ) == sum(range(N)) + assert count_people(engine) == N + 5 + + def test_execute_batches_reports_planning_errors(graph): _, person, knows = graph engine = CypherEngine(CONFIG, {"Person": person, "KNOWS": knows})