feat(ingest): coordinate bounded materializer workers with etcd leases

.github/workflows/rust-test.yml

@@ -171,6 +171,8 @@ jobs:
           cargo clippy -p lance-context-ingest --features http --all-targets -- -D warnings
           cargo clippy -p lance-context-ingest --features http-client,materializer --all-targets -- -D warnings
           cargo test -p lance-context-ingest --features http-client,materializer --lib
+      - name: Materializer worker feature checks
+        run: cargo clippy -p lance-context-ingest --features worker --all-targets -- -D warnings
       - name: Start etcd for HA scheduler test
         run: |
           ETCD_VERSION=3.7.0
@@ -198,6 +200,7 @@ jobs:
         run: |
           cargo test -p lance-context-master -p lance-context-server -p lance-context-merge -- --ignored
           cargo test -p lance-context-core registry --lib -- --ignored
+          cargo test -p lance-context-ingest --features worker --lib tests::worker -- --ignored
 
   coverage:
     runs-on: ubuntu-24.04
@@ -263,6 +266,8 @@ jobs:
             -p lance-context-master -p lance-context-server -p lance-context-merge -- --ignored
           cargo llvm-cov --no-report --all-features \
             -p lance-context-core --lib -- registry --ignored
+          cargo llvm-cov --no-report --all-features \
+            -p lance-context-ingest --lib -- tests::worker --ignored
       - name: Merge coverage reports
         run: cargo llvm-cov report --lcov --output-path lcov.info
       - name: Save coverage report

crates/lance-context-ingest/Cargo.toml

@@ -7,11 +7,13 @@ publish = false
 description = "Durable rollout ingress over the independent lance-context WAL"
 
 [features]
+worker = ["materializer", "dep:etcd-client"]
 http = ["dep:axum", "dep:http-body-util"]
 http-client = ["http", "dep:reqwest", "lance-context-wal/client-journal"]
 materializer = ["lance-context-wal/lance-integration", "dep:lance", "dep:lance-table", "dep:object_store", "dep:futures"]
 
 [dependencies]
+etcd-client = { version = "0.19", features = ["tls"], optional = true }
 axum = { version = "0.8", optional = true, default-features = false, features = ["json", "http1", "tokio"] }
 http-body-util = { version = "0.1", optional = true }
 reqwest = { version = "0.12", optional = true, default-features = false, features = ["json", "rustls-tls"] }

crates/lance-context-ingest/src/lib.rs

@@ -34,6 +34,8 @@ pub mod http;
 #[cfg(feature = "materializer")]
 pub mod materializer;
 pub mod receipt;
+#[cfg(feature = "worker")]
+pub mod worker;
 
 /// Resolved by the application's authorized catalog, not an uploaded URI/UUID.
 #[derive(Clone, Debug)]

crates/lance-context-ingest/src/materializer.rs

@@ -54,7 +54,7 @@ impl Options {
             })
             .context("materializer workspace overflow")
     }
-    fn validate(&self, ingress: &Ingress, budget: &MergeMemoryBudget) -> Result<()> {
+    pub(super) 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(),
@@ -183,31 +183,9 @@ impl Publisher {
         );
         Ok(view)
     }
-    fn reader(&self) -> Result<p::lance::ManagedManifest> {
-        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<u64>) -> Result<Dataset> {
-        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)
+        read_dataset(&self.ingress, &self.storage, version).await
     }
 
     /// Select a complete bounded prefix, pin every source, write immutable Lance
@@ -627,7 +605,41 @@ impl Publisher {
     }
 }
 
