diff --git a/.github/workflows/rust-test.yml b/.github/workflows/rust-test.yml index b8916625..6271fbd1 100644 --- a/.github/workflows/rust-test.yml +++ b/.github/workflows/rust-test.yml @@ -158,6 +158,10 @@ jobs: run: | sudo apt update sudo apt install -y clang libclang-dev protobuf-compiler + - name: Rollout materialization through the WAL client + run: | + cargo clippy -p lance-context-ingest --features materializer --all-targets -- -D warnings + cargo test -p lance-context-ingest --features materializer --lib - name: Start etcd for HA scheduler test run: | ETCD_VERSION=3.7.0 diff --git a/Cargo.lock b/Cargo.lock index 83f3d314..03a07f0a 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4184,10 +4184,12 @@ dependencies = [ "arrow-schema", "base64 0.22.1", "bytes", + "futures", "lance", "lance-context-core", "lance-context-wal", "lance-table", + "object_store 0.14.2", "rcgen", "ring", "serde", diff --git a/crates/lance-context-ingest/Cargo.toml b/crates/lance-context-ingest/Cargo.toml index 83d723b3..c2575654 100644 --- a/crates/lance-context-ingest/Cargo.toml +++ b/crates/lance-context-ingest/Cargo.toml @@ -6,7 +6,14 @@ license = "Apache-2.0" publish = false description = "Durable rollout ingress over the independent lance-context WAL" +[features] +materializer = ["lance-context-wal/lance-integration", "dep:lance", "dep:lance-table", "dep:object_store", "dep:futures"] + [dependencies] +lance = { version = "13.0.0", optional = true } +lance-table = { version = "13.0.0", optional = true } +object_store = { version = "0.14.2", optional = true } +futures = { version = "0.3", optional = true } anyhow = "1" arrow-array = "58" arrow-schema = "58" diff --git a/crates/lance-context-ingest/src/lib.rs b/crates/lance-context-ingest/src/lib.rs index 222ae07b..6de32eca 100644 --- a/crates/lance-context-ingest/src/lib.rs +++ b/crates/lance-context-ingest/src/lib.rs @@ -5,7 +5,7 @@ //! migrate or replace legacy merge/upsert tables, or authenticate an HTTP caller. //! One atomic append must fit one virtual bucket; reject mixed buckets before ACK. use anyhow::{bail, ensure, Context, Result}; -use arrow_array::{Array, StringArray}; +use arrow_array::{Array, RecordBatch, StringArray}; use arrow_schema::SchemaRef; use bytes::Bytes; use lance_context_core::{rollout_schema, MergeMemoryBudget, MergeReservation}; @@ -29,6 +29,8 @@ use serde::{Deserialize, Serialize}; use std::{collections::HashSet, sync::Arc, time::Duration}; use uuid::Uuid; +#[cfg(feature = "materializer")] +pub mod materializer; pub mod receipt; /// Resolved by the application's authorized catalog, not an uploaded URI/UUID. @@ -500,12 +502,39 @@ fn prepare( layout: &Layout, ) -> Result { let batches = ipc::decode(body.clone(), schema)?; + let rows = validate_rows( + &batches, + &upload.scope.contract, + layout, + producer.open.bucket, + )?; + let batch = Batch::new( + upload.scope.table, + producer.open.bucket, + producer.id, + sequence, + body, + ); + Ok(PreparedAppend { + scope: upload.scope, + batch, + rows, + _reservation: upload.reservation, + }) +} + +fn validate_rows( + batches: &[RecordBatch], + contract: &IngestContract, + layout: &Layout, + expected_bucket: u32, +) -> Result { let mut ids = HashSet::new(); let mut rows = 0u32; let mut previous_session = None; - for batch in &batches { + for batch in batches { ensure!( - batch.num_rows() <= (upload.scope.contract.max_rows - rows) as usize, + batch.num_rows() <= (contract.max_rows - rows) as usize, "upload exceeds registered row limit" ); for (field, column) in batch.schema().fields().iter().zip(batch.columns()) { @@ -532,13 +561,10 @@ fn prepare( // Consecutive rows in a rollout share routing. Borrow the decoded // string and hash only when the session changes; no string copy. if previous_session != Some(session) { - let bucket = layout.session_bucket( - &upload.scope.contract.tenant, - &upload.scope.contract.qualified_table, - session, - )?; + let bucket = + layout.session_bucket(&contract.tenant, &contract.qualified_table, session)?; ensure!( - bucket == producer.open.bucket, + bucket == expected_bucket, "atomic append spans producer bucket" ); previous_session = Some(session); @@ -547,19 +573,7 @@ fn prepare( } } ensure!(rows > 0, "empty append"); - let batch = Batch::new( - upload.scope.table, - producer.open.bucket, - producer.id, - sequence, - body, - ); - Ok(PreparedAppend { - scope: upload.scope, - batch, - rows, - _reservation: upload.reservation, - }) + Ok(rows) } #[cfg(test)] diff --git a/crates/lance-context-ingest/src/materializer.rs b/crates/lance-context-ingest/src/materializer.rs new file mode 100644 index 00000000..488aa429 --- /dev/null +++ b/crates/lance-context-ingest/src/materializer.rs @@ -0,0 +1,660 @@ +//! Recoverable publication of session-sharded rollout records into real Lance. +//! Scheduling ownership is supplied by the application. Durable catalog epochs, +//! fixed transaction targets and signed source closure enforce correctness. +use super::*; +use futures::{stream, StreamExt, TryStreamExt}; +use lance::{dataset::builder::DatasetBuilder, Dataset}; +use lance_context_wal::replication::{publication as p, source}; +use object_store::path::Path as ObjectPath; +use std::collections::{BTreeMap, HashMap}; + +mod stage; + +/// Resolve this from the authorized namespace, never from an append request. +#[derive(Clone)] +pub struct Storage { + pub uri: String, + pub base: ObjectPath, + pub options: HashMap, + pub session: Arc, +} + +#[derive(Clone, Debug)] +pub struct Options { + pub max_source_bytes: usize, + pub max_records: usize, + pub fetch_parallelism: usize, + pub write_parallelism: usize, + /// Each poll bounds scanning of other tables sharing the same source group. + pub scan_pages_per_group: usize, + /// Catalog metadata only; retiring it never deletes immutable blob objects. + pub retained_manifest_versions: usize, +} +impl Default for Options { + fn default() -> Self { + Self { + max_source_bytes: 64 << 20, + max_records: 128, + fetch_parallelism: 4, + write_parallelism: 2, + scan_pages_per_group: 4, + retained_manifest_versions: 64, + } + } +} +impl Options { + fn workspace(&self, contract: &IngestContract, bytes: usize, records: usize) -> Result { + bytes + .checked_mul(3) + .and_then(|v| { + records + .checked_mul(contract.max_rows as usize) + .and_then(|n| n.checked_mul(64)) + .and_then(|n| v.checked_add(n)) + }) + .context("materializer workspace overflow") + } + fn validate(&self, ingress: &Ingress, budget: &MergeMemoryBudget) -> Result<()> { + ensure!( + self.max_source_bytes >= ingress.scope.contract.max_record_bytes as usize + && self.max_source_bytes as u64 <= ingress.client.max_plan_source_bytes(), + "materializer source cap differs from WAL/ingress policy" + ); + ensure!( + (1..=p::MAX_INPUTS).contains(&self.max_records) + && (1..=8).contains(&self.fetch_parallelism) + && (1..=8).contains(&self.write_parallelism) + && (1..=32).contains(&self.scan_pages_per_group) + && (2..=p::MAX_MANIFEST_VERSIONS - 2).contains(&self.retained_manifest_versions), + "invalid materializer concurrency or history bounds" + ); + ensure!( + self.workspace( + &ingress.scope.contract, + self.max_source_bytes, + self.max_records + )? <= budget.limit(), + "materializer maximum workspace exceeds shared budget" + ); + Ok(()) + } +} + +/// Persist this acquisition identity before its first RPC, or derive it from +/// an existing durable scheduling claim. A timeout does not justify a takeover. +#[derive(Clone, Debug)] +pub struct Acquisition { + pub owner: Uuid, + pub request: Uuid, + pub expected_epoch: u64, +} + +pub struct Publisher { + ingress: Ingress, + storage: Storage, + options: Options, + budget: Arc, + authority: p::Authority, + // Only skip pages proven to contain no records for this table. This hint + // never advances a durable source checkpoint or skips a selected record. + empty_cursors: BTreeMap, + plan_attempt: Option, + prepared_attempt: Option, +} + +impl Publisher { + pub async fn acquire( + ingress: Ingress, + storage: Storage, + options: Options, + budget: Arc, + acquisition: Acquisition, + ) -> Result { + options.validate(&ingress, &budget)?; + ensure!( + !storage.uri.is_empty() && !acquisition.owner.is_nil() && !acquisition.request.is_nil(), + "invalid publisher configuration" + ); + let applied = commit( + &ingress.client, + Command::Acquire(p::Acquire { + table: ingress.scope.table, + destination: ingress.scope.destination.clone(), + ingest: ingress.scope.contract.clone(), + owner: acquisition.owner, + request: acquisition.request, + expected_epoch: acquisition.expected_epoch, + }), + ) + .await?; + let Applied::Publisher(p::Reply::Acquired(authority)) = applied else { + bail!("publisher acquisition rejected: {applied:?}"); + }; + ensure!( + authority.cluster == ingress.scope.cluster + && authority.layout == ingress.scope.layout + && authority.group == ingress.client.catalog_group() + && authority.table == ingress.scope.table + && authority.owner == acquisition.owner + && authority.epoch + == acquisition + .expected_epoch + .checked_add(1) + .context("publisher epoch exhausted")?, + "foreign publisher authority" + ); + let value = Self { + ingress, + storage, + options, + budget, + authority, + empty_cursors: BTreeMap::new(), + plan_attempt: None, + prepared_attempt: None, + }; + value.state().await?; + Ok(value) + } + pub fn authority(&self) -> &p::Authority { + &self.authority + } + pub async fn state(&self) -> Result> { + let query::Value::Publisher(Some(view)) = query( + &self.ingress, + self.authority.group, + query::Query::Publisher { + table: self.authority.table, + }, + ) + .await? + else { + bail!("publisher binding unavailable"); + }; + ensure!( + view.authority == self.authority + && view.destination == self.ingress.scope.destination + && view.ingest == self.ingress.scope.contract + && view + .manifests + .as_ref() + .is_some_and(|m| m.incarnation == self.ingress.scope.incarnation), + "publisher fenced or table incarnation changed" + ); + Ok(view) + } + fn reader(&self) -> Result { + p::lance::ManagedManifest::reader( + Arc::new(self.ingress.client.manifest_catalog()), + self.storage.base.clone(), + &self.ingress.layout, + self.authority.table, + self.ingress.scope.incarnation, + ) + } + /// Readers of this managed table must use its catalog-selected manifest. + pub async fn read_dataset(&self, version: Option) -> Result { + let mut builder = DatasetBuilder::from_uri(&self.storage.uri) + .with_storage_options(self.storage.options.clone()) + .with_session(self.storage.session.clone()) + .with_commit_handler(Arc::new(self.reader()?)); + if let Some(version) = version { + builder = builder.with_version(version); + } + let dataset = builder.load().await?; + let schema: arrow_schema::Schema = dataset.schema().into(); + ensure!( + schema == rollout_schema(), + "managed dataset is not the bound rollout schema" + ); + Ok(dataset) + } + + /// Select a complete bounded prefix, pin every source, write immutable Lance + /// files and record descriptors. Existing prepared files are reused. + pub async fn prepare_next(&mut self) -> Result>> { + let mut view = self.state().await?; + if view.progress.pending.is_none() { + self.retire_history(&view).await?; + if self.plan_attempt.is_none() { + self.plan_attempt = self.select(&view).await?; + } + let Some(plan) = self.plan_attempt.clone() else { + return Ok(None); + }; + let reply = commit( + &self.ingress.client, + Command::Plan(p::PlanRequest { + authority: self.authority.clone(), + plan, + }), + ) + .await?; + ensure!( + matches!(reply, Applied::Publisher(p::Reply::Planned { .. })), + "publication selection rejected: {reply:?}" + ); + view = self.state().await?; + } + self.plan_attempt = None; + if let Some(prepared) = view.progress.prepared { + self.prepared_attempt = None; + return Ok(Some(prepared)); + } + let pending = view.progress.pending.context("publication plan missing")?; + ensure!( + pending.plan.source_bytes()? <= self.options.max_source_bytes as u64 + && pending + .plan + .sources + .values() + .map(|s| s.records.len()) + .sum::() + <= self.options.max_records, + "recovered plan exceeds this worker's configured limits" + ); + if self + .prepared_attempt + .as_ref() + .is_some_and(|p| p.plan_index != pending.index) + { + self.prepared_attempt = None; + } + if self.prepared_attempt.is_none() { + self.pin(&pending).await?; + let dataset = self.read_dataset(Some(pending.plan.read_version)).await?; + let bytes = usize::try_from(pending.plan.source_bytes()?)?; + let count = pending.plan.sources.values().map(|s| s.records.len()).sum(); + let reservation = Arc::new( + self.budget + .reserve( + self.options + .workspace(&self.ingress.scope.contract, bytes, count)?, + ) + .await, + ); + let descriptors = stage::materialize( + self.ingress.clone(), + self.options.clone(), + pending.clone(), + dataset, + reservation, + ) + .await?; + descriptors.validate(self.ingress.client.max_descriptor_bytes())?; + self.prepared_attempt = Some(p::PrepareRequest { + authority: self.authority.clone(), + transaction: pending.plan.transaction, + plan_index: pending.index, + descriptors, + }); + } + let reply = commit( + &self.ingress.client, + Command::Prepare(self.prepared_attempt.as_ref().unwrap().clone()), + ) + .await?; + ensure!( + matches!(reply, Applied::Publisher(p::Reply::Prepared { .. })), + "materialization result rejected: {reply:?}" + ); + Ok(self.state().await?.progress.prepared) + } + + /// Publish the fixed target and confirm its exact identity. Errors preserve + /// both the prepared files and source pins for explicit reconciliation. + pub async fn commit_pending(&self) -> Result>> { + let view = self.state().await?; + if let Some(published) = view.progress.published { + return Ok(Some(published)); + } + let Some(pending) = view.progress.pending else { + return Ok(None); + }; + let prepared = view + .progress + .prepared + .context("pending plan has no prepared files")?; + let manifests = view.manifests.context("managed manifest index missing")?; + let target = pending.plan.target_version()?; + if !manifests.versions.contains_key(&target) { + let predecessor = manifests + .versions + .get(&pending.plan.read_version) + .context("publication predecessor unavailable")?; + let base = self.read_dataset(Some(pending.plan.read_version)).await?; + stage::publish(self, &pending, &prepared, predecessor, base).await?; + } + // A lost commit reply is reconciled by this quorum-selected record; + // absence on a local replica is never treated as a failed publication. + let selected = self + .state() + .await? + .manifests + .context("manifest index missing")? + .versions + .get(&target) + .cloned() + .context("fixed target not selected")?; + ensure!( + selected.transaction == pending.plan.transaction + && selected.plan_index == pending.index + && selected.prepared_index == prepared.index + && selected.descriptor_digest == prepared.descriptors.digest, + "fixed target belongs to another publication" + ); + let reply = commit( + &self.ingress.client, + Command::Confirm(p::ConfirmPublication { + authority: self.authority.clone(), + incarnation: self.ingress.scope.incarnation, + transaction: pending.plan.transaction, + plan_index: pending.index, + prepared_index: prepared.index, + descriptor_digest: prepared.descriptors.digest, + target, + manifest_identity: selected.identity, + }), + ) + .await?; + let Applied::Publisher(p::Reply::Published(published)) = reply else { + bail!("publication confirmation rejected: {reply:?}"); + }; + Ok(Some(published)) + } + + /// Retire payloads only with catalog-issued closure authorization, then + /// advance independent source cursors using their quorum closure proofs. + pub async fn finish_pending(&self) -> Result>> { + let view = self.state().await?; + let Some(pending) = view.progress.pending else { + return Ok(None); + }; + let published = view + .progress + .published + .context("publication has not been confirmed")?; + let published = &published; + let pending = &pending; + let proofs = stream::iter( + pending + .plan + .sources + .keys() + .copied() + .map(|group| async move { + let query::Value::ClosePermit(Some(permit)) = query( + &self.ingress, + self.authority.group, + query::Query::ClosePermit { + table: self.authority.table, + transaction: pending.plan.transaction, + plan_index: pending.index, + source_group: group, + }, + ) + .await? + else { + bail!("source closure authorization unavailable"); + }; + let expected = permit.closure().clone(); + let closed = commit(&self.ingress.client, Command::Close(*permit)).await?; + ensure!( + matches!(closed, Applied::Source(source::Reply::Closed { .. })), + "source closure rejected: {closed:?}" + ); + let query::Value::ClosedProof(Some(proof)) = query( + &self.ingress, + group, + query::Query::ClosedProof { + table: self.authority.table, + }, + ) + .await? + else { + bail!("source closure proof unavailable"); + }; + ensure!( + proof.closed().permit.closure() == &expected + && expected.published_index == published.index + && proof.source_cut().group == group, + "foreign source closure proof" + ); + Ok::<_, anyhow::Error>((group, *proof)) + }), + ) + .buffer_unordered(self.options.fetch_parallelism) + .try_collect::>() + .await?; + let reply = commit( + &self.ingress.client, + Command::Advance(p::AdvancePublication { + authority: self.authority.clone(), + published_index: published.index, + proofs, + }), + ) + .await?; + let Applied::Publisher(p::Reply::Advanced(advanced)) = reply else { + bail!("publication advancement rejected: {reply:?}"); + }; + Ok(Some(advanced)) + } + pub async fn run_once(&mut self) -> Result>> { + if self.prepare_next().await?.is_none() { + return Ok(None); + } + self.commit_pending().await?; + self.finish_pending().await + } + + async fn pin(&self, pending: &p::PendingPlan) -> Result<()> { + stream::iter( + pending + .plan + .sources + .keys() + .copied() + .map(|group| async move { + let query::Value::Pin(Some(permit)) = query( + &self.ingress, + self.authority.group, + query::Query::Pin { + table: self.authority.table, + transaction: pending.plan.transaction, + plan_index: pending.index, + source_group: group, + }, + ) + .await? + else { + bail!("source pin authorization unavailable"); + }; + let selected = permit.selection(); + ensure!( + selected.table == self.authority.table + && selected.transaction == pending.plan.transaction + && selected.plan_index == pending.index + && selected.read_version == pending.plan.read_version + && selected.group == group + && selected.ingest == self.ingress.scope.contract + && selected.range == pending.plan.sources[&group], + "pin differs from selected prefix" + ); + let reply = commit(&self.ingress.client, Command::Pin(*permit)).await?; + ensure!( + matches!(reply, Applied::Source(source::Reply::Pinned { .. })), + "source pin rejected: {reply:?}" + ); + Ok::<_, anyhow::Error>(()) + }), + ) + .buffer_unordered(self.options.fetch_parallelism) + .try_collect::>() + .await?; + Ok(()) + } + async fn retire_history(&self, view: &p::Table) -> Result<()> { + let manifests = view + .manifests + .as_ref() + .context("manifest history missing")?; + if manifests.versions.len() <= self.options.retained_manifest_versions { + return Ok(()); + } + let retain_from = manifests + .head + .saturating_sub(self.options.retained_manifest_versions as u64 - 2) + .max(manifests.initial_version); + let reply = commit( + &self.ingress.client, + Command::RetireHistory(p::RetireManifestHistory { + authority: self.authority.clone(), + incarnation: self.ingress.scope.incarnation, + retain_from, + }), + ) + .await?; + ensure!( + matches!(reply, Applied::Publisher(p::Reply::ManifestHistory(_))), + "manifest history retirement rejected: {reply:?}" + ); + Ok(()) + } + async fn select(&mut self, view: &p::Table) -> Result> { + let mut groups = self + .ingress + .layout + .groups + .keys() + .copied() + .collect::>(); + if let Some(last) = view + .progress + .advanced + .as_ref() + .and_then(|a| a.sources.keys().next_back()) + { + let offset = groups.iter().position(|g| g > last).unwrap_or(0); + groups.rotate_left(offset); + } + let mut sources = BTreeMap::new(); + let mut bytes = 0usize; + let mut count = 0usize; + let mut next_group = groups[0]; + for (ordinal, &group) in groups.iter().enumerate() { + let after = view.progress.checkpoint.get(&group).copied().unwrap_or(0); + let mut cursor = self + .empty_cursors + .get(&group) + .copied() + .unwrap_or(after) + .max(after); + let mut records = Vec::new(); + let mut through = cursor; + for _ in 0..self.options.scan_pages_per_group { + let query::Value::Accepted(page) = query( + &self.ingress, + group, + query::Query::Accepted { + after: cursor, + limit: 128, + }, + ) + .await? + else { + bail!("source page unavailable at quorum"); + }; + let mut stopped = false; + for (index, meta) in page.entries { + if meta.stream != self.authority.table { + continue; + } + let size = usize::try_from(meta.body_bytes)?; + ensure!( + meta.kind == RecordKind::Data + && size <= self.ingress.scope.contract.max_record_bytes as usize + && self.ingress.layout.group_for_bucket(meta.bucket)? == group, + "foreign source metadata" + ); + if size > self.options.max_source_bytes - bytes + || count == self.options.max_records + { + stopped = true; + break; + } + bytes += size; + count += 1; + through = index; + records.push((index, meta)); + } + if stopped { + break; + } + through = page.through; + if records.is_empty() { + self.empty_cursors.insert(group, through); + } + if page.through == cursor || !records.is_empty() { + break; + } + cursor = page.through; + } + if !records.is_empty() { + sources.insert( + group, + p::SourceRange { + after, + through, + records, + }, + ); + next_group = groups[(ordinal + 1) % groups.len()]; + } + if bytes == self.options.max_source_bytes || count == self.options.max_records { + break; + } + } + if sources.is_empty() { + return Ok(None); + } + Ok(Some(p::Plan { + transaction: Uuid::new_v4(), + read_version: view.progress.minimum_read_version, + sources, + next_group, + })) + } +} + +async fn query(ingress: &Ingress, group: u64, value: query::Query) -> Result { + let response = ingress + .client + .request_routed(Request::Query { + group, + query: value, + }) + .await?; + let Reply::Query(query::Outcome::Observed { cut, value }) = response.reply else { + bail!("quorum observation unavailable"); + }; + ensure!( + cut.cluster == ingress.scope.cluster + && cut.layout == ingress.scope.layout + && cut.group == group + && cut.term > 0 + && cut.quorum_index > 0 + && cut.applied >= cut.quorum_index, + "foreign quorum observation" + ); + Ok(value) +} +async fn commit(client: &Client, command: Command) -> Result { + let response = client + .request_routed(Request::Commit(Box::new(command))) + .await?; + let Reply::Commit(Outcome::Committed(applied)) = response.reply else { + bail!("durable command outcome unknown; reconcile the original operation"); + }; + Ok(applied) +} diff --git a/crates/lance-context-ingest/src/materializer/stage.rs b/crates/lance-context-ingest/src/materializer/stage.rs new file mode 100644 index 00000000..da9dcdea --- /dev/null +++ b/crates/lance-context-ingest/src/materializer/stage.rs @@ -0,0 +1,156 @@ +use super::*; +use lance::dataset::{ + transaction::{Operation, Transaction}, + CommitBuilder, InsertBuilder, WriteMode, WriteParams, +}; +use lance_table::format::Fragment; + +/// A cancelled caller detaches this bounded job rather than cancelling its +/// admitted file writers. The reservation stays with those writers until done. +pub(super) async fn materialize( + ingress: Ingress, + options: Options, + pending: Arc, + dataset: Dataset, + reservation: Arc, +) -> Result { + tokio::spawn(async move { + let _reservation = reservation.clone(); + let inputs = pending + .plan + .sources + .iter() + .flat_map(|(&group, range)| { + range + .records + .iter() + .map(move |(index, meta)| (group, *index, meta.clone())) + }) + .collect::>(); + let decoded = stream::iter(inputs.into_iter().map(|(group, index, meta)| { + let ingress = ingress.clone(); + let reservation = reservation.clone(); + async move { + // Fetch is a local committed read. Retry other configured voters + // only for this same pinned record; none can change its identity. + let mut fetched = None; + for &node in &ingress.layout.groups[&group] { + if let Ok(response) = ingress + .client + .request(node, Request::Fetch { group, index }) + .await + { + if let Reply::Record(Some(observed)) = &response.reply { + ensure!( + observed == &meta, + "fetched record differs from durable source pin" + ); + fetched = Some(response); + break; + } + } + } + let response = fetched.context("pinned source unavailable on every voter")?; + tokio::task::spawn_blocking(move || { + let _reservation = reservation; + meta.validate_payload( + &response.body, + ingress.scope.contract.max_record_bytes as usize, + )?; + let batches = ipc::decode(response.body.clone(), ingress.schema.clone())?; + validate_rows( + &batches, + &ingress.scope.contract, + &ingress.layout, + meta.bucket, + )?; + Ok::<_, anyhow::Error>(batches) + }) + .await? + } + })) + .buffered(options.fetch_parallelism) + .try_collect::>() + .await?; + // Share Arrow buffers across a bounded set of file writers. Each source + // record remains whole; writer count cannot change retry/plan identity. + let writers = options.write_parallelism.min(decoded.len()); + ensure!(writers > 0, "empty pinned source selection"); + let mut partitions = (0..writers).map(|_| Vec::new()).collect::>(); + for (i, batches) in decoded.into_iter().enumerate() { + partitions[i % writers].extend(batches); + } + let dataset = Arc::new(dataset); + let tasks = partitions.into_iter().map(|batches| { + let dataset = dataset.clone(); + let reservation = reservation.clone(); + async move { + let _reservation = reservation; + let staged = InsertBuilder::new(dataset) + .with_params(&WriteParams { + mode: WriteMode::Append, + ..Default::default() + }) + .execute_uncommitted(batches) + .await?; + let Operation::Append { fragments } = staged.operation else { + bail!("materialization did not produce append fragments"); + }; + Ok::<_, anyhow::Error>(fragments) + } + }); + // Join every admitted writer even if another writer fails. Completed + // orphan files are retained; no unverified deletion is attempted here. + let results = futures::future::join_all(tasks).await; + let fragments = results + .into_iter() + .collect::>>()? + .into_iter() + .flatten() + .collect::>(); + p::Descriptors::lance_fragments(Bytes::from(serde_json::to_vec(&fragments)?)) + }) + .await? +} + +pub(super) async fn publish( + publisher: &Publisher, + pending: &p::PendingPlan, + prepared: &p::Prepared, + predecessor: &p::ManifestRecord, + base: Dataset, +) -> Result<()> { + prepared + .descriptors + .validate(publisher.ingress.client.max_descriptor_bytes())?; + let fragments: Vec = serde_json::from_slice(&prepared.descriptors.bytes)?; + ensure!(!fragments.is_empty(), "empty prepared fragment list"); + let mut transaction = Transaction::new( + pending.plan.read_version, + Operation::Append { fragments }, + None, + ); + transaction.uuid = pending.plan.transaction.to_string(); + let handler = p::lance::ManagedManifest::new( + Arc::new(publisher.ingress.client.manifest_catalog()), + publisher.storage.base.clone(), + &publisher.ingress.layout, + p::lance::ManifestCommit { + authority: publisher.authority.clone(), + incarnation: publisher.ingress.scope.incarnation, + transaction: pending.plan.transaction, + plan_index: pending.index, + prepared_index: prepared.index, + descriptor_digest: prepared.descriptors.digest, + predecessor: predecessor.identity, + target: pending.plan.target_version()?, + }, + )?; + CommitBuilder::new(Arc::new(base)) + .with_commit_handler(Arc::new(handler)) + .with_max_retries(0) + .with_skip_auto_cleanup(true) + .execute(transaction) + .await?; + Ok(()) +} diff --git a/crates/lance-context-ingest/src/tests.rs b/crates/lance-context-ingest/src/tests.rs index d6d2473e..74e53740 100644 --- a/crates/lance-context-ingest/src/tests.rs +++ b/crates/lance-context-ingest/src/tests.rs @@ -10,6 +10,8 @@ use lance_table::io::commit::{ConditionalPutCommitHandler, ManifestNamingScheme} use receipt::ReceiptKeys; use std::collections::BTreeMap; +#[cfg(feature = "materializer")] +mod materializer; mod network; fn layout() -> Arc { diff --git a/crates/lance-context-ingest/src/tests/materializer.rs b/crates/lance-context-ingest/src/tests/materializer.rs new file mode 100644 index 00000000..87789de3 --- /dev/null +++ b/crates/lance-context-ingest/src/tests/materializer.rs @@ -0,0 +1,291 @@ +use super::*; +use crate::materializer::{Acquisition, Options, Publisher, Storage}; +use futures::TryStreamExt; +use object_store::path::Path as ObjectPath; +use std::collections::HashMap; + +async fn publisher( + ingress: &Ingress, + storage: &Storage, + budget: &Arc, + epoch: u64, +) -> Result { + Publisher::acquire( + ingress.clone(), + storage.clone(), + Options { + max_source_bytes: 1 << 20, + max_records: 8, + write_parallelism: 2, + retained_manifest_versions: 2, + ..Options::default() + }, + budget.clone(), + Acquisition { + owner: Uuid::new_v4(), + request: Uuid::new_v4(), + expected_epoch: epoch, + }, + ) + .await +} +async fn append( + ingress: &Ingress, + open: &ProducerOpen, + sequence: u64, + ids: &[&str], +) -> Result<(Bytes, AcceptedAppend)> { + let producer = ingress.open_producer(open.clone()).await?; + let session = session(open.bucket); + let body = ipc_file(&[batch(ids, &vec![session.as_str(); ids.len()])?])?; + let prepared = ingress + .prepare( + ingress.reserve_upload(body.len()).await?, + body.clone(), + &producer, + sequence, + ) + .await?; + let accepted = ingress.append(&prepared).await?; + Ok((body, accepted)) +} +async fn verify( + publisher: &Publisher, + expected: &BTreeMap, + version: u64, +) -> Result<()> { + let dataset = publisher.read_dataset(None).await?; + assert_eq!(dataset.version().version, version); + let batches = dataset + .scan() + .try_into_stream() + .await? + .try_collect::>() + .await?; + let mut found = BTreeMap::new(); + for batch in batches { + let ids = batch + .column_by_name("id") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + for row in 0..batch.num_rows() { + let key = ids.value(row).to_owned(); + assert!( + found.insert(key.clone(), ()).is_none(), + "duplicate published id" + ); + let actual = batch.slice(row, 1); + let expected = &expected[&key]; + assert_eq!(actual.schema(), expected.schema()); + for (left, right) in actual.columns().iter().zip(expected.columns()) { + assert_eq!(left, right, "column mismatch for {key}"); + } + } + } + assert_eq!( + found.keys().collect::>(), + expected.keys().collect::>() + ); + Ok(()) +} +fn remember(expected: &mut BTreeMap, body: Bytes) -> Result<()> { + for batch in ipc::decode(body, Arc::new(rollout_schema()))? { + let ids = batch + .column_by_name("id") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + for row in 0..batch.num_rows() { + assert!(expected + .insert(ids.value(row).to_owned(), batch.slice(row, 1)) + .is_none()); + } + } + Ok(()) +} +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn real_rollout_files_resume_after_preparation_and_publication_restart_then_retire( +) -> Result<()> { + let root = tempfile::tempdir()?; + let path = root.path().join("rollouts.lance"); + let c = network::Cluster::open_signed(root.path(), layout()).await?; + bootstrap(&c, &path, &binding()).await?; + let base = Dataset::open(path.to_str().unwrap()).await?; + let object_base = base + .manifest_location() + .path + .as_ref() + .rsplit_once("/_versions/") + .context("manifest base")? + .0 + .to_owned(); + let storage = Storage { + uri: path.to_str().unwrap().into(), + base: ObjectPath::parse(object_base)?, + options: HashMap::new(), + session: Arc::new(lance::session::Session::default()), + }; + drop(base); + let budget = MergeMemoryBudget::new(32 << 20); + let ingress = connect(&c, binding()).await?; + let mut expected = BTreeMap::new(); + let mut originals = Vec::new(); + for bucket in 0..4 { + let open = ProducerOpen { + bucket, + slot: bucket, + expected_generation: 0, + operation: Uuid::new_v4(), + }; + let ids = [format!("{bucket}-a"), format!("{bucket}-b")]; + let (body, receipt) = append(&ingress, &open, 1, &[&ids[0], &ids[1]]).await?; + remember(&mut expected, body.clone())?; + originals.push((open, body, receipt)); + } + let mut first = publisher(&ingress, &storage, &budget, 1).await?; + let saved = first + .prepare_next() + .await? + .context("no staged publication")?; + assert_eq!( + first + .state() + .await? + .progress + .pending + .as_ref() + .unwrap() + .plan + .sources + .len(), + 2 + ); + assert_eq!(budget.reserved(), 0); + let old_authority = first.authority().clone(); + drop((first, ingress)); + c.shutdown().await?; + + let c = network::Cluster::open_signed(root.path(), layout()).await?; + let ingress = connect(&c, binding()).await?; + let mut second = publisher(&ingress, &storage, &budget, 2).await?; + assert_eq!( + second.prepare_next().await?.unwrap(), + saved, + "recovery rewrote durable prepared files" + ); + let old = c + .client + .request_routed(Request::Commit(Box::new(Command::Prepare( + p::PrepareRequest { + authority: old_authority, + transaction: saved.transaction, + plan_index: saved.plan_index, + descriptors: saved.descriptors.clone(), + }, + )))) + .await?; + assert!(matches!( + old.reply, + Reply::Commit(Outcome::Committed(Applied::Publisher(p::Reply::Rejected( + p::Rejection::Authority + )))) + )); + let published = second + .commit_pending() + .await? + .context("publication missing")?; + verify(&second, &expected, 2).await?; + for (_, _, receipt) in &originals { + assert_eq!( + ingress.publication_status(receipt).await?, + PublicationStatus::Pending + ); + } + // Crash boundary: the real manifest is selected and confirmed, but source + // closure/cursor advancement has not yet run. New authority must resume it. + drop((second, ingress)); + c.shutdown().await?; + + let c = network::Cluster::open_signed(root.path(), layout()).await?; + let ingress = connect(&c, binding()).await?; + let mut third = publisher(&ingress, &storage, &budget, 3).await?; + assert_eq!(third.commit_pending().await?.unwrap(), published); + let advanced = third + .finish_pending() + .await? + .context("advancement missing")?; + assert_eq!(advanced.sources.len(), 2); + for (open, body, receipt) in &originals { + assert_eq!( + ingress.publication_status(receipt).await?, + PublicationStatus::Published { + minimum_read_version: 2 + } + ); + let producer = ingress.open_producer(open.clone()).await?; + let exact = ingress + .prepare( + ingress.reserve_upload(body.len()).await?, + body.clone(), + &producer, + 1, + ) + .await?; + assert_eq!(ingress.append(&exact).await?.position(), receipt.position()); + for node in 1..=3 { + tokio::time::timeout(Duration::from_secs(10), async { + loop { + let response = c + .client + .request( + node, + Request::Fetch { + group: receipt.position().group, + index: receipt.position().index, + }, + ) + .await?; + if matches!(response.reply, Reply::Record(None)) { + break; + } + tokio::time::sleep(Duration::from_millis(10)).await; + } + Ok::<_, anyhow::Error>(()) + }) + .await??; + } + } + assert!(third.run_once().await?.is_none()); + verify(&third, &expected, 2).await?; + for sequence in 2..=5 { + for (open, _, _) in &originals { + let id = format!("{}-round-{sequence}", open.bucket); + let (body, _) = append(&ingress, open, sequence, &[&id]).await?; + remember(&mut expected, body)?; + } + assert!(third.run_once().await?.is_some()); + verify(&third, &expected, sequence + 1).await?; + assert!( + third.state().await?.manifests.unwrap().versions.len() <= 3, + "catalog history grew beyond policy" + ); + } + assert_eq!(budget.reserved(), 0); + drop((third, ingress)); + c.shutdown().await?; + + let c = network::Cluster::open_signed(root.path(), layout()).await?; + let ingress = connect(&c, binding()).await?; + let mut final_owner = publisher(&ingress, &storage, &budget, 4).await?; + assert!(final_owner.run_once().await?.is_none()); + verify(&final_owner, &expected, 6).await?; + assert!( + final_owner.read_dataset(Some(2)).await.is_err(), + "retired history unexpectedly resolved" + ); + drop(final_owner); + c.shutdown().await +} diff --git a/crates/lance-context-ingest/src/tests/network.rs b/crates/lance-context-ingest/src/tests/network.rs index 083f76f6..246d60a4 100644 --- a/crates/lance-context-ingest/src/tests/network.rs +++ b/crates/lance-context-ingest/src/tests/network.rs @@ -73,11 +73,59 @@ pub(super) struct Cluster { pub client: Arc, nodes: Vec<(Runtime, peer::Transport, client::Server)>, } +type Signers = ( + Arc, + BTreeMap>, +); impl Cluster { pub async fn open(root: &Path, layout: Arc) -> Result { + Self::open_inner(root, layout, None).await + } + #[cfg(feature = "materializer")] + pub async fn open_signed(root: &Path, layout: Arc) -> Result { + use lance_context_wal::replication::source::{CatalogSigner, SourceSigner}; + fn key(path: &Path) -> Result> { + if !path.exists() { + let key = ring::signature::Ed25519KeyPair::generate_pkcs8( + &ring::rand::SystemRandom::new(), + ) + .map_err(|_| anyhow::anyhow!("test signing key generation"))?; + std::fs::write(path, key.as_ref())?; + } + Ok(std::fs::read(path)?) + } + let catalog = Arc::new(CatalogSigner::from_pkcs8(&key(&root.join("catalog.pk8"))?)?); + let sources = layout + .groups + .keys() + .map(|&group| { + Ok(( + group, + Arc::new(SourceSigner::from_pkcs8(&key( + &root.join(format!("source-{group}.pk8")) + )?)?), + )) + }) + .collect::>>()?; + Self::open_inner(root, layout, Some((catalog, sources))).await + } + async fn open_inner( + root: &Path, + layout: Arc, + signers: Option, + ) -> Result { let tls = certificates()?; let limits = Limits { producer_slots: 16, + catalog_key: signers.as_ref().map(|s| s.0.public_key()), + source_keys: signers + .as_ref() + .map(|s| { + s.1.iter() + .map(|(&group, key)| (group, key.public_key())) + .collect() + }) + .unwrap_or_default(), ..Limits::default() }; let mut peers = BTreeMap::new(); @@ -108,7 +156,7 @@ impl Cluster { } let mut nodes = Vec::new(); for (id, listener, data, control) in listeners { - let node = Node::open( + let mut node = Node::open( &root.join(id.to_string()), layout.clone(), id, @@ -119,6 +167,12 @@ impl Cluster { data_bytes: 128 << 20, headroom_bytes: 256 << 20, })?; + if let Some((catalog, sources)) = &signers { + node.set_catalog_signer(catalog.clone())?; + for (&group, key) in sources { + node.set_source_signer(group, key.clone())?; + } + } let (runtime, outbox) = Runtime::spawn(node, runtime::Options::default())?; let mut server = tls[&id].server.clone(); server.alpn_protocols = vec![peer::ALPN.to_vec()]; diff --git a/crates/lance-context-wal/src/replication/runtime/client.rs b/crates/lance-context-wal/src/replication/runtime/client.rs index 59f5105c..f8e6de2f 100644 --- a/crates/lance-context-wal/src/replication/runtime/client.rs +++ b/crates/lance-context-wal/src/replication/runtime/client.rs @@ -457,6 +457,8 @@ pub struct Client { layout: [u8; 32], policy: [u8; 32], max_record: usize, + max_plan_source: u64, + max_descriptors: u64, producer_slots: u32, next: AtomicU64, voters: BTreeMap>, @@ -470,6 +472,12 @@ impl Client { pub fn max_record_bytes(&self) -> usize { self.max_record } + pub fn max_plan_source_bytes(&self) -> u64 { + self.max_plan_source + } + pub fn max_descriptor_bytes(&self) -> u64 { + self.max_descriptors + } pub fn producer_slots(&self) -> u32 { self.producer_slots } @@ -511,6 +519,8 @@ impl Client { layout: layout.fingerprint()?, policy: Sha256::digest(limits.identity()?).into(), max_record: limits.max_record_bytes as usize, + max_plan_source: limits.max_plan_source_bytes, + max_descriptors: limits.max_descriptor_bytes, producer_slots: limits.producer_slots, connector: TlsConnector::from(tls), budgets: Budgets::new(&options), diff --git a/docs/design/wal-rollout-materializer.md b/docs/design/wal-rollout-materializer.md new file mode 100644 index 00000000..3b863db8 --- /dev/null +++ b/docs/design/wal-rollout-materializer.md @@ -0,0 +1,83 @@ +# Rollout materialization through the public WAL transport + +The optional `lance-context-ingest/materializer` feature supplies a publisher for +the fixed rollout schema. It reads source records over the same authenticated, +bounded WAL client used by ingestion, writes real Lance files, publishes through +the managed manifest catalog, and retires source payloads after durable source +closure. No Lance OSS changes are required. + +The application supplies an authorized `Ingress`, resolved storage URI/object +path/options, a shared Lance session, a shared materialization memory budget and +a persisted publisher-acquisition identity. The storage namespace and dataset +incarnation must already be bound and bootstrapped. Catalog epochs fence stale +owners; a scheduling lease alone does not authorize a publication. + +## Publication and recovery + +`prepare_next` selects a bounded prefix across source groups, records one fixed +transaction/read version, obtains catalog-issued source permits, and commits +each source pin before fetching payloads. It matches each fetched record against +the pin, verifies its digest off the async executor, and validates the actual +rollout IPC/schema/session bucket. Existing input buffers remain shared through +Arrow decoding; alignment repair and Lance encoding can allocate. + +The publisher limits source bytes, record count, fetch concurrency, file writers +and per-group scan pages. A bounded number of writers stage immutable files; +their descriptors are then recorded in the replicated catalog. An uncertain +plan/prepare RPC retains the same in-memory attempt. After a new authority takes +over, a durable prepared result is reused instead of writing replacement files. +Files produced before a prepare is committed can become orphans. + +`commit_pending` uses the recorded fragments and transaction UUID with the +managed commit handler. The target remains exactly read-version plus one; +automatic rebase and cleanup are disabled. A selected manifest is resolved by a +quorum catalog read and checked against the transaction, plan, preparation and +descriptor digest before confirmation. An uncertain commit never authorizes +source deletion or moving the transaction to a new version. + +`finish_pending` obtains closure permits from the confirmed catalog, commits each +source closure and obtains its quorum proof. The catalog verifies the proofs +before advancing its independent source positions. Ingress receipts become +published at that advancement. A crash after manifest selection or after some +source closures resumes the same pending publication; exact last-batch retries +retain their original ACK positions even after the payload is retired. + +`run_once` combines these phases. An error is reconciled using the durable state +on the next call. A return of `None` means there is currently no pending work; +it is not proof that an earlier uncertain operation was rejected. The application +still supplies bounded scheduling, backoff, shutdown and table discovery. + +## Resource and retention boundaries + +One materialization reserves space for source buffers, decoding and row metadata +before payload fetch. Admitted file-writing work retains the reservation if its +caller is cancelled, and joins all its file writers even when one fails. These +are logical application bounds, not an allocator or encoder/cgroup hard limit. +The application must also bound its Lance encoder and metadata caches. + +Pages containing only other tables can advance an in-memory scan hint; they do +not advance the durable table checkpoint. A selected record is never skipped by +that hint. Each scan is bounded, and publication rotates through source groups. +This is not a global uniqueness index: globally unique immutable row IDs remain +the caller's explicit contract, as on the ingress adapter. + +Catalog manifest history is retired to a bounded suffix while retaining the +bootstrap record and active dependencies. Retirement does not delete any blob +object. Old-version reads can fail once catalog history is retired. Physical +WAL log/image collection belongs to the WAL service; coordinated blob orphan +collection and legacy-table migration remain separate integration work. + +## Qualification scope + +The materializer regression uses real TLS-connected replicas and a real Lance +dataset with the complete rollout schema. It restarts the original replica +directories after file preparation and again after manifest publication, takes +over publisher authority, verifies complete row values, retires the source +payloads and retries the original appends. Additional publications exercise +bounded manifest history, followed by another full-directory restart. + +The regression and feature Clippy checks pass, including the existing ingress +and TLS catalog-reader regressions. This is not a sustained throughput result, +a Kubernetes deployment or an HTTP ingestion integration. + +Tracking: #378.