From 6b12a91c084976fab3311b06ef57b59f4f906322 Mon Sep 17 00:00:00 2001 From: Beinan Date: Mon, 19 Oct 2026 03:38:25 +0000 Subject: [PATCH] feat: add shared WAL proposal admission budget --- .github/workflows/rust-test.yml | 4 +- .github/workflows/style.yml | 2 +- Cargo.lock | 1 + crates/lance-context-wal/Cargo.toml | 4 + crates/lance-context-wal/README.md | 17 ++ crates/lance-context-wal/src/admission.rs | 251 ++++++++++++++++++++++ crates/lance-context-wal/src/lib.rs | 3 + 7 files changed, 279 insertions(+), 3 deletions(-) create mode 100644 crates/lance-context-wal/src/admission.rs diff --git a/.github/workflows/rust-test.yml b/.github/workflows/rust-test.yml index d5a85bee..c9e10233 100644 --- a/.github/workflows/rust-test.yml +++ b/.github/workflows/rust-test.yml @@ -63,8 +63,8 @@ jobs: run: | cargo run -p lance-context-wal --features lab --bin context-wal-lab -- codec --payload-bytes 1024 --rows 8 --iterations 2 cargo run -p lance-context-wal --features lab --bin context-wal-lab -- disk --directory "$RUNNER_TEMP/wal-disk-smoke" --payload-bytes 1024 --batch 4 --writers 2 --groups-per-writer 2 --rounds 8 - - name: WAL placement compatibility tests - run: cargo test -p lance-context-wal --features placement --lib + - name: WAL placement and shared admission tests + run: cargo test -p lance-context-wal --features placement,admission --lib - name: Start etcd for HA scheduler test run: | ETCD_VERSION=3.7.0 diff --git a/.github/workflows/style.yml b/.github/workflows/style.yml index 80772f6e..43c97b18 100644 --- a/.github/workflows/style.yml +++ b/.github/workflows/style.yml @@ -56,7 +56,7 @@ jobs: - name: Clippy (lance-context, default features) run: cargo clippy --manifest-path crates/lance-context/Cargo.toml --all-targets -- -D warnings - name: Clippy (WAL experiments) - run: cargo clippy -p lance-context-wal --features lab,placement --all-targets -- -D warnings + run: cargo clippy -p lance-context-wal --features lab,placement,admission --all-targets -- -D warnings typos: name: Spell Check diff --git a/Cargo.lock b/Cargo.lock index af53d6b3..62dba986 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4184,6 +4184,7 @@ dependencies = [ "serde", "serde_json", "sha2 0.10.9", + "tempfile", "uuid", ] diff --git a/crates/lance-context-wal/Cargo.toml b/crates/lance-context-wal/Cargo.toml index 00abd309..58b5958f 100644 --- a/crates/lance-context-wal/Cargo.toml +++ b/crates/lance-context-wal/Cargo.toml @@ -8,6 +8,7 @@ description = "Independent Rust WAL experiments owned by lance-context" [features] placement = ["dep:anyhow", "dep:serde", "dep:serde_json", "dep:sha2", "dep:uuid"] +admission = ["dep:anyhow", "dep:serde"] lab = ["dep:anyhow", "dep:arrow-array", "dep:arrow-buffer", "dep:arrow-ipc", "dep:arrow-schema", "dep:bytes", "dep:clap", "dep:hdrhistogram", "dep:raft-engine", "dep:raft-proto", "dep:serde_json"] [dependencies] @@ -26,6 +27,9 @@ serde = { version = "1", features = ["derive"], optional = true } sha2 = { version = "0.10", optional = true } uuid = { version = "1", features = ["serde"], optional = true } +[dev-dependencies] +tempfile = "3" + [[bin]] name = "context-wal-lab" path = "src/bin/lab.rs" diff --git a/crates/lance-context-wal/README.md b/crates/lance-context-wal/README.md index b9adde50..83ffd276 100644 --- a/crates/lance-context-wal/README.md +++ b/crates/lance-context-wal/README.md @@ -19,6 +19,23 @@ Compatibility tests preserve a layout fingerprint recorded by a real three-node experiment and a fixed session-hash vector. Other tests reject ambiguous ownership and ensure routes survive JSON serialization and restart. +The independent `admission` feature provides one proposal budget shared by all +groups on a node. Data proposals leave 25% of the configured budget (up to +256 MiB) for control traffic. Reservations release on drop; encoded engine-write +counters account for peer, control and rewrite traffic between directory scans. +Refresh measures file lengths and allocated blocks, rather than estimating +segment sizes. Successful admission does not walk the directory. + +This is proposal admission, not filesystem space reservation or a limit on +already accepted replication traffic. Callers must supply a monotonic write +counter covering every group and protect file lifecycle changes with the shared +lock. Snapshot collection and replica-pressure coordination are still required +for sustained operation. No production writer is connected to this module yet. + +```bash +cargo test -p lance-context-wal --features admission --lib +``` + ```bash cargo build --release -p lance-context-wal --features lab --bin context-wal-lab target/release/context-wal-lab codec --payload-bytes 262144 --rows 8 --iterations 30 diff --git a/crates/lance-context-wal/src/admission.rs b/crates/lance-context-wal/src/admission.rs new file mode 100644 index 00000000..62e7597a --- /dev/null +++ b/crates/lance-context-wal/src/admission.rs @@ -0,0 +1,251 @@ +//! Node-wide proposal admission, shared by every physical group on one engine. +//! This reserves a CONTROL WINDOW, not filesystem extents or a bound on future +//! catch-up traffic. Replication must still persist already accepted entries; +//! snapshot/GC and replica-pressure coordination are required for sustained use. +use std::path::{Path, PathBuf}; +use std::sync::{ + atomic::{AtomicU64, Ordering}, + Arc, Mutex, RwLock, TryLockError, +}; + +use anyhow::{ensure, Result}; +use serde::{Deserialize, Serialize}; + +#[derive(Clone, Debug, Serialize, Deserialize)] +pub struct Snapshot { + pub measured_disk_bytes: u64, + pub accounted_disk_bytes: u64, + pub reserved_proposal_bytes: u64, + pub data_limit: u64, + pub max_log_bytes: u64, +} +#[derive(Debug)] +struct Accounting { + measured: u64, + counter_at_measurement: u64, + pending: u64, +} +#[derive(Debug)] +pub struct Budget { + directory: PathBuf, + written: Arc, + data_limit: u64, + max: u64, + accounting: Mutex, + filesystem: Arc>, +} +impl Budget { + /// `written` must count every encoded engine append and rewrite, across all + /// groups and peer/control traffic. It must never reset during this budget's + /// lifetime. Files are direct children of `directory`; no recursive scan. + /// Hold `filesystem` for writing during file deletion/purge/rename. Appends + /// may overlap refresh, provided the counter includes their encoded bytes. + pub fn new( + directory: &Path, + written: Arc, + max: usize, + filesystem: Arc>, + ) -> Result> { + ensure!(max >= 4 * 1024 * 1024, "log budget must be at least 4 MiB"); + let counter = written.load(Ordering::Acquire); + let measured = retained_disk_bytes(directory)?; + // Reserve 25%, up to 256 MiB, from new producer proposals. It remains + // available for publisher/control records and previously accepted Raft + // traffic. This is explicitly not a guarantee against a full filesystem. + let reserve = (max / 4).min(256 * 1024 * 1024); + Ok(Arc::new(Self { + directory: directory.to_owned(), + written, + filesystem, + max: max as u64, + data_limit: (max - reserve) as u64, + accounting: Mutex::new(Accounting { + measured, + counter_at_measurement: counter, + pending: 0, + }), + })) + } + fn accounted(&self, state: &Accounting) -> u64 { + state.measured.saturating_add( + self.written + .load(Ordering::Acquire) + .saturating_sub(state.counter_at_measurement), + ) + } + /// No directory walk on successful admission. The engine's exact encoded + /// write counts include peer traffic and control/HardState writes across all + /// groups; pending proposals are charged atomically on top of that baseline. + pub fn reserve(self: &Arc, bytes: usize, data: bool) -> Result { + let bytes = u64::try_from(bytes)?; + let mut state = self.accounting.lock().unwrap(); + let limit = if data { self.data_limit } else { self.max }; + ensure!( + self.accounted(&state) + .saturating_add(state.pending) + .saturating_add(bytes) + <= limit, + "node proposal disk budget exhausted; data admission preserves control window" + ); + state.pending += bytes; + Ok(Reservation { + budget: self.clone(), + bytes, + }) + } + /// Refresh after persistence and on explicit status reads. Sampling the + /// counter BEFORE reading directory metadata makes overlapping appends count + /// twice at worst; it cannot subtract a write the measurement did not see. + /// Engine append writes must not acquire the accounting mutex: independent + /// group writes can still coalesce in the engine. + pub fn refresh(&self) -> Result<()> { + let _files = match self.filesystem.try_read() { + Ok(guard) => guard, + Err(TryLockError::WouldBlock) => return Ok(()), + Err(TryLockError::Poisoned(_)) => anyhow::bail!("maintenance filesystem lock poisoned"), + }; + self.refresh_exclusive() + } + /// Caller owns a directory lifecycle read/write guard. No unlink/purge can + /// overlap this measurement or its accounting update. Engine appends can. + fn refresh_exclusive(&self) -> Result<()> { + let counter = self.written.load(Ordering::Acquire); + let measured = retained_disk_bytes(&self.directory)?; + let mut state = self.accounting.lock().unwrap(); + if counter >= state.counter_at_measurement { + state.measured = measured; + state.counter_at_measurement = counter; + } + Ok(()) + } + pub fn snapshot(&self) -> Snapshot { + let state = self.accounting.lock().unwrap(); + Snapshot { + measured_disk_bytes: state.measured, + accounted_disk_bytes: self.accounted(&state), + reserved_proposal_bytes: state.pending, + data_limit: self.data_limit, + max_log_bytes: self.max, + } + } +} +#[derive(Debug)] +pub struct Reservation { + budget: Arc, + bytes: u64, +} +impl Drop for Reservation { + fn drop(&mut self) { + self.budget.accounting.lock().unwrap().pending -= self.bytes; + } +} + +fn retained_disk_bytes(directory: &Path) -> Result { + let mut bytes = 0u64; + for entry in std::fs::read_dir(directory)? { + let entry = entry?; + ensure!( + !entry.file_type()?.is_symlink(), + "WAL accounting directory contains a symlink" + ); + let metadata = entry.metadata()?; + if !metadata.is_file() { + continue; + } + let mut occupied = metadata.len(); + #[cfg(unix)] + { + use std::os::unix::fs::MetadataExt; + occupied = occupied.max(metadata.blocks().saturating_mul(512)); + } + bytes = bytes + .checked_add(occupied) + .ok_or_else(|| anyhow::anyhow!("disk usage overflow"))?; + } + Ok(bytes) +} + +#[cfg(test)] +mod tests { + use super::*; + #[test] + fn racing_groups_share_admission_and_preserve_control_window() -> Result<()> { + let directory = tempfile::tempdir()?; + let written = Arc::new(AtomicU64::new(0)); + let budget = Budget::new( + directory.path(), + written.clone(), + 8 * 1024 * 1024, + Arc::new(RwLock::new(())), + )?; + budget.refresh()?; + let gate = std::sync::Barrier::new(8); + let reservations = std::thread::scope(|scope| { + let threads: Vec<_> = (0..8) + .map(|_| { + let (budget, gate) = (budget.clone(), &gate); + scope.spawn(move || { + gate.wait(); + budget.reserve(1024 * 1024, true).ok() + }) + }) + .collect(); + threads + .into_iter() + .filter_map(|t| t.join().unwrap()) + .collect::>() + }); + assert_eq!(reservations.len(), 6); + assert!(budget.reserve(1, true).is_err()); + let control = budget.reserve(1024 * 1024, false)?; + assert_eq!(budget.snapshot().reserved_proposal_bytes, 7 * 1024 * 1024); + drop(reservations); + drop(control); + assert_eq!(budget.snapshot().reserved_proposal_bytes, 0); + // Peer/HardState traffic must be included even though it never passed + // through producer admission and may be written by another local group. + written.fetch_add(6 * 1024 * 1024, Ordering::Release); + assert!(budget.reserve(1, true).is_err()); + assert!(budget.reserve(1024, false).is_ok()); + Ok(()) + } + + #[test] + fn collection_refresh_releases_measured_bytes_but_not_pending_proposals() -> Result<()> { + let directory = tempfile::tempdir()?; + let path = directory.path().join("segment"); + std::fs::write(&path, vec![7; 6 * 1024 * 1024])?; + let written = Arc::new(AtomicU64::new(6 * 1024 * 1024)); + let lifecycle = Arc::new(RwLock::new(())); + let budget = Budget::new( + directory.path(), + written.clone(), + 8 * 1024 * 1024, + lifecycle.clone(), + )?; + assert!(budget.reserve(1, true).is_err()); + let pending = budget.reserve(1024, false)?; + { + let _guard = lifecycle.write().unwrap(); + std::fs::remove_file(&path)?; + budget.refresh()?; // Maintenance owns the lifecycle lock: no scan. + assert!(budget.reserve(1, true).is_err()); + } + budget.refresh()?; + assert_eq!(budget.snapshot().accounted_disk_bytes, 0); + assert_eq!(budget.snapshot().reserved_proposal_bytes, 1024); + assert!(budget.reserve(6 * 1024 * 1024, true).is_err()); + drop(pending); + let all = budget.reserve(6 * 1024 * 1024, true)?; + assert!(budget.reserve(1, true).is_err()); + drop(all); + // Peer appends bypassing producer admission still count without a scan. + std::fs::write(&path, vec![9; 6 * 1024 * 1024])?; + written.fetch_add(6 * 1024 * 1024, Ordering::Release); + assert!(budget.reserve(1, true).is_err()); + budget.refresh()?; + assert!(budget.reserve(1, true).is_err()); + assert!(budget.reserve(usize::MAX, false).is_err()); + Ok(()) + } +} diff --git a/crates/lance-context-wal/src/lib.rs b/crates/lance-context-wal/src/lib.rs index 4ae2efac..5ac9820f 100644 --- a/crates/lance-context-wal/src/lib.rs +++ b/crates/lance-context-wal/src/lib.rs @@ -7,3 +7,6 @@ #[cfg(feature = "placement")] pub mod placement; + +#[cfg(feature = "admission")] +pub mod admission;