diff --git a/crates/lance-context-core/src/generic_store.rs b/crates/lance-context-core/src/generic_store.rs index baaf1465..15896be2 100644 --- a/crates/lance-context-core/src/generic_store.rs +++ b/crates/lance-context-core/src/generic_store.rs @@ -38,6 +38,16 @@ use lance_context_api::schema_spec::{SchemaSpec, ID_COLUMN}; /// reopened without the caller re-declaring its columns. const SCHEMA_SPEC_KEY: &str = "lance-context:schema_spec"; +/// Schema-metadata key holding the store's seal-on-add mode. +/// +/// Persisted for the same reason as the schema: it is a property *of the +/// store*, chosen at creation, not of whoever happens to open it. Without this +/// a store created with `seal_on_add: true` silently loses read-your-write the +/// first time it is reopened (after an LRU eviction, or a process restart) — +/// no error, no data loss, just "sometimes I can't read what I just wrote", +/// which is close to undiagnosable in production. +const SEAL_ON_ADD_KEY: &str = "lance-context:seal_on_add"; + /// Configuration for opening a [`GenericStore`]. #[derive(Debug, Clone, Default)] pub struct GenericStoreOptions { @@ -71,6 +81,9 @@ pub struct GenericStore { /// not against `spec.to_arrow()`, so a dataset whose physical field order /// differs still writes correctly. schema: Arc, + /// Whether `add` seals before returning. Resolved from the dataset on a + /// reopen, from the caller on a create. + seal_on_add: bool, } impl GenericStore { @@ -116,7 +129,7 @@ impl GenericStore { // key column against whatever it is given, and an empty placeholder // would fail that check. let create_schema = match &spec { - Some(spec) => Arc::new(schema_with_spec(spec)?), + Some(spec) => Arc::new(schema_with_spec(spec, options.seal_on_add)?), None if create_if_missing => { return Err(ArrowError::SchemaError( "a schema is required to create a store".to_string(), @@ -134,6 +147,14 @@ impl GenericStore { } }; + // A reopen honors the mode the store was *created* with; only a create + // takes it from the caller. Read it before opening, since the base + // needs it up front. + let seal_on_add = match &spec { + Some(_) => options.seal_on_add, + None => seal_on_add_from_schema(&create_schema).unwrap_or(options.seal_on_add), + }; + let base = StorageBase::open( uri, StorageBaseOptions { @@ -148,7 +169,7 @@ impl GenericStore { // User schemas are immutable in v1, so there is nothing to // evolve an older base table to. latest_schema: None, - seal_on_put: options.seal_on_add, + seal_on_put: seal_on_add, }, create_if_missing, ) @@ -172,6 +193,7 @@ impl GenericStore { base, spec: persisted, schema, + seal_on_add, }) } @@ -181,6 +203,13 @@ impl GenericStore { &self.spec } + /// Whether `add` seals before returning, so its rows are immediately + /// readable. Persisted with the store, so a reopen preserves it. + #[must_use] + pub fn seal_on_add(&self) -> bool { + self.seal_on_add + } + /// URI of the underlying Lance dataset. #[must_use] pub fn uri(&self) -> &str { @@ -358,17 +387,28 @@ impl GenericStore { } } -/// Attach the serialized spec to the Arrow schema so it round-trips on open. -fn schema_with_spec(spec: &SchemaSpec) -> Result { +/// Attach the serialized spec and the seal mode to the Arrow schema, so both +/// round-trip on open. +fn schema_with_spec(spec: &SchemaSpec, seal_on_add: bool) -> Result { let schema = spec.to_arrow()?; let encoded = serde_json::to_string(spec) .map_err(|error| ArrowError::SchemaError(format!("could not serialize schema: {error}")))?; let mut metadata = schema.metadata().clone(); metadata.insert(SCHEMA_SPEC_KEY.to_string(), encoded); + metadata.insert(SEAL_ON_ADD_KEY.to_string(), seal_on_add.to_string()); Ok(schema.with_metadata(metadata)) } +/// Read the persisted seal mode. `None` for a store written before this was +/// persisted, which then falls back to the caller's option. +fn seal_on_add_from_schema(schema: &Schema) -> Option { + schema + .metadata() + .get(SEAL_ON_ADD_KEY) + .and_then(|raw| raw.parse().ok()) +} + /// Read the spec back out of a dataset's schema metadata. fn spec_from_schema(schema: &Schema) -> Result { let encoded = schema.metadata().get(SCHEMA_SPEC_KEY).ok_or_else(|| { @@ -494,6 +534,63 @@ mod tests { }); } + #[test] + fn seal_mode_survives_a_reopen() { + // The mode is a property of the store, not of whoever opens it. Before + // it was persisted, a reopen silently reverted to the caller's default + // and the store lost read-your-write with no error. + let dir = TempDir::new().unwrap(); + let uri = dir.path().to_string_lossy().to_string(); + let rt = tokio::runtime::Runtime::new().unwrap(); + rt.block_on(async { + { + let store = GenericStore::open(&uri, spec(), sealing()).await.unwrap(); + assert!(store.seal_on_add()); + } + + // Reopen with the *opposite* default: the persisted value wins. + let reopened = GenericStore::open_existing(&uri, GenericStoreOptions::default()) + .await + .unwrap(); + assert!( + reopened.seal_on_add(), + "a store created with seal_on_add must keep it across a reopen" + ); + + // And it is behavioral, not just a flag: the write is visible + // without an explicit flush. + reopened.add(&[row(json!({"id": "r1"}))]).await.unwrap(); + assert_eq!(reopened.list(None, None).await.unwrap().len(), 1); + }); + } + + #[test] + fn deferred_seal_mode_also_survives_a_reopen() { + let dir = TempDir::new().unwrap(); + let uri = dir.path().to_string_lossy().to_string(); + let rt = tokio::runtime::Runtime::new().unwrap(); + rt.block_on(async { + { + GenericStore::open(&uri, spec(), GenericStoreOptions::default()) + .await + .unwrap(); + } + + // Reopen with the opposite default; the persisted `false` wins. + let reopened = GenericStore::open_existing(&uri, sealing()).await.unwrap(); + assert!(!reopened.seal_on_add()); + + reopened.add(&[row(json!({"id": "r1"}))]).await.unwrap(); + assert_eq!( + reopened.list(None, None).await.unwrap().len(), + 0, + "a deferred-seal store must not start sealing after a reopen" + ); + reopened.flush().await.unwrap(); + assert_eq!(reopened.list(None, None).await.unwrap().len(), 1); + }); + } + #[test] fn reopening_with_a_conflicting_schema_is_rejected() { // Silently reinterpreting existing data under a new schema is the diff --git a/crates/lance-context-server/src/config.rs b/crates/lance-context-server/src/config.rs index 49f29d67..509f0ff8 100644 --- a/crates/lance-context-server/src/config.rs +++ b/crates/lance-context-server/src/config.rs @@ -44,16 +44,29 @@ pub struct ServerConfig { /// regardless of count, so stale generations are folded in even on /// low-traffic shards that never cross the count threshold. `0` (the default) /// disables the timer. + /// + /// # Applies to every store kind, despite the name + /// + /// The sweeper visits rollout, datagen and generic stores alike — every + /// MemWAL-backed store accumulates generations, and only a merge reclaims + /// them. The `ROLLOUT_` prefix is historical: these variables predate the + /// other store kinds, and renaming them would silently break existing + /// deployment manifests for no functional gain. #[arg(long, env = "ROLLOUT_CLEANUP_INTERVAL_SECS", default_value = "0")] pub rollout_cleanup_interval_secs: u64, - /// Interval, in seconds, at which the sweeper flushes each resident rollout - /// store's active MemWAL memtable into a queryable generation. Rollout - /// appends are durable on return (the WAL entry is persisted to object - /// storage) but are not visible to reads until the memtable is flushed, so + /// Interval, in seconds, at which the sweeper flushes each resident store's + /// active MemWAL memtable into a queryable generation. A deferred-seal + /// append is durable on return (the WAL entry is persisted to object + /// storage) but is not visible to reads until the memtable is flushed, so /// this interval bounds read-after-write latency. Decoupling flush from the - /// append path is what lets concurrent appends run without serializing behind - /// a per-append seal. Default `30`. + /// append path is what lets concurrent appends run without serializing + /// behind a per-append seal. Default `30`. + /// + /// Like the cleanup interval, this applies to **every** store kind despite + /// the `ROLLOUT_` prefix. Rollout and generic stores default to a deferred + /// seal and depend on it; datagen seals on each append, so its pass is a + /// no-op in steady state. /// /// `0` disables periodic flush, leaving the cleanup sweeper /// (`ROLLOUT_CLEANUP_INTERVAL_SECS`) as the only thing that seals memtables @@ -120,3 +133,57 @@ impl ServerConfig { .filter(|value| !value.is_empty()) } } + +#[cfg(test)] +mod tests { + use super::*; + use clap::CommandFactory; + + /// Every tunable must stay an *optional flag* with a default. + /// + /// Regression guard: dropping a field's `#[arg(..)]` attribute silently + /// turns it into a required positional argument, so the binary still + /// compiles and every unit test still passes — it just refuses to start + /// with the documented flags. That is only caught by actually running the + /// server, which nothing in the Rust suite does. + #[test] + fn every_config_field_is_an_optional_flag() { + let command = ServerConfig::command(); + + let positionals: Vec<_> = command + .get_positionals() + .map(|arg| arg.get_id().to_string()) + .collect(); + assert!( + positionals.is_empty(), + "config fields must be flags, not positional args; found {positionals:?} \ + (a missing #[arg(..)] attribute does this)" + ); + + for arg in command.get_arguments() { + if arg.get_id() == "help" || arg.get_id() == "version" { + continue; + } + assert!( + arg.get_long().is_some(), + "'{}' has no long flag", + arg.get_id() + ); + assert!( + arg.get_default_values().len() == 1 || !arg.is_required_set(), + "'{}' must have a default or be optional", + arg.get_id() + ); + } + } + + /// The server must start with no arguments beyond a data directory — the + /// documented defaults have to actually be defaults. + #[test] + fn parses_with_only_a_data_dir() { + let config = ServerConfig::try_parse_from(["lance-context-server", "--data-dir", "/tmp/x"]) + .expect("the documented minimal invocation must parse"); + assert_eq!(config.rollout_flush_interval_secs, 30); + assert_eq!(config.rollout_cleanup_interval_secs, 0); + } +} diff --git a/crates/lance-context-server/src/main.rs b/crates/lance-context-server/src/main.rs index 102567e7..9ca9f4fc 100644 --- a/crates/lance-context-server/src/main.rs +++ b/crates/lance-context-server/src/main.rs @@ -2,6 +2,7 @@ mod config; mod error; mod routes; mod state; +mod sweeper; use std::sync::Arc; @@ -58,9 +59,10 @@ async fn main() { if state.rollout_flush_interval_secs == 0 && state.rollout_cleanup_interval_secs == 0 { tracing::warn!( "ROLLOUT_FLUSH_INTERVAL_SECS=0 and ROLLOUT_CLEANUP_INTERVAL_SECS=0: \ - nothing will seal MemWAL memtables, so rollout appends will be durable \ - but invisible to reads until this process restarts. Set at least one \ - of them to a non-zero interval." + nothing will seal MemWAL memtables, so deferred-seal appends (rollout \ + and generic stores) will be durable but invisible to reads until this \ + process restarts. Pending WAL generations will also never be merged, \ + for every store kind. Set at least one of them to a non-zero interval." ); } diff --git a/crates/lance-context-server/src/routes/generic.rs b/crates/lance-context-server/src/routes/generic.rs index 959e1d6a..4d50a284 100644 --- a/crates/lance-context-server/src/routes/generic.rs +++ b/crates/lance-context-server/src/routes/generic.rs @@ -464,6 +464,105 @@ mod tests { assert!(matches!(err, AppError::NotFound(_)), "{err:?}"); } + /// The bug this guards: `seal_on_add` used to live only in the open + /// options, so a store created with it lost read-your-write the first time + /// it was evicted from the LRU and reopened — silently, with no error and + /// no data loss, just "sometimes I can't read what I just wrote". + #[tokio::test] + async fn seal_mode_survives_eviction_and_reopen() { + let (state, _dir) = test_state().await; + let _ = create_generic_store( + State(state.clone()), + Json(CreateGenericStoreRequest { + name: "s1".to_string(), + schema: spec(), + storage_options: None, + seal_on_add: true, + }), + ) + .await + .unwrap(); + + // Evict, forcing the next request to reopen from object storage. This + // is exactly what the LRU does under capacity pressure. + state.generic_stores.lock().await.pop("s1"); + + let _ = add_rows( + State(state.clone()), + Path("s1".to_string()), + Json(AddRowsRequest { + rows: vec![row(serde_json::json!({"id": "r1"}))], + }), + ) + .await + .unwrap(); + + let Json(listed) = list_rows( + State(state), + Path("s1".to_string()), + Query(ListRowsQuery { + limit: None, + offset: None, + filter: None, + }), + ) + .await + .unwrap(); + assert_eq!( + listed.rows.len(), + 1, + "a reopened store must keep the seal mode it was created with" + ); + } + + /// The mirror case: a store created with the deferred default must not + /// start sealing after a reopen either. + #[tokio::test] + async fn deferred_seal_mode_also_survives_eviction() { + let (state, _dir) = test_state().await; + let _ = create_generic_store( + State(state.clone()), + Json(CreateGenericStoreRequest { + name: "s1".to_string(), + schema: spec(), + storage_options: None, + seal_on_add: false, + }), + ) + .await + .unwrap(); + + state.generic_stores.lock().await.pop("s1"); + + let _ = add_rows( + State(state.clone()), + Path("s1".to_string()), + Json(AddRowsRequest { + rows: vec![row(serde_json::json!({"id": "r1"}))], + }), + ) + .await + .unwrap(); + + let query = || ListRowsQuery { + limit: None, + offset: None, + filter: None, + }; + let Json(listed) = list_rows(State(state.clone()), Path("s1".to_string()), Query(query())) + .await + .unwrap(); + assert_eq!(listed.rows.len(), 0, "deferred seal must stay deferred"); + + flush_generic_store(State(state.clone()), Path("s1".to_string())) + .await + .unwrap(); + let Json(after) = list_rows(State(state), Path("s1".to_string()), Query(query())) + .await + .unwrap(); + assert_eq!(after.rows.len(), 1); + } + #[tokio::test] async fn deferred_seal_needs_a_flush_and_wal_merges() { let (state, _dir) = test_state().await; diff --git a/crates/lance-context-server/src/state.rs b/crates/lance-context-server/src/state.rs index 51b7b7a0..82cde808 100644 --- a/crates/lance-context-server/src/state.rs +++ b/crates/lance-context-server/src/state.rs @@ -15,6 +15,7 @@ use tokio::task::JoinHandle; use crate::config::ServerConfig; use crate::error::AppError; +use crate::sweeper; /// Default upper bound on resident rollout-store handles when the config does /// not specify one. Sized for peak *concurrent* experiments, not the total @@ -636,9 +637,9 @@ impl AppState { } let uri = self.generic_uri(name); - // `seal_on_add` is a per-store property that is not persisted, so a - // reopened store takes the server default. Callers needing - // read-your-write should flush explicitly. + // The seal mode is persisted in the dataset, so `open_existing` restores + // whatever the store was created with; the option passed here is only a + // fallback for datasets written before that was persisted. let opened = GenericStore::open_existing(&uri, self.generic_store_options(false)) .await .map_err(AppError::from_lance)?; @@ -656,11 +657,14 @@ impl AppState { /// /// This replaces the former one-timer-per-store model, which does not scale /// to hundreds of thousands of per-experiment datasets. On each tick the - /// sweeper snapshots the *resident* rollout stores (those in the LRU) and - /// folds each one's flushed MemWAL generations into its base table, guarding - /// every pass with a timeout so a wedged store cannot stall the sweeper (see - /// [`RolloutStore::cleanup_own_shard`] and the timeout added in the periodic - /// cleanup hardening). + /// sweeper snapshots the *resident* stores of **every kind** (those in the + /// LRUs) and folds each one's flushed MemWAL generations into its base + /// table, guarding every pass with a timeout so a wedged store cannot stall + /// the sweeper. + /// + /// Covering all kinds is not cosmetic: it previously visited only rollout + /// stores, so a datagen store's generations accumulated forever and every + /// read unioned all of them. /// /// Cold stores (evicted from the LRU) are not swept: having no new writes, /// they accumulate no new generations, and are swept again the moment a @@ -683,102 +687,32 @@ impl AppState { let Some(state) = weak.upgrade() else { return; }; - // Snapshot resident stores without holding the LRU lock across - // the awaits below. - let resident: Vec<(String, Arc>)> = { - let cache = state.rollout_stores.lock().await; - cache - .iter() - .map(|(name, store)| (name.clone(), store.clone())) - .collect() - }; - for (name, store) in resident { - // Same read-lock/write-lock split as the flush sweeper: seal - // and read under the shared lock so appends keep flowing, - // then commit under the exclusive lock. Threshold 1 — merge - // whatever is pending, which is what makes this the *time* - // half of the "time OR count" trigger. - let prepared = { - let guard = store.read().await; - match tokio::time::timeout(pass_timeout, guard.prepare_cleanup_merge()) - .await - { - Ok(Ok(prepared)) => prepared, - Ok(Err(e)) => { - metrics::counter!("rollout_wal_cleanup_total", "result" => "failed") - .increment(1); - tracing::warn!( - store = %name, - error = %e, - "global sweeper WAL cleanup failed" - ); - continue; - } - Err(_elapsed) => { - metrics::counter!("rollout_wal_cleanup_total", "result" => "timeout") - .increment(1); - tracing::warn!( - store = %name, - "global sweeper WAL cleanup timed out; abandoning this store this tick" - ); - continue; - } - } - }; - let Some((manifest_store, manifest, prepared)) = prepared else { - continue; - }; - let mut guard = store.write().await; - match tokio::time::timeout( - pass_timeout, - guard.commit_prepared_merge(&manifest_store, &manifest, prepared), - ) - .await - { - Ok(Ok(n)) => { - metrics::counter!("rollout_wal_cleanup_total", "result" => "merged") - .increment(1); - metrics::counter!("rollout_wal_generations_reclaimed_total") - .increment(n as u64); - tracing::info!( - store = %name, - reclaimed = n, - "global sweeper merged flushed generations" - ) - } - Ok(Err(e)) => { - metrics::counter!("rollout_wal_cleanup_total", "result" => "failed") - .increment(1); - tracing::warn!( - store = %name, - error = %e, - "global sweeper WAL cleanup failed" - ) - } - Err(_elapsed) => { - metrics::counter!("rollout_wal_cleanup_total", "result" => "timeout") - .increment(1); - tracing::warn!( - store = %name, - "global sweeper WAL cleanup timed out; abandoning this store this tick" - ) - } - } - } + // Every kind, not just rollout: generations accumulate in each + // one, and only the base table absorbs them. + sweeper::merge_pass(sweeper::resident(&state.rollout_stores).await, pass_timeout) + .await; + sweeper::merge_pass(sweeper::resident(&state.datagen_stores).await, pass_timeout) + .await; + sweeper::merge_pass(sweeper::resident(&state.generic_stores).await, pass_timeout) + .await; } })) } /// Spawn the process-wide MemWAL flush sweeper. /// - /// Rollout appends are durable on return (the WAL entry is persisted) but not - /// visible to reads until the active memtable is flushed into a queryable - /// generation. Rather than flush on every append — which would serialize - /// concurrent writes behind a per-append seal — this sweeper flushes each - /// resident store on a fixed interval, bounding read-after-write latency - /// while keeping the append path concurrent. + /// A deferred-seal append is durable on return (the WAL entry is persisted) + /// but not visible to reads until the active memtable is flushed into a + /// queryable generation. Rather than flush on every append — which would + /// serialize concurrent writes behind a per-append seal — this sweeper + /// flushes each resident store on a fixed interval, bounding + /// read-after-write latency while keeping the append path concurrent. /// - /// After flushing a store it also runs the count-triggered merge + /// Applies to every store kind. Rollout and generic stores default to a + /// deferred seal and genuinely depend on this; datagen seals on each append, + /// so its pass is a no-op in steady state and is kept only for symmetry. + /// + /// For rollout it also runs the count-triggered merge /// ([`RolloutStore::maybe_merge_own_shard`]): the read-amplification bound /// that formerly lived on the append path now rides this timer. The heavier /// time-based cleanup/merge remains on [`Self::spawn_global_sweeper`]. @@ -802,98 +736,19 @@ impl AppState { let Some(state) = weak.upgrade() else { return; }; - let resident: Vec<(String, Arc>)> = { - let cache = state.rollout_stores.lock().await; - cache - .iter() - .map(|(name, store)| (name.clone(), store.clone())) - .collect() - }; - for (name, store) in resident { - // Flush under a read lock so concurrent appends are not blocked. - { - let guard = store.read().await; - match tokio::time::timeout(pass_timeout, guard.flush()).await { - Ok(Ok(())) => {} - Ok(Err(e)) => { - metrics::counter!("rollout_wal_flush_total", "result" => "failed") - .increment(1); - tracing::warn!(store = %name, error = %e, "flush sweeper failed"); - continue; - } - Err(_elapsed) => { - metrics::counter!("rollout_wal_flush_total", "result" => "timeout") - .increment(1); - tracing::warn!(store = %name, "flush sweeper timed out"); - continue; - } - } - metrics::counter!("rollout_wal_flush_total", "result" => "ok").increment(1); - } - // Count-triggered merge (no-op unless the threshold is set - // and met). - // - // Split by lock scope: the expensive half (sealing the - // memtable and reading every flushed generation out of - // object storage) runs under the *read* lock, so appends - // keep flowing through it. Only the short commit — append to - // the base table, drain the manifest, delete the merged - // dirs — takes the write lock. Previously the whole merge - // held the write lock, which stalled every concurrent append - // for its full duration. - let threshold = state.rollout_merge_after_generations; - if threshold == 0 { - continue; - } - let prepared = { - let guard = store.read().await; - match tokio::time::timeout( - pass_timeout, - guard.prepare_merge_if_ready(threshold), - ) - .await - { - Ok(Ok(prepared)) => prepared, - Ok(Err(e)) => { - tracing::warn!(store = %name, error = %e, "flush sweeper merge failed"); - continue; - } - Err(_elapsed) => { - tracing::warn!(store = %name, "flush sweeper merge timed out"); - continue; - } - } - }; - let Some((manifest_store, manifest, prepared)) = prepared else { - continue; - }; - let mut guard = store.write().await; - match tokio::time::timeout( - pass_timeout, - guard.commit_prepared_merge(&manifest_store, &manifest, prepared), - ) - .await - { - Ok(Ok(n)) => { - metrics::counter!("rollout_wal_generations_reclaimed_total") - .increment(n as u64); - tracing::info!( - store = %name, - reclaimed = n, - "flush sweeper count-merged flushed generations" - ); - } - Ok(Err(e)) => { - tracing::warn!(store = %name, error = %e, "flush sweeper merge failed"); - } - Err(_elapsed) => { - tracing::warn!(store = %name, "flush sweeper merge timed out"); - } - } - } + sweeper::flush_pass(sweeper::resident(&state.rollout_stores).await, pass_timeout) + .await; + // Datagen seals on every append, so its pass is a no-op in + // steady state; generic stores default to a deferred seal and + // genuinely depend on this. + sweeper::flush_pass(sweeper::resident(&state.datagen_stores).await, pass_timeout) + .await; + sweeper::flush_pass(sweeper::resident(&state.generic_stores).await, pass_timeout) + .await; } })) } + /// /// [`RolloutStore`]'s writer ([`ShardWriter`]) has no `Drop`, so its /// background tasks are only reclaimed by an explicit `close().await`. On the @@ -1026,6 +881,144 @@ mod tests { handle.abort(); } + /// One minimal, well-formed datagen event. + fn datagen_event() -> lance_context_core::DatagenEvent { + use lance_context_core::{ + datagen_event_id, DatagenEvent, DatagenEventType, DatagenStepKind, + DATAGEN_SCHEMA_VERSION, + }; + DatagenEvent { + event_id: datagen_event_id("item-1", "ckpt-1", 0), + item_id: "item-1".to_string(), + root_item_id: "item-1".to_string(), + parent_item_id: None, + item_seq: 1, + checkpoint_id: "ckpt-1".to_string(), + event_type: DatagenEventType::StepCompleted, + step_name: Some("grade".to_string()), + step_kind: Some(DatagenStepKind::Leaf), + step_index: Some(0), + enclosing_step: None, + selector_step: None, + attempt: 0, + run_id: "run-1".to_string(), + writer_epoch: "writer-1".to_string(), + field_name: None, + field_type: None, + codec_version: None, + value: None, + query_tags: None, + status: None, + error_type: None, + error_dump: None, + traceback: None, + event_ts: chrono::Utc::now(), + schema_version: DATAGEN_SCHEMA_VERSION, + } + } + + /// The sweepers must visit **every** store kind, not just rollout. Before + /// this, a datagen store's flushed generations accumulated forever (nothing + /// merged them) and a generic store with the default deferred seal stayed + /// invisible until someone called `/flush` by hand. + /// + /// Drives one pass directly rather than waiting on the timer, so the test + /// asserts the traversal rather than the tick. + #[tokio::test] + async fn sweeper_passes_cover_datagen_and_generic_stores() { + use lance_context_api::{ColumnSpec, ColumnType, SchemaSpec, ID_COLUMN}; + use lance_context_core::{DatagenStore, GenericStore, GenericStoreOptions}; + use std::time::Duration; + + use crate::sweeper; + + let dir = TempDir::new().unwrap(); + let state = state_with_interval(&dir, 0).await; + let timeout = Duration::from_secs(30); + + // --- generic: deferred seal, so the flush pass is what publishes it --- + let schema = SchemaSpec::new(vec![( + ID_COLUMN.to_string(), + ColumnSpec::required(ColumnType::String { large: false }), + )]); + let generic_uri = state.generic_uri("g1"); + let generic = GenericStore::open( + &generic_uri, + schema, + GenericStoreOptions { + seal_on_add: false, + ..Default::default() + }, + ) + .await + .unwrap(); + let mut rows = serde_json::Map::new(); + rows.insert("id".to_string(), serde_json::json!("r1")); + generic.add(&[rows]).await.unwrap(); + assert_eq!( + generic.list(None, None).await.unwrap().len(), + 0, + "precondition: a deferred-seal write is not yet visible" + ); + + let generic = Arc::new(RwLock::new(generic)); + state + .register_generic("g1", &generic_uri, generic.clone()) + .await + .unwrap(); + + sweeper::flush_pass(sweeper::resident(&state.generic_stores).await, timeout).await; + assert_eq!( + generic.read().await.list(None, None).await.unwrap().len(), + 1, + "the flush pass must publish a generic store's buffered rows" + ); + + // --- datagen: seals on append, so the merge pass is what it needs --- + let datagen_uri = state.datagen_uri("d1"); + let mut datagen = DatagenStore::open(&datagen_uri).await.unwrap(); + datagen.append(&[datagen_event()]).await.unwrap(); + assert!( + datagen.pending_wal_generations().await.unwrap() > 0, + "precondition: datagen seals on append, so a generation is pending" + ); + let datagen = Arc::new(RwLock::new(datagen)); + state + .register_datagen("d1", &datagen_uri, datagen.clone()) + .await + .unwrap(); + + // Both kinds have pending generations now; the merge pass drains them. + sweeper::merge_pass(sweeper::resident(&state.generic_stores).await, timeout).await; + assert_eq!( + generic + .read() + .await + .pending_wal_generations() + .await + .unwrap(), + 0, + "the merge pass must fold a generic store's generations into the base table" + ); + // Rows survive the merge. + assert_eq!( + generic.read().await.list(None, None).await.unwrap().len(), + 1 + ); + + sweeper::merge_pass(sweeper::resident(&state.datagen_stores).await, timeout).await; + assert_eq!( + datagen + .read() + .await + .pending_wal_generations() + .await + .unwrap(), + 0, + "the merge pass must reach datagen stores too" + ); + } + /// `shutdown` drains resident rollout writers by awaiting `close()` on each /// (idempotent: `close` is a no-op when no writer is resident), so /// `ShardWriter` background tasks are reclaimed deterministically instead of diff --git a/crates/lance-context-server/src/sweeper.rs b/crates/lance-context-server/src/sweeper.rs new file mode 100644 index 00000000..0500e257 --- /dev/null +++ b/crates/lance-context-server/src/sweeper.rs @@ -0,0 +1,197 @@ +//! Periodic MemWAL maintenance, uniformly across every store kind. +//! +//! Two things have to happen on a timer for a MemWAL-backed store: +//! +//! - **flush** — seal the active memtable so durable-but-invisible rows become +//! readable. Only matters for a store that defers the seal. +//! - **merge** — fold flushed generations back into the base table. Matters for +//! *every* store: without it, `_mem_wal/` generations accumulate forever and +//! every read unions all of them. +//! +//! Both sweepers used to hardcode `rollout_stores`, so datagen and generic +//! stores were never swept — datagen's generations grew without bound, and a +//! generic store with the default deferred seal stayed invisible until someone +//! called `/flush` by hand. Rather than copy the loop a third time (the exact +//! duplication issue #214 removed one layer down), the traversal is generic +//! over [`Sweepable`] and each store kind supplies a thin impl. + +use std::sync::Arc; +use std::time::Duration; + +use lance_context_core::{DatagenStore, GenericStore, RolloutStore}; +use lru::LruCache; +use tokio::sync::{Mutex, RwLock}; + +/// A store the sweepers can maintain. +/// +/// Implemented on `Arc>` rather than on the store itself so each +/// kind decides its own locking. That is load-bearing for rollout, whose merge +/// deliberately splits into a shared-lock prepare (the expensive object-storage +/// reads, during which appends keep flowing) and a brief exclusive-lock commit. +/// A trait over `&mut Store` would have forced the exclusive lock across the +/// whole merge and quietly stalled the write path. +pub(crate) trait Sweepable: Send + Sync + 'static { + /// Human-readable kind, for log and metric labels. + fn kind() -> &'static str; + + /// Seal the active memtable. A no-op for a store that seals on write. + fn flush(&self) -> impl std::future::Future> + Send; + + /// Fold pending flushed generations into the base table; returns how many + /// were reclaimed. + fn merge_wal(&self) -> impl std::future::Future> + Send; +} + +impl Sweepable for Arc> { + fn kind() -> &'static str { + "rollout" + } + + async fn flush(&self) -> Result<(), String> { + // Read lock: `flush` is `&self`, so concurrent appends are not blocked. + let guard = self.read().await; + let result = guard.flush().await.map_err(|e| e.to_string()); + if result.is_ok() { + // The count-triggered merge rides this timer; it is a no-op unless + // the threshold is configured and met. + drop(guard); + let mut guard = self.write().await; + guard + .maybe_merge_own_shard() + .await + .map_err(|e| e.to_string())?; + } + result + } + + async fn merge_wal(&self) -> Result { + // The prepare/commit split: seal and read every flushed generation + // under the *shared* lock so appends keep running, then take the + // exclusive lock only for the short commit. + let prepared = { + let guard = self.read().await; + guard + .prepare_cleanup_merge() + .await + .map_err(|e| e.to_string())? + }; + let Some((manifest_store, manifest, prepared)) = prepared else { + return Ok(0); + }; + let mut guard = self.write().await; + guard + .commit_prepared_merge(&manifest_store, &manifest, prepared) + .await + .map_err(|e| e.to_string()) + } +} + +impl Sweepable for Arc> { + fn kind() -> &'static str { + "datagen" + } + + async fn flush(&self) -> Result<(), String> { + // Datagen seals on every append, so this is a no-op in steady state. + // Kept for symmetry, and it still drains anything a fenced writer left. + let guard = self.read().await; + guard.flush().await.map_err(|e| e.to_string()) + } + + async fn merge_wal(&self) -> Result { + let mut guard = self.write().await; + guard.cleanup_own_shard().await.map_err(|e| e.to_string()) + } +} + +impl Sweepable for Arc> { + fn kind() -> &'static str { + "generic" + } + + async fn flush(&self) -> Result<(), String> { + let guard = self.read().await; + guard.flush().await.map_err(|e| e.to_string()) + } + + async fn merge_wal(&self) -> Result { + let mut guard = self.write().await; + guard.cleanup_wal().await.map_err(|e| e.to_string()) + } +} + +/// Snapshot a cache's resident entries without holding its lock across the +/// awaits that follow. +pub(crate) async fn resident(cache: &Mutex>) -> Vec<(String, S)> { + cache + .lock() + .await + .iter() + .map(|(name, store)| (name.clone(), store.clone())) + .collect() +} + +/// Flush every resident store of one kind, bounding each by `pass_timeout` so a +/// single wedged store cannot stall the rest. +/// +/// Metric names keep their historical `rollout_` prefix so existing dashboards +/// and alerts keep working; the new `kind` label is what distinguishes the +/// store types. Renaming them would be a silent breakage for anyone graphing +/// these today. +pub(crate) async fn flush_pass(stores: Vec<(String, S)>, pass_timeout: Duration) { + let kind = S::kind(); + for (name, store) in stores { + match tokio::time::timeout(pass_timeout, store.flush()).await { + Ok(Ok(())) => { + metrics::counter!("rollout_wal_flush_total", "result" => "ok", "kind" => kind) + .increment(1); + } + Ok(Err(error)) => { + metrics::counter!("rollout_wal_flush_total", "result" => "failed", "kind" => kind) + .increment(1); + tracing::warn!(store = %name, kind, %error, "flush sweeper failed"); + } + Err(_elapsed) => { + metrics::counter!("rollout_wal_flush_total", "result" => "timeout", "kind" => kind) + .increment(1); + tracing::warn!(store = %name, kind, "flush sweeper timed out"); + } + } + } +} + +/// Merge every resident store's pending generations, same timeout discipline. +pub(crate) async fn merge_pass(stores: Vec<(String, S)>, pass_timeout: Duration) { + let kind = S::kind(); + for (name, store) in stores { + match tokio::time::timeout(pass_timeout, store.merge_wal()).await { + Ok(Ok(0)) => {} + Ok(Ok(reclaimed)) => { + metrics::counter!("rollout_wal_cleanup_total", "result" => "merged", "kind" => kind) + .increment(1); + metrics::counter!("rollout_wal_generations_reclaimed_total", "kind" => kind) + .increment(reclaimed as u64); + tracing::info!( + store = %name, + kind, + reclaimed, + "sweeper merged flushed generations" + ); + } + Ok(Err(error)) => { + metrics::counter!("rollout_wal_cleanup_total", "result" => "failed", "kind" => kind) + .increment(1); + tracing::warn!(store = %name, kind, %error, "sweeper WAL cleanup failed"); + } + Err(_elapsed) => { + metrics::counter!("rollout_wal_cleanup_total", "result" => "timeout", "kind" => kind) + .increment(1); + tracing::warn!( + store = %name, + kind, + "sweeper WAL cleanup timed out; abandoning this store this tick" + ); + } + } + } +} diff --git a/python/tests/test_generic_store.py b/python/tests/test_generic_store.py index ef92d8ae..636e60a9 100644 --- a/python/tests/test_generic_store.py +++ b/python/tests/test_generic_store.py @@ -190,3 +190,32 @@ def test_shorthand_and_full_type_forms_agree(tmp_path: Path) -> None: ) shorthand.add([{"id": "a", "n": 7}]) assert shorthand.list()[0]["n"] == 7 + + +def test_seal_mode_survives_a_reopen(tmp_path: Path) -> None: + # `seal_on_add` is a property of the store, not of whoever opens it. It used + # to live only in the open options, so a reopened store silently reverted to + # the caller's default and lost read-your-write with no error. + uri = str(tmp_path / "sealed.lance") + store = GenericStore.open(uri, schema=SCHEMA, seal_on_add=True) + store.add([{"id": "r1"}]) + del store + + # Reopen with the opposite default: the persisted value must win. + reopened = GenericStore.open_existing(uri, seal_on_add=False) + reopened.add([{"id": "r2"}]) + assert len(reopened.list()) == 2, ( + "a store created with seal_on_add must keep it across a reopen" + ) + + +def test_deferred_seal_mode_survives_a_reopen(tmp_path: Path) -> None: + uri = str(tmp_path / "deferred.lance") + GenericStore.open(uri, schema=SCHEMA, seal_on_add=False) + + reopened = GenericStore.open_existing(uri, seal_on_add=True) + reopened.add([{"id": "r1"}]) + assert reopened.list() == [], "a deferred-seal store must stay deferred" + + reopened.flush() + assert len(reopened.list()) == 1