-async fn query(ingress: &Ingress, group: u64, value: query::Query) -> Result<query::Value> {
+/// Open a catalog-selected dataset without acquiring publisher ownership.
+/// The application resolves and authorizes both the binding and storage first.
+pub async fn read_dataset(
+    ingress: &Ingress,
+    storage: &Storage,
+    version: Option<u64>,
+) -> Result<Dataset> {
+    let reader = p::lance::ManagedManifest::reader(
+        Arc::new(ingress.client.manifest_catalog()),
+        storage.base.clone(),
+        &ingress.layout,
+        ingress.scope.table,
+        ingress.scope.incarnation,
+    )?;
+    let mut builder = DatasetBuilder::from_uri(&storage.uri)
+        .with_storage_options(storage.options.clone())
+        .with_session(storage.session.clone())
+        .with_commit_handler(Arc::new(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)
+}
+
+pub(super) async fn query(
+    ingress: &Ingress,
+    group: u64,
+    value: query::Query,
+) -> Result<query::Value> {
     let response = ingress
         .client
         .request_routed(Request::Query {

crates/lance-context-ingest/src/tests.rs

@@ -15,6 +15,8 @@ mod http;
 #[cfg(feature = "materializer")]
 mod materializer;
 mod network;
+#[cfg(feature = "worker")]
+mod worker;
 
 fn layout() -> Arc<Layout> {
     Arc::new(Layout {

crates/lance-context-ingest/src/tests/worker.rs

@@ -0,0 +1,233 @@
+use super::*;
+use crate::{
+    materializer::{self as publication, Storage},
+    worker::{Assignment, Claims, Options, Worker},
+};
+use futures::TryStreamExt;
+use std::{collections::HashMap, path::Path};
+
+async fn fixture() -> Result<(etcd_client::Client, Claims, String)> {
+    let endpoints = std::env::var("ETCD_TEST_ENDPOINTS").context("ETCD_TEST_ENDPOINTS required")?;
+    let client =
+        etcd_client::Client::connect(endpoints.split(',').collect::<Vec<_>>(), None).await?;
+    let prefix = format!("/lance-context-test/ingest-worker/{}", Uuid::new_v4());
+    let claims = Claims::new(client.clone(), prefix.clone(), 10, Duration::from_secs(3))?;
+    Ok((client, claims, prefix))
+}
+async fn storage(path: &Path) -> Result<Storage> {
+    let dataset = Dataset::open(path.to_str().context("test URI")?).await?;
+    let base = dataset
+        .manifest_location()
+        .path
+        .as_ref()
+        .rsplit_once("/_versions/")
+        .context("manifest base")?
+        .0
+        .to_owned();
+    Ok(Storage {
+        uri: path.to_str().unwrap().into(),
+        base: object_store::path::Path::parse(base)?,
+        options: HashMap::new(),
+        session: Arc::new(lance::session::Session::default()),
+    })
+}
+fn options() -> Options {
+    Options {
+        parallel_tables: 1,
+        publications_per_claim: 4,
+        poll_interval: Duration::from_millis(50),
+        publication: publication::Options {
+            max_source_bytes: 1 << 20,
+            max_records: 8,
+            retained_manifest_versions: 2,
+            ..Default::default()
+        },
+    }
+}
+async fn append(
+    ingress: &Ingress,
+    producer: &Producer,
+    sequence: u64,
+    id: &str,
+) -> Result<(Bytes, AcceptedAppend)> {
+    let session = session(0);
+    let body = ipc_file(&[batch(&[id], &[&session])?])?;
+    let prepared = ingress
+        .prepare(
+            ingress.reserve_upload(body.len()).await?,
+            body.clone(),
+            producer,
+            sequence,
+        )
+        .await?;
+    Ok((body, ingress.append(&prepared).await?))
+}
+#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
+#[ignore = "requires ETCD_TEST_ENDPOINTS; executed by the etcd CI suite"]
+async fn expired_claim_cannot_steal_successor_epoch_and_workers_resume_prepared_lance() -> Result<()>
+{
+    let (mut etcd, claims, prefix) = fixture().await?;
+    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 ingress = connect(&c, binding()).await?;
+    let storage = storage(&path).await?;
+    let assignment = Assignment {
+        ingress: ingress.clone(),
+        storage: storage.clone(),
+    };
+    let budget = MergeMemoryBudget::new(32 << 20);
+    let producer = ingress
+        .open_producer(ProducerOpen {
+            bucket: 0,
+            slot: 0,
+            expected_generation: 0,
+            operation: Uuid::new_v4(),
+        })
+        .await?;
+    let (first, receipt) = append(&ingress, &producer, 1, "before-takeover").await?;
+    let mut expected = BTreeMap::new();
+    materializer::remember(&mut expected, first)?;
+    let old = claims
+        .take(&ingress)
+        .await?
+        .context("first claim unavailable")?;
+    assert!(
+        claims.take(&ingress).await?.is_none(),
+        "overlapping table claim granted"
+    );
+    let mut old_publisher = old
+        .acquire(assignment.clone(), options().publication, budget.clone())
+        .await?
+        .context("first claim invalid")?;
+    let prepared = old_publisher
+        .prepare_next()
+        .await?
+        .context("prepared result missing")?;
+    let old_epoch = old_publisher.authority().epoch;
+    let key = format!("{prefix}/{}/{}", layout().cluster, binding().table);
+    let lease = etcd.get(key, None).await?.kvs()[0].lease();
+    etcd.lease_revoke(lease).await?;
+    let successor = claims
+        .take(&ingress)
+        .await?
+        .context("successor claim unavailable")?;
+    let mut publisher = successor
+        .acquire(assignment.clone(), options().publication, budget.clone())
+        .await?
+        .context("successor claim invalid")?;
+    assert_eq!(publisher.authority().epoch, old_epoch + 1);
+    // This deliberately calls the OLD claim after the successor has committed
+    // its newer epoch. It must fail the lease check and leave that epoch intact.
+    assert!(old
+        .acquire(assignment.clone(), options().publication, budget.clone())
+        .await?
+        .is_none());
+    assert_eq!(publisher.state().await?.authority, *publisher.authority());
+    assert!(old_publisher.run_once().await.is_err());
+    assert_eq!(publisher.prepare_next().await?.unwrap(), prepared);
+    let _ = old.release().await;
+    assert!(
+        successor.valid().await?,
+        "old lease cleanup removed successor"
+    );
+    assert!(publisher.run_once().await?.is_some());
+    assert_eq!(
+        ingress
+            .wait_published(&receipt, Duration::from_secs(5))
+            .await?,
+        2
+    );
+    materializer::verify(&publisher, &expected, 2).await?;
+    successor.release().await?;
+    drop((old_publisher, publisher));
+
+    let (second, receipt) = append(&ingress, &producer, 2, "automatic-worker").await?;
+    materializer::remember(&mut expected, second)?;
+    let workers = [
+        Worker::start(
+            claims.clone(),
+            vec![assignment.clone()],
+            options(),
+            budget.clone(),
+        )?,
+        Worker::start(claims.clone(), vec![assignment], options(), budget.clone())?,
+    ];
+    assert_eq!(
+        ingress
+            .wait_published(&receipt, Duration::from_secs(15))
+            .await?,
+        3
+    );
+    assert!(workers.iter().all(|w| w.observation().active <= 1));
+    for worker in workers {
+        worker.shutdown().await?;
+    }
+    assert_eq!(budget.reserved(), 0);
+    // Read independently of either stopped publisher; no acquisition for reads.
+    let dataset = publication::read_dataset(&ingress, &storage, None).await?;
+    assert_eq!(dataset.version().version, 3);
+    let mut found = BTreeMap::new();
+    for batch in dataset
+        .scan()
+        .try_into_stream()
+        .await?
+        .try_collect::<Vec<_>>()
+        .await?
+    {
+        let ids = batch
+            .column_by_name("id")
+            .unwrap()
+            .as_any()
+            .downcast_ref::<StringArray>()
+            .unwrap();
+        for row in 0..batch.num_rows() {
+            let id = ids.value(row);
+            assert!(found.insert(id.to_owned(), ()).is_none());
+            assert_eq!(batch.slice(row, 1), expected[id]);
+        }
+    }
+    assert_eq!(found.len(), expected.len());
+    etcd.delete(
+        prefix,
+        Some(etcd_client::DeleteOptions::new().with_prefix()),
+    )
+    .await?;
+    drop((dataset, ingress));
+    c.shutdown().await
+}
+
+#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
+#[ignore = "requires ETCD_TEST_ENDPOINTS; executed by the etcd CI suite"]
+async fn abandoned_claim_expires_without_a_detached_delete() -> Result<()> {
+    let (mut etcd, claims, prefix) = fixture().await?;
+    let root = tempfile::tempdir()?;
+    let c = network::Cluster::open_signed(root.path(), layout()).await?;
+    bootstrap(&c, &root.path().join("rollouts.lance"), &binding()).await?;
+    let ingress = connect(&c, binding()).await?;
+    let first = claims.take(&ingress).await?.context("claim unavailable")?;
+    drop(first);
+    assert!(
+        claims.take(&ingress).await?.is_none(),
+        "drop deleted the claim instead of letting its lease expire"
+    );
+    let successor = tokio::time::timeout(Duration::from_secs(20), async {
+        loop {
+            if let Some(claim) = claims.take(&ingress).await? {
+                return Ok::<_, anyhow::Error>(claim);
+            }
+            tokio::time::sleep(Duration::from_millis(250)).await;
+        }
+    })
+    .await??;
+    assert!(successor.valid().await?);
+    successor.release().await?;
+    etcd.delete(
+        prefix,
+        Some(etcd_client::DeleteOptions::new().with_prefix()),
+    )
+    .await?;
+    drop(ingress);
+    c.shutdown().await
+}