diff --git a/.github/workflows/python-publish.yml b/.github/workflows/python-publish.yml index 643f652c..558e4451 100644 --- a/.github/workflows/python-publish.yml +++ b/.github/workflows/python-publish.yml @@ -7,6 +7,7 @@ on: pull_request: paths: - .github/workflows/python-publish.yml + - scripts/ci/install-manylinux-deps.sh - Cargo.lock - python/Cargo.lock - python/pyproject.toml @@ -74,12 +75,7 @@ jobs: target: ${{ matrix.target }} rust-toolchain: stable before-script-linux: | - if command -v yum >/dev/null 2>&1; then - yum install -y curl make openssl-devel perl - elif command -v apt-get >/dev/null 2>&1; then - apt-get update - apt-get install -y curl libssl-dev ca-certificates make perl - fi + bash ../scripts/ci/install-manylinux-deps.sh PROTOC_VERSION=25.6 curl -LsS -o /tmp/protoc.zip "https://github.com/protocolbuffers/protobuf/releases/download/v${PROTOC_VERSION}/protoc-${PROTOC_VERSION}-linux-x86_64.zip" python3 - <<'PY' diff --git a/.github/workflows/python-test.yml b/.github/workflows/python-test.yml index 368ffcf6..e38f1a93 100644 --- a/.github/workflows/python-test.yml +++ b/.github/workflows/python-test.yml @@ -8,6 +8,7 @@ on: - python/** - crates/** - .github/workflows/python-test.yml + - scripts/ci/install-manylinux-deps.sh concurrency: group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }} @@ -61,12 +62,7 @@ jobs: rust-toolchain: stable sccache: "true" before-script-linux: | - if command -v yum >/dev/null 2>&1; then - yum install -y curl make openssl-devel perl - elif command -v apt-get >/dev/null 2>&1; then - apt-get update - apt-get install -y curl libssl-dev ca-certificates make perl - fi + bash ../scripts/ci/install-manylinux-deps.sh PROTOC_VERSION=25.6 curl -LsS -o /tmp/protoc.zip "https://github.com/protocolbuffers/protobuf/releases/download/v${PROTOC_VERSION}/protoc-${PROTOC_VERSION}-linux-x86_64.zip" python3 - <<'PY' diff --git a/crates/lance-context-wal/Cargo.toml b/crates/lance-context-wal/Cargo.toml index 7870029e..18e29e4b 100644 --- a/crates/lance-context-wal/Cargo.toml +++ b/crates/lance-context-wal/Cargo.toml @@ -13,7 +13,7 @@ runtime = ["replication", "dep:tokio", "tokio/sync", "tokio/rt"] lance-integration = ["replication", "ipc", "arrow-contract", "dep:lance", "dep:lance-table", "dep:lance-io", "dep:object_store", "dep:async-trait", "dep:futures", "dep:serde_json"] client-journal = ["dep:anyhow", "dep:raft-engine", "raft-engine/internals", "dep:bytes", "dep:serde", "dep:serde_json", "dep:tokio", "tokio/sync"] replication = ["raft-log", "record", "contract", "dep:bincode", "dep:slog", "dep:ring", "uuid/v4", "serde/rc"] -raft-log = ["placement", "dep:raft", "dep:raft-engine", "dep:raft-proto", "dep:protobuf"] +raft-log = ["placement", "admission", "dep:raft", "dep:raft-engine", "dep:raft-proto", "dep:protobuf"] vectored-io = ["dep:bytes", "dep:tokio"] runner = ["dep:anyhow", "dep:clap", "dep:libc", "dep:serde", "dep:serde_json", "dep:sha2"] contract = ["dep:serde", "dep:sha2"] diff --git a/crates/lance-context-wal/src/admission.rs b/crates/lance-context-wal/src/admission.rs index a1e45ec8..52fd99e2 100644 --- a/crates/lance-context-wal/src/admission.rs +++ b/crates/lance-context-wal/src/admission.rs @@ -4,7 +4,7 @@ //! snapshot/GC and replica-pressure coordination are required for sustained use. use std::path::{Path, PathBuf}; use std::sync::{ - atomic::{AtomicU64, Ordering}, + atomic::{AtomicBool, AtomicU64, Ordering}, Arc, Mutex, RwLock, TryLockError, }; @@ -16,7 +16,9 @@ pub struct Snapshot { pub measured_disk_bytes: u64, pub accounted_disk_bytes: u64, pub reserved_proposal_bytes: u64, + pub reserved_file_bytes: u64, pub data_limit: u64, + pub file_limit: u64, pub max_log_bytes: u64, } #[derive(Debug)] @@ -24,12 +26,15 @@ struct Accounting { measured: u64, counter_at_measurement: u64, pending: u64, + pending_files: u64, } #[derive(Debug)] pub struct Budget { directory: PathBuf, written: Arc, data_limit: AtomicU64, + file_limit: AtomicU64, + failed: AtomicBool, max: u64, accounting: Mutex, filesystem: Arc>, @@ -59,10 +64,13 @@ impl Budget { filesystem, max: max as u64, data_limit: AtomicU64::new((max - reserve) as u64), + file_limit: AtomicU64::new(max as u64), + failed: AtomicBool::new(false), accounting: Mutex::new(Accounting { measured, counter_at_measurement: counter, pending: 0, + pending_files: 0, }), })) } @@ -91,32 +99,88 @@ impl Budget { /// 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 { + self.reserve_below(bytes, false, || { + if data { + self.data_limit.load(Ordering::Acquire) + } else { + self.max + } + }) + } + fn reserve_below( + self: &Arc, + bytes: usize, + file: bool, + limit: impl FnOnce() -> u64, + ) -> Result { let bytes = u64::try_from(bytes)?; let mut state = self.accounting.lock().unwrap(); - let limit = if data { - self.data_limit.load(Ordering::Acquire) - } else { - self.max - }; + ensure!( + !self.failed.load(Ordering::Acquire), + "disk accounting failed" + ); ensure!( self.accounted(&state) - .saturating_add(state.pending) - .saturating_add(bytes) - <= limit, + .checked_add(state.pending) + .and_then(|n| n.checked_add(bytes)) + .is_some_and(|n| n <= limit()), "node proposal disk budget exhausted; data admission preserves control window" ); state.pending += bytes; + if file { + state.pending_files += bytes; + } Ok(Reservation { budget: self.clone(), bytes, + file, }) } + + /// Preserve a separate log/control/rewrite window from new image files. + /// Apply before exposing the node. Already held reservations remain charged. + pub fn limit_files(&self, bytes: usize) -> Result<()> { + let bytes = u64::try_from(bytes)?; + let _state = self.accounting.lock().unwrap(); + ensure!( + bytes > 0 && bytes <= self.file_limit.load(Ordering::Acquire), + "invalid file ceiling" + ); + self.file_limit.store(bytes, Ordering::Release); + Ok(()) + } + + /// Reserve complete file lengths before creation, rounded independently to + /// 4 KiB. Keep the owner until all writes (including buffered drops) finish. + /// Drop converts the whole reservation to counted growth: interrupted and + /// installed files remain charged until a later directory measurement. + /// Filesystem allocation can exceed this rounding; this is admission with + /// reserved headroom, not an OS quota or a physical free-space guarantee. + pub fn reserve_files(self: &Arc, lengths: &[u64]) -> Result { + let bytes = lengths.iter().try_fold(0u64, |sum, length| { + let rounded = length + .checked_add(4095) + .map(|n| n / 4096 * 4096) + .ok_or_else(|| anyhow::anyhow!("file reservation overflow"))?; + sum.checked_add(rounded) + .ok_or_else(|| anyhow::anyhow!("file reservation overflow")) + })?; + Ok(FileReservation(self.reserve_below( + usize::try_from(bytes)?, + true, + || self.file_limit.load(Ordering::Acquire), + )?)) + } /// 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<()> { + ensure!( + !self.failed.load(Ordering::Acquire), + "disk accounting failed" + ); let _files = match self.filesystem.try_read() { Ok(guard) => guard, Err(TryLockError::WouldBlock) => return Ok(()), @@ -126,7 +190,7 @@ impl Budget { } /// 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<()> { + pub(crate) 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(); @@ -141,20 +205,47 @@ impl Budget { Snapshot { measured_disk_bytes: state.measured, accounted_disk_bytes: self.accounted(&state), - reserved_proposal_bytes: state.pending, + reserved_proposal_bytes: state.pending - state.pending_files, + reserved_file_bytes: state.pending_files, data_limit: self.data_limit.load(Ordering::Acquire), + file_limit: self.file_limit.load(Ordering::Acquire), max_log_bytes: self.max, } } } + +/// A physical-file operation's charge outlives network cancellation. Unlike a +/// proposal reservation, dropping this owner does not return retained disk space. +pub struct FileReservation(Reservation); +impl Drop for FileReservation { + fn drop(&mut self) { + let budget = &self.0.budget; + // Publish growth before Reservation's drop releases its pending charge. + // A racing refresh sees the existing files or retains this growth delta. + if budget + .written + .try_update(Ordering::AcqRel, Ordering::Acquire, |n| { + n.checked_add(self.0.bytes) + }) + .is_err() + { + budget.failed.store(true, Ordering::Release); + } + } +} #[derive(Debug)] pub struct Reservation { budget: Arc, bytes: u64, + file: bool, } impl Drop for Reservation { fn drop(&mut self) { - self.budget.accounting.lock().unwrap().pending -= self.bytes; + let mut state = self.budget.accounting.lock().unwrap(); + state.pending -= self.bytes; + if self.file { + state.pending_files -= self.bytes; + } } } @@ -187,6 +278,51 @@ fn retained_disk_bytes(directory: &Path) -> Result { mod tests { use super::*; #[test] + fn interrupted_files_stay_charged_until_measurement_and_actual_collection() -> Result<()> { + let directory = tempfile::tempdir()?; + let lifecycle = Arc::new(RwLock::new(())); + let budget = Budget::new( + directory.path(), + Arc::new(AtomicU64::new(0)), + 8 << 20, + lifecycle.clone(), + )?; + budget.limit_data(4 << 20)?; + budget.limit_files(5 << 20)?; + let receiving = budget.reserve_files(&[3 << 20, 1])?; + assert!(budget.reserve(1 << 20, true).is_err()); + let path = directory.path().join("partial.pack"); + std::fs::write(&path, vec![7; 1 << 20])?; + // A measurement while the writer still exists cannot return its credit. + budget.refresh()?; + drop(receiving); + assert_eq!(budget.snapshot().reserved_proposal_bytes, 0); + assert!(budget.snapshot().accounted_disk_bytes >= 3 << 20); + assert!(budget.reserve(2 << 20, true).is_err()); + budget.refresh()?; + assert_eq!(budget.snapshot().accounted_disk_bytes, 1 << 20); + let reopened = Budget::new( + directory.path(), + Arc::new(AtomicU64::new(0)), + 8 << 20, + lifecycle.clone(), + )?; + assert_eq!(reopened.snapshot().accounted_disk_bytes, 1 << 20); + { + let _files = lifecycle.write().unwrap(); + std::fs::remove_file(path)?; + budget.refresh_exclusive()?; + } + assert_eq!(budget.snapshot().accounted_disk_bytes, 0); + assert!(budget.reserve_files(&[u64::MAX]).is_err()); + assert!(budget.reserve_files(&[5 << 20, 1]).is_err()); + assert!( + budget.reserve(6 << 20, false).is_ok(), + "control window remains independent" + ); + Ok(()) + } + #[test] fn lower_ceiling_keeps_pending_proposals_and_maintenance_charged() -> Result<()> { let directory = tempfile::tempdir()?; let budget = Budget::new( diff --git a/crates/lance-context-wal/src/raft_log.rs b/crates/lance-context-wal/src/raft_log.rs index 71822be1..d1ac224b 100644 --- a/crates/lance-context-wal/src/raft_log.rs +++ b/crates/lance-context-wal/src/raft_log.rs @@ -17,8 +17,10 @@ use std::collections::{BTreeMap, BTreeSet}; use std::path::Path; +#[cfg(any(test, feature = "replication"))] +use std::path::PathBuf; use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; -use std::sync::{Arc, Mutex, RwLock, RwLockReadGuard}; +use std::sync::{Arc, Mutex, OnceLock, RwLock, RwLockReadGuard}; use anyhow::{ensure, Context, Result}; use protobuf::Message; @@ -28,7 +30,9 @@ use raft_proto::eraftpb::{ConfState, Entry, EntryType, HardState, Snapshot}; use crate::placement::Layout; +mod budget; mod checkpoint; +pub use budget::{DiskMeasurement, DiskOptions}; mod maintenance; pub use maintenance::LogMaintenance; #[cfg(test)] @@ -60,8 +64,15 @@ impl MessageExt for EntryIndex { struct NodeEngine { engine: LogEngine, + #[cfg(any(test, feature = "replication"))] + directory: PathBuf, + #[cfg(feature = "replication")] + target_file_bytes: u64, + written: Arc, + budget: OnceLock>, failed: AtomicBool, - maintenance_gate: RwLock<()>, + maintenance_gate: Arc>, + filesystem: Arc>, append_gate: RwLock<()>, sync_gate: Mutex<()>, synced: AtomicU64, @@ -145,7 +156,19 @@ impl NodeEngine { match std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { self.engine.write(batch, sync) })) { - Ok(Ok(_)) => Ok(()), + Ok(Ok(bytes)) => { + if self + .written + .try_update(Ordering::AcqRel, Ordering::Acquire, |n| { + n.checked_add(bytes as u64) + }) + .is_err() + { + self.failed.store(true, Ordering::Release); + anyhow::bail!("physical write counter exhausted; reopen the node log"); + } + Ok(()) + } Ok(Err(error)) => { self.failed.store(true, Ordering::Release); Err(error).context("uncertain engine write; node log must be reopened") @@ -405,8 +428,15 @@ impl PhysicalLog { let node_id = node; let node = Arc::new(NodeEngine { engine, + #[cfg(any(test, feature = "replication"))] + directory: PathBuf::from(directory), + #[cfg(feature = "replication")] + target_file_bytes: options.target_file_bytes, + written: Arc::new(AtomicU64::new(0)), + budget: OnceLock::new(), failed: AtomicBool::new(false), - maintenance_gate: RwLock::new(()), + maintenance_gate: Arc::new(RwLock::new(())), + filesystem: Arc::new(RwLock::new(())), append_gate: RwLock::new(()), sync_gate: Mutex::new(()), synced: AtomicU64::new(0), diff --git a/crates/lance-context-wal/src/raft_log/budget.rs b/crates/lance-context-wal/src/raft_log/budget.rs new file mode 100644 index 00000000..1cd3e600 --- /dev/null +++ b/crates/lance-context-wal/src/raft_log/budget.rs @@ -0,0 +1,177 @@ +//! One physical directory's admission and measurement lifecycle. +use super::*; +#[cfg(any(test, feature = "replication"))] +use crate::admission::Budget; +use crate::admission::Snapshot; + +/// Local admission policy, independent of replicated layout/state. Producer +/// proposals stop at `data_bytes`; new images preserve `headroom_bytes` for +/// accepted Raft/control traffic and physical rewrites. Peer writes must still +/// persist and can consume this window. This is not a hard filesystem quota. +#[derive(Clone, Copy, Debug)] +pub struct DiskOptions { + pub max_bytes: usize, + pub data_bytes: usize, + pub headroom_bytes: usize, +} +#[cfg(any(test, feature = "replication"))] +impl DiskOptions { + pub(crate) fn validate(&self) -> Result<()> { + ensure!( + self.max_bytes >= 4 << 20 + && self.headroom_bytes > 0 + && self.data_bytes > 0 + && self.data_bytes < self.max_bytes + && self.headroom_bytes <= self.max_bytes - self.data_bytes, + "invalid node disk admission policy" + ); + Ok(()) + } +} + +/// A blocking measurement task retaining the engine's exclusive directory lock. +/// No payload reads/decoding. Appends overlap safely via their growth counter; +/// native rewrite and image deletion are excluded by the lifecycle guard. +pub struct DiskMeasurement { + node: Arc, +} +impl DiskMeasurement { + pub fn run(&self) -> Result { + self.node.check()?; + let budget = self + .node + .budget + .get() + .context("node disk budget not configured")?; + budget.refresh()?; + self.node.check()?; + Ok(budget.snapshot()) + } +} + +impl GroupLog { + #[cfg(feature = "replication")] + pub(crate) fn allocation_headroom(&self) -> Result { + self.node + .target_file_bytes + .checked_mul(2) + .context("log file headroom overflow") + } + #[cfg(any(test, feature = "replication"))] + pub(crate) fn configure_disk(&self, options: DiskOptions) -> Result<()> { + options.validate()?; + let _guard = self + .node + .maintenance_gate + .write() + .map_err(|_| anyhow::anyhow!("maintenance gate poisoned"))?; + self.node.check()?; + ensure!( + self.node.budget.get().is_none(), + "node disk budget already configured" + ); + let _files = self + .node + .filesystem + .write() + .map_err(|_| anyhow::anyhow!("filesystem lifecycle lock poisoned"))?; + let budget = Budget::new( + &self.node.directory, + self.node.written.clone(), + options.max_bytes, + self.node.filesystem.clone(), + )?; + budget.limit_data(options.data_bytes)?; + budget.limit_files(options.max_bytes - options.headroom_bytes)?; + self.node + .budget + .set(budget) + .map_err(|_| anyhow::anyhow!("node disk budget already configured"))?; + Ok(()) + } + #[cfg(any(test, feature = "replication"))] + pub(crate) fn disk_budget(&self) -> Option<&Arc> { + self.node.budget.get() + } + pub fn disk_measurement(&self) -> Result { + ensure!( + self.node.budget.get().is_some(), + "node disk budget not configured" + ); + Ok(DiskMeasurement { + node: self.node.clone(), + }) + } + #[cfg(any(test, feature = "replication"))] + pub(crate) fn image_collection_guard(&self) -> Result> { + self.node + .filesystem + .write() + .map_err(|_| anyhow::anyhow!("maintenance gate poisoned")) + } + #[cfg(feature = "replication")] + pub(crate) fn refresh_disk_exclusive(&self) -> Result<()> { + if let Some(budget) = self.node.budget.get() { + budget.refresh_exclusive()?; + } + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + #[test] + fn image_collection_lifecycle_does_not_block_a_real_log_append() -> Result<()> { + let layout = Layout { + cluster: uuid::Uuid::from_u128(1), + nodes: (1..=3) + .map(|id| (id, format!("wal-{id}.example.invalid:7000"))) + .collect(), + groups: BTreeMap::from([(10, vec![1, 2, 3]), (20, vec![1, 2, 3])]), + buckets: vec![10, 20], + }; + let directory = tempfile::tempdir()?; + let log = PhysicalLog::open(directory.path(), &layout, 1)?; + let image_owner = log.group(10)?; + image_owner.configure_disk(DiskOptions { + max_bytes: 512 << 20, + data_bytes: 128 << 20, + headroom_bytes: 256 << 20, + })?; + let lifecycle = image_owner.image_collection_guard()?; + let writer = log.group(20)?; + let (sent, received) = std::sync::mpsc::sync_channel(1); + let thread = std::thread::spawn(move || { + let result = writer.persist( + &[Entry { + index: 1, + term: 1, + data: vec![7; 4096].into(), + ..Default::default() + }], + Some(&HardState { + term: 1, + vote: 1, + commit: 1, + ..Default::default() + }), + ); + let _ = sent.send(result); + }); + let result = received.recv_timeout(std::time::Duration::from_secs(10)); + drop(lifecycle); + thread.join().unwrap(); + result.context("image deletion lifecycle blocked an independent log append")??; + assert_eq!(log.group(20)?.initial_state()?.hard_state.commit, 1); + assert!( + image_owner + .disk_budget() + .unwrap() + .snapshot() + .accounted_disk_bytes + > 4096 + ); + Ok(()) + } +} diff --git a/crates/lance-context-wal/src/raft_log/maintenance.rs b/crates/lance-context-wal/src/raft_log/maintenance.rs index 8dd1c6cc..b88deb4b 100644 --- a/crates/lance-context-wal/src/raft_log/maintenance.rs +++ b/crates/lance-context-wal/src/raft_log/maintenance.rs @@ -21,6 +21,19 @@ impl LogMaintenance { /// per-call byte/time bound, so callers must measure its latency and reserve /// rewrite headroom. This API alone does not impose a disk quota. pub fn run(&self) -> Result> { + // Admission failure occurs before touching native storage and must not + // fence a healthy node. Existing accepted traffic still owns its window. + let _reservation = self + .node + .budget + .get() + .map(|budget| { + budget.reserve( + usize::try_from(budget.snapshot().accounted_disk_bytes)?, + false, + ) + }) + .transpose()?; let result = self.run_locked(); if result.is_err() { self.node.failed.store(true, Ordering::Release); @@ -48,6 +61,11 @@ impl LogMaintenance { // The maintenance gate drained pending appends and checkpoint writes. // Establish an explicit durable marker before the native rewrite too. self.node.sync_locked()?; + let _files = self + .node + .filesystem + .write() + .map_err(|_| anyhow::anyhow!("filesystem lifecycle lock poisoned"))?; let outcome = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { self.node.engine.purge_expired_files() })); @@ -65,6 +83,11 @@ impl LogMaintenance { groups.retain(|group| *group != 0); groups.sort_unstable(); groups.dedup(); + if let Some(budget) = self.node.budget.get() { + // Native rewrites are outside Engine::write. While every appender + // and image collector is excluded, replace the measured baseline. + budget.refresh_exclusive()?; + } self.node.check()?; Ok(groups) } diff --git a/crates/lance-context-wal/src/replication.rs b/crates/lance-context-wal/src/replication.rs index f86f17e2..a31914cb 100644 --- a/crates/lance-context-wal/src/replication.rs +++ b/crates/lance-context-wal/src/replication.rs @@ -8,8 +8,9 @@ //! The owner schedules ticks, bounded network I/O and calls on one thread. This //! component supports explicit application checkpoints and prepared snapshot //! transfer and explicit collection of unreferenced application images. Raft -//! segment reclamation, production scheduling and bounded disk admission remain -//! separate responsibilities. +//! segment reclamation and production scheduling remain separate responsibilities. +//! Optional shared node disk admission covers proposals and image files while +//! preserving a window for accepted peer traffic; it is not a filesystem quota. //! Message scope checks are not authentication. Do not expose it to untrusted //! transport or enable ingestion without the remaining admission/lifecycle work. @@ -246,7 +247,7 @@ struct Group { raw: RawNode, state: State, // Index AND term prevents suffix replacement from completing another input. - pending: BTreeMap<(u64, u64), (u64, BatchMeta)>, + pending: BTreeMap<(u64, u64), (u64, BatchMeta, Option)>, reads: query::Tracker, image: Option>, incoming: Option, @@ -270,6 +271,44 @@ pub struct Node { } impl Node { + /// Configure one shared physical-directory budget before splitting group + /// owners or starting transports. Reopens reconstruct installed/orphan/log + /// usage from the same directory. This does not alter replicated identity. + pub fn with_disk_budget(self, options: crate::raft_log::DiskOptions) -> Result { + options.validate()?; + ensure!( + self.groups.len() + == self + .layout + .groups + .values() + .filter(|v| v.contains(&self.id)) + .count(), + "configure disk admission before splitting groups" + ); + let allocation = self.images.store().allocation_headroom()?; + let reserve = self + .limits + .max_uncommitted_bytes + .checked_mul(self.groups.len() as u64) + .and_then(|n| n.checked_add(allocation)) + .context("replication headroom overflow")?; + ensure!( + options.headroom_bytes as u64 >= reserve, + "headroom must cover hosted uncommitted windows and two log files" + ); + self.images.configure_disk(options)?; + Ok(self) + } + + pub fn disk_budget(&self) -> Option { + self.images.store().disk_budget().map(|b| b.snapshot()) + } + + pub fn prepare_disk_measurement(&self) -> Result { + self.healthy()?; + self.images.store().disk_measurement() + } /// Give each group to an independent worker while retaining the shared /// physical engine and its sync/failure fence. Consumes the aggregate owner; /// it never clones a RawNode or permits two owners of one group. @@ -679,6 +718,19 @@ impl Node { ); let index = g.raw.raft.raft_log.last_index() + 1; let term = g.raw.raft.term; + let reservation = g + .raw + .store() + .disk_budget() + .map(|budget| { + let bytes = batch + .body + .len() + .checked_add(4096) + .context("proposal disk reservation overflow")?; + budget.reserve(bytes, batch.meta.kind == crate::record::RecordKind::Data) + }) + .transpose()?; // RawNode::propose accepts Vec and would copy the opaque Bytes. This is // the equivalent raft-rs proposal path with shared entry.data ownership. let entry = Entry { @@ -695,7 +747,8 @@ impl Node { g.raw.raft.raft_log.last_index() == index, "proposal was not queued" ); - g.pending.insert((index, term), (request, batch.meta)); + g.pending + .insert((index, term), (request, batch.meta, reservation)); Ok(()) } @@ -819,7 +872,7 @@ impl Node { }) .collect(); for key in uncertain { - let (request, _) = group.pending.remove(&key).unwrap(); + let (request, _, _reservation) = group.pending.remove(&key).unwrap(); events.completions.push(Completion { group: *id, request, @@ -901,7 +954,9 @@ fn apply( "application precedes durable commit" ); let result = group.state.apply(&entry, limits)?; - if let Some((request, meta)) = group.pending.remove(&(entry.index, entry.term)) { + if let Some((request, meta, _reservation)) = + group.pending.remove(&(entry.index, entry.term)) + { ensure!( state::decode_meta(&entry.context)? == meta, "pending proposal identity changed" diff --git a/crates/lance-context-wal/src/replication/checkpoint.rs b/crates/lance-context-wal/src/replication/checkpoint.rs index 9a95d2e1..efdab081 100644 --- a/crates/lance-context-wal/src/replication/checkpoint.rs +++ b/crates/lance-context-wal/src/replication/checkpoint.rs @@ -462,6 +462,100 @@ fn locations(state: &State) -> Result<(BTreeMap, u64)> { Ok((locations, offset)) } +fn write_metadata( + metadata: &mut impl Write, + descriptor: &Descriptor, + state: &State, + limits: &Limits, +) -> Result<()> { + metadata.write_all(magic(limits))?; + write_frame(metadata, &codec().serialize(&descriptor.binding)?)?; + write_u64(metadata, descriptor.records)?; + write_u64(metadata, descriptor.producers)?; + write_u64(metadata, descriptor.tables)?; + write_u64(metadata, descriptor.pins)?; + write_u64(metadata, descriptor.closed)?; + write_u64(metadata, descriptor.retired)?; + for (index, meta) in &state.accepted { + write_meta(metadata, *index, meta)?; + } + for (meta, index) in state.producers.values() { + write_meta(metadata, *index, meta)?; + } + for table in state.catalog.tables.values() { + write_frame(metadata, &publication::encode_table(table)?)?; + } + for pin in state.pins.values() { + write_frame(metadata, &source::encode_pin(pin)?)?; + } + for closed in state.closed.values() { + write_frame(metadata, &source::encode_closed(closed)?)?; + } + for (key, retired) in &state.retired { + write_frame(metadata, &state::encode_retired(*key, *retired)?)?; + } + // A flat, policy-sized cell table; decoding never trusts an allocation count + // supplied by the file. Zero-slot groups retain the exact legacy layout. + if limits.producer_slots > 0 { + write_u64(metadata, u64::from(limits.producer_slots))?; + for cell in &state.producer_slots { + write_frame(metadata, &producer::encode_cell(cell)?)?; + } + } + Ok(()) +} + +#[derive(Default)] +struct Count(u64); +impl Write for Count { + fn write(&mut self, bytes: &[u8]) -> std::io::Result { + self.0 = self + .0 + .checked_add(bytes.len() as u64) + .ok_or_else(|| std::io::Error::other("image length overflow"))?; + Ok(bytes.len()) + } + fn flush(&mut self) -> std::io::Result<()> { + Ok(()) + } +} + +/// Count only bounded metadata; payloads remain opaque and are never decoded or +/// copied here. The writer uses the same encoder, so this is its exact length. +pub(super) fn build_sizes( + layout: &Layout, + limits: &Limits, + cut: &Cut, + state: &State, +) -> Result<[u64; 2]> { + let payload_bytes = state.accepted.values().try_fold(0u64, |n, meta| { + n.checked_add(meta.body_bytes) + .context("image payload overflow") + })?; + let descriptor = Descriptor { + format: format(limits), + binding: Binding::new(layout, limits, cut.group, cut.index, cut.term)?, + id: cut.id, + records: state.accepted.len() as u64, + producers: state.producers.len() as u64, + tables: state.catalog.tables.len() as u64, + pins: state.pins.len() as u64, + closed: state.closed.len() as u64, + retired: state.retired.len() as u64, + metadata_bytes: 0, + payload_bytes, + metadata_digest: [0; 32], + payload_digest: [0; 32], + }; + let mut count = Count::default(); + write_metadata(&mut count, &descriptor, state, limits)?; + ensure!( + count.0 <= metadata_limit(limits), + "image metadata exceeds policy" + ); + Ok([count.0, payload_bytes]) +} + /// Inputs must be a captured, immutable applied cut. The owner retains its log /// and prior image while this worker reads one bounded body at a time. pub(super) fn build( @@ -499,40 +593,7 @@ pub(super) fn build( }; let paths = descriptor.paths(directory); let mut metadata = Writer::new(&paths[0], metadata_limit(limits))?; - metadata.write_all(magic(limits))?; - write_frame(&mut metadata, &codec().serialize(&descriptor.binding)?)?; - write_u64(&mut metadata, descriptor.records)?; - write_u64(&mut metadata, descriptor.producers)?; - write_u64(&mut metadata, descriptor.tables)?; - write_u64(&mut metadata, descriptor.pins)?; - write_u64(&mut metadata, descriptor.closed)?; - write_u64(&mut metadata, descriptor.retired)?; - for (index, meta) in &state.accepted { - write_meta(&mut metadata, *index, meta)?; - } - for (meta, index) in state.producers.values() { - write_meta(&mut metadata, *index, meta)?; - } - for table in state.catalog.tables.values() { - write_frame(&mut metadata, &publication::encode_table(table)?)?; - } - for pin in state.pins.values() { - write_frame(&mut metadata, &source::encode_pin(pin)?)?; - } - for closed in state.closed.values() { - write_frame(&mut metadata, &source::encode_closed(closed)?)?; - } - for (key, retired) in &state.retired { - write_frame(&mut metadata, &state::encode_retired(*key, *retired)?)?; - } - // A flat, policy-sized cell table; decoding never trusts an allocation count - // supplied by the file. Zero-slot groups retain the exact legacy layout. - if limits.producer_slots > 0 { - write_u64(&mut metadata, u64::from(limits.producer_slots))?; - for cell in &state.producer_slots { - write_frame(&mut metadata, &producer::encode_cell(cell)?)?; - } - } + write_metadata(&mut metadata, &descriptor, &state, limits)?; (descriptor.metadata_bytes, descriptor.metadata_digest) = metadata.finish()?; let mut pack = Writer::new(&paths[1], payload_bytes)?; for (index, meta) in &state.accepted { diff --git a/crates/lance-context-wal/src/replication/images.rs b/crates/lance-context-wal/src/replication/images.rs index 7ea671ca..6b2ec329 100644 --- a/crates/lance-context-wal/src/replication/images.rs +++ b/crates/lance-context-wal/src/replication/images.rs @@ -31,6 +31,16 @@ struct Registry { installed: BTreeMap, } impl Owner { + pub(super) fn configure_disk( + self: &Arc, + options: crate::raft_log::DiskOptions, + ) -> Result<()> { + let _permit = self.acquire()?; + self.store.configure_disk(options) + } + pub(super) fn store(&self) -> &GroupLog { + &self.store + } pub fn new(directory: &Path, groups: BTreeSet, store: GroupLog) -> Result> { Ok(Arc::new(Self { directory: directory.canonicalize()?, @@ -111,6 +121,14 @@ impl CheckpointTask { state, } = self; let group = cut.group; + let _disk = permit + .0 + .store + .disk_budget() + .map(|budget| { + budget.reserve_files(&checkpoint::build_sizes(&layout, &limits, &cut, &state)?) + }) + .transpose()?; let loaded = checkpoint::build( &permit.0.directory, &layout, @@ -151,9 +169,13 @@ impl PreparedCheckpoint { /// Bounded transfer state. Socket I/O stays outside the Raft owner; move this /// object to the admitted worker and feed exact, sequential file chunks. pub struct SnapshotReceiver { - permit: Permit, packet: Packet, receiver: checkpoint::Receiver, + // Declared after receiver: buffered file writes must finish before the + // pending charge turns into counted retained growth, including on Drop. + _disk: Option, + // Release the image/engine owner only after buffered writers and accounting. + permit: Permit, } impl SnapshotReceiver { pub fn write(&mut self, part: Part, offset: u64, bytes: &[u8]) -> Result<()> { @@ -298,8 +320,14 @@ impl Node { } pub fn receive_snapshot(&self, packet: Packet, local_id: Uuid) -> Result { - self.snapshot_sizes(&packet)?; + let sizes = self.snapshot_sizes(&packet)?; let permit = self.images.acquire()?; + let disk = self + .images + .store + .disk_budget() + .map(|budget| budget.reserve_files(&sizes)) + .transpose()?; let receiver = checkpoint::Receiver::new( &self.images.directory, packet.message.get_snapshot(), @@ -312,6 +340,7 @@ impl Node { permit, packet, receiver, + _disk: disk, }) } @@ -397,7 +426,7 @@ pub(super) fn install_received( prepared.permit.0.installed(&prepared.loaded.image)?; g.state = prepared.loaded.state; g.image = Some(prepared.loaded.image); - for (_, (request, _)) in std::mem::take(&mut g.pending) { + for (_, (request, _, _reservation)) in std::mem::take(&mut g.pending) { events.completions.push(Completion { group: id, request, diff --git a/crates/lance-context-wal/src/replication/images/collection.rs b/crates/lance-context-wal/src/replication/images/collection.rs index d01b6562..11fb686a 100644 --- a/crates/lance-context-wal/src/replication/images/collection.rs +++ b/crates/lance-context-wal/src/replication/images/collection.rs @@ -42,6 +42,7 @@ impl ImageCollection { "invalid image collection batch size" ); self.permit.0.store.initial_state()?; + let _files = self.permit.0.store.image_collection_guard()?; let mut progress = ImageCollectionProgress::default(); for _ in 0..max_entries { if self.finished { @@ -88,6 +89,9 @@ impl ImageCollection { File::open(&self.permit.0.directory)?.sync_all()?; } progress.finished = self.finished; + if self.finished { + self.permit.0.store.refresh_disk_exclusive()?; + } Ok(progress) } } diff --git a/crates/lance-context-wal/src/replication/runtime.rs b/crates/lance-context-wal/src/replication/runtime.rs index 67cb3f10..50df6f42 100644 --- a/crates/lance-context-wal/src/replication/runtime.rs +++ b/crates/lance-context-wal/src/replication/runtime.rs @@ -8,7 +8,7 @@ //! A transport must authenticate peers/callers, enforce frame limits before //! allocation, hold outbound ownership through transmission, and implement //! snapshot transfer and maintenance. This module is the worker layer, not a -//! network listener, disk admission policy, or complete ingestion service. +//! network listener, automatic maintenance scheduler, or complete ingestion service. use super::{producer, publication as p, query, source, Node, Outcome, Packet, Status}; use crate::record::{Batch, RecordKind}; use anyhow::{ensure, Context, Result}; @@ -270,6 +270,7 @@ impl Outbound { #[derive(Clone)] pub struct Handle { + disk: Option>, routes: Arc>, shared: Arc, catalog: u64, @@ -449,8 +450,8 @@ pub struct Runtime { } impl Runtime { /// Pass an already opened, configured node with any required signing keys. - /// Queue credits do not establish physical disk bounds. Integrate node-wide - /// disk admission, maintenance and snapshot transfer before exposing ingress. + /// Queue credits do not establish physical disk bounds. Configure the Node's + /// shared disk admission, maintenance and snapshot transfer before ingress. pub fn spawn(node: Node, options: Options) -> Result<(Self, async_mpsc::Receiver)> { options.validate()?; node.healthy()?; @@ -466,6 +467,7 @@ impl Runtime { }); let (outbox, receive) = async_mpsc::channel(options.peer_requests); let mut handle = Handle { + disk: node.images.store().disk_budget().cloned(), routes: Arc::new(BTreeMap::new()), shared: shared.clone(), catalog: p::control_group(&node.layout), diff --git a/crates/lance-context-wal/src/replication/runtime/images.rs b/crates/lance-context-wal/src/replication/runtime/images.rs index ffa5761b..6e03a4bc 100644 --- a/crates/lance-context-wal/src/replication/runtime/images.rs +++ b/crates/lance-context-wal/src/replication/runtime/images.rs @@ -111,6 +111,15 @@ impl Worker { } impl Handle { + pub fn disk_budget(&self) -> Option { + self.disk.as_ref().map(|budget| budget.snapshot()) + } + /// Capture a blocking directory measurement retaining the physical engine. + /// Run this on the maintenance worker, never on an async executor thread. + pub fn prepare_disk_measurement(&self) -> Result> { + let group = *self.routes.keys().next().context("no hosted group")?; + self.local(group, |node| node.prepare_disk_measurement()) + } /// Validate exact incoming image lengths without creating files. The /// authenticated transport uses this before reserving physical disk space. pub fn snapshot_sizes(&self, packet: Packet) -> Result> { diff --git a/crates/lance-context-wal/src/replication/runtime/peer/snapshot.rs b/crates/lance-context-wal/src/replication/runtime/peer/snapshot.rs index b3afcfd1..5dd2f9ec 100644 --- a/crates/lance-context-wal/src/replication/runtime/peer/snapshot.rs +++ b/crates/lance-context-wal/src/replication/runtime/peer/snapshot.rs @@ -1,8 +1,8 @@ //! Authenticated image streams on separate sockets and one joined file worker. -//! Snapshot I/O never occupies ordinary peer sockets. Disk admission is supplied -//! by the embedding service and must account for installed and orphaned files, -//! concurrent log growth and maintenance headroom; concurrent-transfer credits -//! alone are not a physical disk quota. +//! Snapshot I/O never occupies ordinary peer sockets. A configured node disk +//! budget is required; it reserves exact file lengths before creation and keeps +//! installed and orphaned files charged alongside log growth. An optional extra +//! admission policy may impose stricter service limits. This is not an OS quota. use super::*; use crate::replication::checkpoint::{Image, Part, CHUNK_BYTES}; use std::{future::Future, thread}; @@ -31,12 +31,13 @@ impl Incoming { } } -/// Required physical disk admission, executed on the bounded file worker. +/// Optional additional admission, executed on the bounded file worker. +/// The required node budget is enforced independently before file creation. /// The returned owner spans pending file work, even when its network task is /// cancelled. Dropping it MUST NOT release space still used by installed or /// orphaned files. Reclaim that charge only after actual collection. Reopening /// the service must recover accounting from its existing files. There is no -/// permissive default implementation. +/// weakening of the node budget when this additional policy is absent. pub trait Admission: Send + Sync + 'static { fn reserve(&self, incoming: &Incoming) -> Result>; } @@ -72,7 +73,7 @@ pub struct Settings { pub client_tls: Arc, pub server_tls: Arc, pub options: Options, - pub admission: Arc, + pub admission: Option>, } #[derive(Default)] @@ -148,6 +149,10 @@ pub struct Transport { } impl Transport { pub fn spawn(handle: Handle, listener: TcpListener, settings: Settings) -> Result { + ensure!( + handle.disk.is_some(), + "snapshot transport requires configured node disk admission" + ); let Settings { peers, client_tls, @@ -367,7 +372,7 @@ async fn listen( fingerprints: Arc>, files: Files, options: Options, - admission: Arc, + admission: Option>, stop: watch::Receiver, statistics: Arc, ) -> Result<()> { @@ -519,7 +524,7 @@ async fn receive_image( peer: u64, files: &Files, options: &Options, - admission: Arc, + admission: Option>, ) -> Result<()> { let packet = read_header(stream, handle, peer).await?; let sizes = handle.snapshot_sizes(packet.clone())?.await?; @@ -534,7 +539,9 @@ async fn receive_image( incoming.bytes()? <= options.max_image_bytes, "snapshot exceeds receive budget" ); - let reservation = files.run(move || admission.reserve(&incoming)).await?; + let reservation = files + .run(move || admission.map(|a| a.reserve(&incoming)).transpose()) + .await?; let mut receiver = handle .receive_snapshot(packet, uuid::Uuid::new_v4())? .await?; diff --git a/crates/lance-context-wal/src/replication/runtime/peer/tests.rs b/crates/lance-context-wal/src/replication/runtime/peer/tests.rs index 5b0bd703..f5ab94bb 100644 --- a/crates/lance-context-wal/src/replication/runtime/peer/tests.rs +++ b/crates/lance-context-wal/src/replication/runtime/peer/tests.rs @@ -106,7 +106,8 @@ impl Cluster { Ok(c) } fn start(&mut self, root: &Path, id: u64, listener: TcpListener) -> Result<()> { - let node = Node::open(&root.join(id.to_string()), layout(), id, Limits::default())?; + let node = Node::open(&root.join(id.to_string()), layout(), id, Limits::default())? + .with_disk_budget(disk_options())?; let (runtime, outbox) = Runtime::spawn(node, super::super::Options::default())?; let handle = runtime.handle(); let peers = self @@ -468,3 +469,11 @@ async fn framing_shares_outbound_payload_and_rejects_unfunded_bodies_before_read assert_eq!(credits.available_permits(), 4 << 20); runtime.shutdown().await } + +fn disk_options() -> crate::raft_log::DiskOptions { + crate::raft_log::DiskOptions { + max_bytes: 512 << 20, + data_bytes: 128 << 20, + headroom_bytes: 256 << 20, + } +} diff --git a/crates/lance-context-wal/src/replication/runtime/peer/tests/process.rs b/crates/lance-context-wal/src/replication/runtime/peer/tests/process.rs index aded8d69..d094144a 100644 --- a/crates/lance-context-wal/src/replication/runtime/peer/tests/process.rs +++ b/crates/lance-context-wal/src/replication/runtime/peer/tests/process.rs @@ -145,7 +145,8 @@ async fn tls_peer_process_fixture() -> Result<()> { layout(), id, Limits::default(), - )?; + )? + .with_disk_budget(super::disk_options())?; let (runtime, outbox) = Runtime::spawn(node, super::super::super::Options::default())?; let handle = runtime.handle(); let peers = addresses @@ -200,9 +201,7 @@ async fn tls_peer_process_fixture() -> Result<()> { .map(PathBuf::from), ..Default::default() }, - admission: Arc::new(super::snapshots::Quota::open( - &root.join(format!("node-{id}")), - )?), + admission: None, }, )?; network.attach_snapshots(images.route())?; diff --git a/crates/lance-context-wal/src/replication/runtime/peer/tests/snapshots.rs b/crates/lance-context-wal/src/replication/runtime/peer/tests/snapshots.rs index 269c7df6..767e3269 100644 --- a/crates/lance-context-wal/src/replication/runtime/peer/tests/snapshots.rs +++ b/crates/lance-context-wal/src/replication/runtime/peer/tests/snapshots.rs @@ -105,7 +105,7 @@ impl Images { transfer_timeout: Duration::from_secs(60), ..Default::default() }, - admission, + admission: Some(admission), }, )?; self.net @@ -179,23 +179,30 @@ async fn append(handle: &Handle, group: u64, sequence: u64) -> Result<(u64, Batc async fn tls_snapshots_catch_up_two_compacted_groups_and_reopen_original_payloads() -> Result<()> { let root = tempfile::tempdir()?; let mut c = Images::open(root.path()).await?; + let mut leaders = BTreeMap::new(); for group in [10, 20] { c.net.nodes[&1].handle.campaign(group)?.await?; - assert_eq!(c.net.leader(group).await?, 1); + leaders.insert(group, c.net.leader(group).await?); } - c.stop(3).await?; + // A campaign is only a request. Choose a voter that is a follower in + // both groups after the actual elections have settled. + let follower = (1..=3) + .find(|id| !leaders.values().any(|leader| leader == id)) + .unwrap(); + c.stop(follower).await?; let mut expected = Vec::new(); for group in [10, 20] { for sequence in 1..=4 { - let (index, batch) = append(&c.net.nodes[&1].handle, group, sequence).await?; + let (index, batch) = + append(&c.net.nodes[&leaders[&group]].handle, group, sequence).await?; expected.push((group, index, batch)); } - drop(checkpoint(&c.net.nodes[&1].handle, group).await?); + drop(checkpoint(&c.net.nodes[&leaders[&group]].handle, group).await?); } - c.start(root.path(), 3).await?; + c.start(root.path(), follower).await?; wait_for(|| { [10, 20].iter().all(|g| { - c.net.nodes[&3] + c.net.nodes[&follower] .handle .observation(*g) .unwrap() @@ -205,16 +212,18 @@ async fn tls_snapshots_catch_up_two_compacted_groups_and_reopen_original_payload }) }) .await?; - assert!( - c.services[&3] + // State installation precedes ACK flushing and transport accounting. + wait_for(|| { + c.services[&follower] .statistics() .installed .load(Ordering::Acquire) >= 2 - ); - assert!(c.quotas[&3].used.load(Ordering::Acquire) >= 8 << 20); + }) + .await?; + assert!(c.quotas[&follower].used.load(Ordering::Acquire) >= 8 << 20); for group in [10, 20] { - let (index, batch) = append(&c.net.nodes[&1].handle, group, 5).await?; + let (index, batch) = append(&c.net.nodes[&leaders[&group]].handle, group, 5).await?; expected.push((group, index, batch)); } wait_for(|| { @@ -238,7 +247,7 @@ async fn tls_snapshots_catch_up_two_compacted_groups_and_reopen_original_payload } } let c = Images::open(root.path()).await?; - assert!(c.quotas[&3].used.load(Ordering::Acquire) >= 8 << 20); + assert!(c.quotas[&follower].used.load(Ordering::Acquire) >= 8 << 20); for group in [10, 20] { // Reopening every engine may already have elected another voter. // A campaign is not a leadership receipt; retry through the actual @@ -263,22 +272,23 @@ async fn stalled_tls_snapshot_does_not_block_healthy_quorum_and_recovers_after_d let root = tempfile::tempdir()?; let mut c = Images::open(root.path()).await?; c.net.nodes[&1].handle.campaign(10)?.await?; - assert_eq!(c.net.leader(10).await?, 1); - c.stop(3).await?; + let leader_id = c.net.leader(10).await?; + let follower = (1..=3).rfind(|id| *id != leader_id).unwrap(); + c.stop(follower).await?; let mut expected = Vec::new(); for sequence in 1..=4 { - expected.push(append(&c.net.nodes[&1].handle, 10, sequence).await?); + expected.push(append(&c.net.nodes[&leader_id].handle, 10, sequence).await?); } - drop(checkpoint(&c.net.nodes[&1].handle, 10).await?); - // Node3 handles ordinary Raft messages, but its separate image socket is a + drop(checkpoint(&c.net.nodes[&leader_id].handle, 10).await?); + // The follower handles ordinary Raft messages, but its separate image socket is a // correctly authenticated peer that never grants receive admission. c.net.start( root.path(), - 3, - TcpListener::bind(c.net.addresses[&3]).await?, + follower, + TcpListener::bind(c.net.addresses[&follower]).await?, )?; - let listener = TcpListener::bind(c.addresses[&3]).await?; - let acceptor = TlsAcceptor::from(image_tls(&c.net.tls[&3]).1); + let listener = TcpListener::bind(c.addresses[&follower]).await?; + let acceptor = TlsAcceptor::from(image_tls(&c.net.tls[&follower]).1); let connected = Arc::new(AtomicBool::new(false)); let observed = connected.clone(); let stalled = tokio::spawn(async move { @@ -295,7 +305,7 @@ async fn stalled_tls_snapshot_does_not_block_healthy_quorum_and_recovers_after_d // Transfers have a60sdeadline; quorum writes must finish much earlier. timeout(Duration::from_secs(10), async { for sequence in 5..=12 { - expected.push(append(&c.net.nodes[&1].handle, 10, sequence).await?); + expected.push(append(&c.net.nodes[&leader_id].handle, 10, sequence).await?); } Ok::<_, anyhow::Error>(()) }) @@ -303,9 +313,9 @@ async fn stalled_tls_snapshot_does_not_block_healthy_quorum_and_recovers_after_d assert!(!stalled.is_finished()); stalled.abort(); let _ = stalled.await; - c.attach(3, TcpListener::bind(c.addresses[&3]).await?)?; + c.attach(follower, TcpListener::bind(c.addresses[&follower]).await?)?; wait_for(|| { - c.net.nodes[&3] + c.net.nodes[&follower] .handle .observation(10) .unwrap() @@ -314,13 +324,14 @@ async fn stalled_tls_snapshot_does_not_block_healthy_quorum_and_recovers_after_d == 12 }) .await?; - assert!( - c.services[&3] + wait_for(|| { + c.services[&follower] .statistics() .installed .load(Ordering::Acquire) > 0 - ); + }) + .await?; c.shutdown().await?; for id in 1..=3 { let node = Node::open( @@ -341,10 +352,13 @@ async fn upload_header( packet: &Packet, ) -> Result> { use protobuf::Message as _; - let connector = TlsConnector::from(image_tls(&c.net.tls[&1]).0); - let tcp = TcpStream::connect(c.addresses[&3]).await?; + let connector = TlsConnector::from(image_tls(&c.net.tls[&packet.message.from]).0); + let tcp = TcpStream::connect(c.addresses[&packet.message.to]).await?; let mut stream = connector - .connect(ServerName::try_from(c.net.tls[&3].name.clone())?, tcp) + .connect( + ServerName::try_from(c.net.tls[&packet.message.to].name.clone())?, + tcp, + ) .await?; let mut header = b"LCWSNAP1".to_vec(); header.extend_from_slice(packet.cluster.as_bytes()); @@ -376,38 +390,39 @@ async fn tls_snapshot_denial_and_interruption_never_install_or_release_disk_char let root = tempfile::tempdir()?; let mut c = Images::open(root.path()).await?; c.net.nodes[&1].handle.campaign(10)?.await?; - assert_eq!(c.net.leader(10).await?, 1); + let leader_id = c.net.leader(10).await?; + let follower = (1..=3).rfind(|id| *id != leader_id).unwrap(); // Retain the actual receiver runtime and snapshot service, but remove its // normal transport so only the deliberate TLS uploads can reach it. let Running { network, runtime, handle, - } = c.net.nodes.remove(&3).unwrap(); + } = c.net.nodes.remove(&follower).unwrap(); network.shutdown().await?; - let (index, expected) = append(&c.net.nodes[&1].handle, 10, 1).await?; - let image = checkpoint(&c.net.nodes[&1].handle, 10).await?; - let leader = &c.net.nodes[&1].handle; + let (index, expected) = append(&c.net.nodes[&leader_id].handle, 10, 1).await?; + let image = checkpoint(&c.net.nodes[&leader_id].handle, 10).await?; + let leader = &c.net.nodes[&leader_id].handle; let packet = Packet { cluster: leader.cluster, layout: leader.layout, policy: leader.policy, group: 10, message: Message { - from: 1, - to: 3, + from: leader_id, + to: follower, term: leader.observation(10)?.status.term, msg_type: MessageType::MsgSnapshot, snapshot: Some(image.snapshot()?).into(), ..Default::default() }, }; - let path = root.path().join("3"); + let path = root.path().join(follower.to_string()); let before = handle.observation(10)?.status.applied; let files = image_files(&path)?; - c.quotas[&3] + c.quotas[&follower] .used - .store(c.quotas[&3].limit, Ordering::Release); + .store(c.quotas[&follower].limit, Ordering::Release); let mut denied = upload_header(&c, &packet).await?; let mut ack = [0; 8]; no_response(&mut denied).await?; @@ -417,7 +432,7 @@ async fn tls_snapshot_denial_and_interruption_never_install_or_release_disk_char // The rejected request created no files or charge; open a fresh fixture // allowance. A partial reception then retains its full conservative charge. - c.quotas[&3].used.store(0, Ordering::Release); + c.quotas[&follower].used.store(0, Ordering::Release); let mut partial = upload_header(&c, &packet).await?; assert_eq!(partial.read_u8().await?, 1); partial.write_all(&[0]).await?; @@ -425,7 +440,7 @@ async fn tls_snapshot_denial_and_interruption_never_install_or_release_disk_char no_response(&mut partial).await?; assert_eq!(handle.observation(10)?.status.applied, before); let bytes = image.bytes(Part::Metadata) + image.bytes(Part::Payload); - assert_eq!(c.quotas[&3].used.load(Ordering::Acquire), bytes); + assert_eq!(c.quotas[&follower].used.load(Ordering::Acquire), bytes); assert!(image_files(&path)? > files); drop(partial); @@ -446,13 +461,13 @@ async fn tls_snapshot_denial_and_interruption_never_install_or_release_disk_char complete.flush().await?; timeout(Duration::from_secs(5), complete.read_exact(&mut ack)).await??; assert_eq!(&ack, b"LCWSYNC1"); - assert_eq!(c.quotas[&3].used.load(Ordering::Acquire), bytes * 2); + assert_eq!(c.quotas[&follower].used.load(Ordering::Acquire), bytes * 2); wait_for(|| handle.observation(10).unwrap().status.accepted_records == 1).await?; drop(complete); drop(image); c.shutdown().await?; runtime.shutdown().await?; - let node = Node::open(&path, layout(), 3, Limits::default())?; + let node = Node::open(&path, layout(), follower, Limits::default())?; assert_eq!(node.read_accepted(10, index)?, Some(expected)); Ok(()) } @@ -485,37 +500,45 @@ async fn snapshot_file_worker_panic_stops_service_and_allows_original_engine_rec let root = tempfile::tempdir()?; let mut c = Images::open(root.path()).await?; c.net.nodes[&1].handle.campaign(10)?.await?; - assert_eq!(c.net.leader(10).await?, 1); - c.stop(3).await?; - let (index, expected) = append(&c.net.nodes[&1].handle, 10, 1).await?; - drop(checkpoint(&c.net.nodes[&1].handle, 10).await?); + let leader_id = c.net.leader(10).await?; + let follower = (1..=3).rfind(|id| *id != leader_id).unwrap(); + c.stop(follower).await?; + let (index, expected) = append(&c.net.nodes[&leader_id].handle, 10, 1).await?; + drop(checkpoint(&c.net.nodes[&leader_id].handle, 10).await?); c.net.start( root.path(), - 3, - TcpListener::bind(c.net.addresses[&3]).await?, + follower, + TcpListener::bind(c.net.addresses[&follower]).await?, )?; c.attach_with( - 3, - TcpListener::bind(c.addresses[&3]).await?, + follower, + TcpListener::bind(c.addresses[&follower]).await?, Arc::new(PanicAdmission), )?; - wait_for(|| c.services[&3].statistics().failed.load(Ordering::Acquire) > 0).await?; - let failed = c.services.remove(&3).unwrap(); + wait_for(|| { + c.services[&follower] + .statistics() + .failed + .load(Ordering::Acquire) + > 0 + }) + .await?; + let failed = c.services.remove(&follower).unwrap(); assert!(timeout(Duration::from_secs(5), failed.shutdown()) .await? .is_err()); assert_eq!( - c.net.nodes[&3] + c.net.nodes[&follower] .handle .observation(10)? .status .accepted_records, 0 ); - c.net.stop(3).await?; - c.start(root.path(), 3).await?; + c.net.stop(follower).await?; + c.start(root.path(), follower).await?; wait_for(|| { - c.net.nodes[&3] + c.net.nodes[&follower] .handle .observation(10) .unwrap() @@ -525,7 +548,12 @@ async fn snapshot_file_worker_panic_stops_service_and_allows_original_engine_rec }) .await?; c.shutdown().await?; - let node = Node::open(&root.path().join("3"), layout(), 3, Limits::default())?; + let node = Node::open( + &root.path().join(follower.to_string()), + layout(), + follower, + Limits::default(), + )?; assert_eq!(node.read_accepted(10, index)?, Some(expected)); Ok(()) } diff --git a/crates/lance-context-wal/src/replication/tests.rs b/crates/lance-context-wal/src/replication/tests.rs index dac0f9e0..bae49d0e 100644 --- a/crates/lance-context-wal/src/replication/tests.rs +++ b/crates/lance-context-wal/src/replication/tests.rs @@ -1,6 +1,8 @@ use super::*; use bytes::Bytes; +#[path = "tests/disk_budget.rs"] +mod disk_budget; #[path = "tests/image_collection.rs"] mod image_collection; diff --git a/crates/lance-context-wal/src/replication/tests/disk_budget.rs b/crates/lance-context-wal/src/replication/tests/disk_budget.rs new file mode 100644 index 00000000..b45563da --- /dev/null +++ b/crates/lance-context-wal/src/replication/tests/disk_budget.rs @@ -0,0 +1,220 @@ +use super::*; +use crate::raft_log::{DiskOptions, LogOptions}; + +fn limits() -> Limits { + Limits { + max_record_bytes: 1 << 20, + max_uncommitted_bytes: 1 << 20, + ..Limits::default() + } +} +fn open(root: &Path, id: u64) -> Result { + Node::open_with_storage_options( + &root.join(id.to_string()), + layout(), + id, + limits(), + LogOptions { + target_file_bytes: 16 << 10, + purge_threshold_bytes: 64 << 10, + }, + )? + .with_disk_budget(DiskOptions { + max_bytes: 32 << 20, + data_bytes: 8 << 20, + headroom_bytes: 8 << 20, + }) +} +fn payload(group: u64, sequence: u64) -> Batch { + Batch::new( + Uuid::from_u128(2), + u32::from(group == 20), + Uuid::from_u128(group.into()), + sequence, + Bytes::from(vec![sequence as u8; 128 << 10]), + ) +} +fn packet(leader: &Node, to: u64, image: &checkpoint::Image) -> Result { + Ok(Packet { + cluster: leader.layout.cluster, + layout: leader.layout_hash, + policy: leader.policy, + group: 10, + message: Message { + from: 1, + to, + term: leader.status(10)?.term, + msg_type: MessageType::MsgSnapshot, + snapshot: Some(image.snapshot()?).into(), + ..Default::default() + }, + }) +} + +#[test] +fn shared_budget_preserves_control_reopens_orphans_and_reclaims_only_deleted_images() -> Result<()> +{ + use checkpoint::Part; + let root = tempfile::tempdir()?; + let mut nodes = (1..=3) + .map(|id| open(root.path(), id)) + .collect::>>()?; + nodes[0].campaign(10)?; + nodes[0].campaign(20)?; + pump(&mut nodes)?; + let mut accepted = Vec::new(); + for group in [10, 20] { + for sequence in 1..=8 { + let token = group * 100 + sequence; + nodes[0].propose(group, token, payload(group, sequence))?; + let events = pump(&mut nodes)?; + let Outcome::Committed(Applied::Accepted { index }) = events + .iter() + .find(|e| e.request == token) + .context("missing append completion")? + .outcome + else { + anyhow::bail!("append rejected"); + }; + accepted.push((group, sequence, index)); + } + } + for node in &nodes { + assert!(node.disk_budget().unwrap().accounted_disk_bytes >= 2 << 20); + } + let prepared = nodes[0] + .prepare_checkpoint(10, Uuid::from_u128(900))? + .run()?; + let image = prepared.image(); + nodes[0].install_checkpoint(prepared)?; + // A captured checkpoint must reserve exactly the same encoder lengths. + let g = &nodes[0].groups[&10]; + let cut = checkpoint::Cut { + group: 10, + id: Uuid::from_u128(901), + index: g.state.applied, + term: g.raw.store().term(g.state.applied)?, + }; + assert_eq!( + checkpoint::build_sizes(&layout(), &limits(), &cut, &g.state)?, + [image.bytes(Part::Metadata), image.bytes(Part::Payload)] + ); + + let before = nodes[2].disk_budget().unwrap().accounted_disk_bytes; + let mut receiver = + nodes[2].receive_snapshot(packet(&nodes[0], 3, &image)?, Uuid::from_u128(902))?; + let expected_charge = nodes[2].disk_budget().unwrap().reserved_file_bytes; + assert!(expected_charge >= image.bytes(Part::Payload)); + receiver.write(Part::Payload, 0, &image.chunk(Part::Payload, 0, 128 << 10)?)?; + drop(receiver); + assert_eq!(nodes[2].disk_budget().unwrap().reserved_file_bytes, 0); + assert!(nodes[2].disk_budget().unwrap().accounted_disk_bytes >= before + expected_charge); + + // The measurement task owns the engine lock even after its Node is gone. + let measure = nodes[2].prepare_disk_measurement()?; + drop(nodes.pop()); + assert!(open(root.path(), 3).is_err()); + assert!(measure.run()?.accounted_disk_bytes >= 2 << 20); + drop(measure); + nodes.push(open(root.path(), 3)?); + // Native recovery may truncate unused preallocated log tails. Validate the + // current physical files, not monotonic disk usage across engine reopens. + let mut occupied = 0; + for entry in std::fs::read_dir(root.path().join("3"))? { + let metadata = entry?.metadata()?; + if !metadata.is_file() { + continue; + } + let mut bytes = metadata.len(); + #[cfg(unix)] + { + use std::os::unix::fs::MetadataExt; + bytes = bytes.max(metadata.blocks() * 512); + } + occupied += bytes; + } + assert_eq!( + nodes[2].disk_budget().unwrap().accounted_disk_bytes, + occupied + ); + let orphan_usage = nodes[2].disk_budget().unwrap().accounted_disk_bytes; + let mut collection = nodes[2].prepare_image_collection()?; + let mut removed = 0; + loop { + let p = collection.run_batch(2)?; + removed += p.removed_files; + if p.finished { + break; + } + } + drop(collection); + assert_eq!(removed, 2); + assert!(nodes[2].disk_budget().unwrap().accounted_disk_bytes < orphan_usage); + + // The same physical budget rejects a new image before either file exists. + let budget = nodes[1].images.store().disk_budget().unwrap(); + let used = budget.snapshot().accounted_disk_bytes; + budget.limit_files(usize::try_from(used + image.bytes(Part::Payload) / 2)?)?; + let files = std::fs::read_dir(root.path().join("2"))?.count(); + assert!(nodes[1] + .receive_snapshot(packet(&nodes[0], 2, &image)?, Uuid::from_u128(903)) + .is_err()); + assert_eq!(std::fs::read_dir(root.path().join("2"))?.count(), files); + + nodes[0].prepare_log_maintenance()?.run()?; + let shared = nodes[0].images.store().disk_budget().unwrap().clone(); + let used = shared.snapshot().accounted_disk_bytes; + shared.limit_data(usize::try_from(used + 4096)?)?; + assert!(nodes[0].propose(20, 999, payload(20, 9)).is_err()); + let acquire = publication::Acquire { + table: Uuid::from_u128(100), + owner: Uuid::from_u128(101), + request: Uuid::from_u128(102), + expected_epoch: 0, + destination: publication::Destination { + namespace: Uuid::from_u128(103), + dataset: "test/table".into(), + }, + ingest: crate::ingest_contract::IngestContract { + tenant: "test".into(), + qualified_table: "table".into(), + schema: [1; 32], + policy: crate::ingest_contract::WritePolicy::ImmutableUniqueIds, + max_record_bytes: 1 << 20, + max_rows: 1024, + }, + }; + nodes[0].acquire_publisher(1000, acquire)?; + assert!(pump(&mut nodes)?.iter().any(|e| e.request == 1000 + && matches!( + e.outcome, + Outcome::Committed(Applied::Publisher(publication::Reply::Acquired(_))) + ))); + drop(image); + drop(nodes); + let mut nodes = (1..=3) + .map(|id| open(root.path(), id)) + .collect::>>()?; + for node in &nodes { + for &(group, sequence, index) in &accepted { + assert_eq!( + node.read_accepted(group, index)?, + Some(payload(group, sequence)) + ); + } + } + nodes[0].campaign(10)?; + nodes[0].campaign(20)?; + pump(&mut nodes)?; + for group in [10, 20] { + nodes[0].propose(group, 2000 + group, payload(group, 8))?; + let expected = accepted + .iter() + .find(|(g, s, _)| *g == group && *s == 8) + .unwrap() + .2; + assert!(pump(&mut nodes)?.iter().any(|e| e.request == 2000 + group + && e.outcome == Outcome::Committed(Applied::Accepted { index: expected }))); + } + Ok(()) +} diff --git a/docs/design/wal-node-disk-admission.md b/docs/design/wal-node-disk-admission.md new file mode 100644 index 00000000..cbac6755 --- /dev/null +++ b/docs/design/wal-node-disk-admission.md @@ -0,0 +1,58 @@ +# Shared node disk admission + +`Node::with_disk_budget(DiskOptions)` binds one admission policy to the physical +engine before its group owners or network services start. All hosted groups, +local image builders, incoming snapshots, orphan files and native maintenance +share that policy. Reopening measures the original directory, including partial +files from an interrupted receiver. Policy is local and does not change the +replicated layout or application identity. + +`data_bytes` stops new data proposals. Their reservation survives until durable +application or an uncertain completion. Control proposals can use the remaining +`max_bytes` window. Every encoded engine write contributes to a shared monotonic +counter, including follower replication, HardState and sync markers. Successful +admission takes a small accounting lock and does not walk the directory or +decode payloads. `prepare_disk_measurement` returns a task for a blocking worker; +it retains the original engine lock across runtime shutdown. + +New image files preserve `headroom_bytes` below the total ceiling. Configuration +requires at least all hosted uncommitted windows plus two native target files. +The local builder counts bounded metadata through its actual encoder and sums +opaque payload lengths before creating either file. A receiver uses its +validated descriptor's exact lengths. Each length is rounded to 4 KiB for +reservation. Pending file work retains its charge across cancellation, then +converts it to counted growth after all buffered writes finish. A later scan may +replace this conservative charge with actual retained lengths/allocation; only +actual deletion returns the space used by an orphan or installed file. + +Physical maintenance reserves an additional current-directory-size window +before native rewriting. Rejection occurs before mutation and leaves the node +healthy. A successful native pass refreshes measured usage while appends are +excluded. Image collection instead holds a separate filesystem lifecycle guard: +ordinary log appends can continue during its bounded deletion batches. A fully +drained image scan refreshes accounting after deletion. Neither lock substitutes +for the image owners that protect installed snapshots and captured readers. + +The snapshot transport requires a node with this policy configured. Its optional +`Admission` callback can add service-specific restrictions; omitting that extra +policy cannot bypass the node's shared reservation. The real three-process TLS +recovery test uses only the built-in policy, including a killed receiver after a +payload chunk, catch-up, leader loss and full restart from original files. + +This is admission with reserved headroom, **not a hard filesystem quota**. +Accepted peer traffic must still persist; followers do not yet advertise disk +credits to leaders. Filesystem allocation and native rewrite peaks can exceed +logical sizes or estimates. Disk space can also be consumed outside this owned +directory. Exhaustion or uncertain persistence still fails closed and requires +recovery. Deployment must provide enough physical space, monitor the remaining +window and schedule measurements/checkpointing/collection. Automatic scheduling, +peer pressure coordination and complete ingestion/materializer integration remain +required before claiming bounded sustained production capacity. + +Tests cover shared data/control limits, rejection before receive-file creation, +partial files through reopen and collection, all payloads and exact retries after +engine restart, and a real append completing while the image deletion lifecycle +is held. These establish the stated accounting and recovery boundaries; they do +not establish a throughput gain or resilience to losing a durable majority. + +Tracking: #378. diff --git a/docs/design/wal-runtime-images.md b/docs/design/wal-runtime-images.md index a4033124..196b3aba 100644 --- a/docs/design/wal-runtime-images.md +++ b/docs/design/wal-runtime-images.md @@ -43,8 +43,9 @@ images keep the native engine lock alive after runtime shutdown. This is the runtime integration boundary. It does not schedule automatic maintenance or impose a disk quota. The optional `peer::snapshot::Transport` -provides separate authenticated image streams; attach it and supply physical -disk admission before enabling compaction. See `wal-snapshot-transport.md`. +provides separate authenticated image streams and requires configured shared +node disk admission; attach it before enabling compaction. See +`wal-snapshot-transport.md` and `wal-node-disk-admission.md`. Tests use real files and owner threads with a test-only image relay. They cover two compacted groups catching up a lagging follower, every original engine diff --git a/docs/design/wal-snapshot-transport.md b/docs/design/wal-snapshot-transport.md index 7369877f..e1ed49f4 100644 --- a/docs/design/wal-snapshot-transport.md +++ b/docs/design/wal-snapshot-transport.md @@ -22,21 +22,23 @@ already running filesystem operation. Dropping a transport requests shutdown; call `shutdown().await` before reopening storage. A stalled image connection does not occupy an ordinary Raft connection. -The embedding service must supply `Admission`; there is no permissive default. -Its accounting must cover existing installed images, orphaned/partial files, -concurrent log growth and maintenance headroom, and recover after restart. A -reservation owner spans pending file work. Releasing the owner must not release -charges for files that still exist. Collection, rather than network completion, -reclaims those charges. `max_image_bytes` and connection counts do not impose a -physical node disk quota. The shared production disk-admission implementation, -automatic maintenance and public ingestion integration remain follow-up work. +The node must configure `with_disk_budget` before starting this transport. Its +shared policy covers installed images, orphaned/partial files, encoded log +growth and reserved maintenance headroom. A reservation spans pending file work; +completion or cancellation does not return space still used by files. An optional +`Admission` callback can impose additional service restrictions. Neither that +callback nor `max_image_bytes` replaces shared accounting. See +`wal-node-disk-admission.md` for its limits: this is not a hard filesystem quota, +and automatic maintenance, peer disk credits and public ingestion integration +remain follow-up work. Tests exercise actual mutual TLS and files, including two compacted groups, reopening all original engines with complete payload and exact retries, stalled snapshot isolation from healthy quorum commits, admission rejection before files, interrupted reception without successful acknowledgment, and retained -orphan charges. The test admission provider is a conservative image-only budget; -it does not implement the missing production-wide disk policy. A three-process qualification also pauses after a real payload chunk is written, +orphan charges. Additional rejection tests use a conservative image-only policy +on top of the node budget. A three-process qualification uses the built-in shared +budget alone and pauses after a real payload chunk is written, kills that receiver, rejoins its original files and endpoints, kills the leader, and restarts all original engines. Every payload and exact retry index must survive. The parent carries only test control requests; TLS carries both Raft diff --git a/scripts/ci/install-manylinux-deps.sh b/scripts/ci/install-manylinux-deps.sh new file mode 100644 index 00000000..3bfce4e6 --- /dev/null +++ b/scripts/ci/install-manylinux-deps.sh @@ -0,0 +1,27 @@ +#!/usr/bin/env bash +set -euo pipefail + +if command -v yum >/dev/null 2>&1; then + yum_options=(--setopt=timeout=30 --setopt=retries=2) + if test -f /etc/almalinux-release; then + # These build dependencies are all in BaseOS/AppStream. Avoid querying EPEL + # and third-party mirrorlists; keep the distribution's GPG and TLS checks. + yum_options+=( + '--disablerepo=*' + --enablerepo=baseos + --enablerepo=appstream + '--setopt=baseos.mirrorlist=' + '--setopt=appstream.mirrorlist=' + '--setopt=baseos.baseurl=https://repo.almalinux.org/almalinux/$releasever/BaseOS/$basearch/os/' + '--setopt=appstream.baseurl=https://repo.almalinux.org/almalinux/$releasever/AppStream/$basearch/os/' + ) + fi + yum "${yum_options[@]}" install -y curl make openssl-devel perl +elif command -v apt-get >/dev/null 2>&1; then + apt-get -o Acquire::Retries=2 -o Acquire::http::Timeout=30 -o Acquire::https::Timeout=30 update + apt-get -o Acquire::Retries=2 -o Acquire::http::Timeout=30 -o Acquire::https::Timeout=30 \ + install -y curl libssl-dev ca-certificates make perl +else + echo 'No supported package manager for manylinux build dependencies' >&2 + exit 1 +fi