fix: reclaim client journal rewrite bytes before fencing

crates/lance-context-wal/Cargo.toml

@@ -9,7 +9,7 @@ description = "Independent Rust WAL experiments owned by lance-context"
 [features]
 # Real Lance materialization/recovery qualification; normal WAL replicas do not link Lance.
 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", "dep:bytes", "dep:serde", "dep:serde_json", "dep:tokio", "tokio/sync"]
+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"]
 vectored-io = ["dep:bytes", "dep:tokio"]

crates/lance-context-wal/src/client_journal.rs

@@ -88,6 +88,7 @@ struct Counters {
     write_ns: AtomicU64,
     maintenance_calls: AtomicU64,
     maintenance_ns: AtomicU64,
+    pressure_rewrites: AtomicU64,
     queue_ns: AtomicU64,
     fenced: AtomicBool,
 }
@@ -100,6 +101,7 @@ pub struct Statistics {
     pub write_seconds: f64,
     pub maintenance_calls: u64,
     pub maintenance_seconds: f64,
+    pub pressure_rewrites: u64,
     pub queue_seconds: f64,
     pub fenced: bool,
 }
@@ -292,6 +294,7 @@ impl Journal {
             write_seconds: get(&c.write_ns) as f64 / 1e9,
             maintenance_calls: get(&c.maintenance_calls),
             maintenance_seconds: get(&c.maintenance_ns) as f64 / 1e9,
+            pressure_rewrites: get(&c.pressure_rewrites),
             queue_seconds: get(&c.queue_ns) as f64 / 1e9,
             fenced: c.fenced.load(Ordering::Acquire),
         }
@@ -353,9 +356,24 @@ fn worker<F: FileSystem>(
                     .fetch_add(nanos(start), Ordering::Relaxed);
                 counters.maintenance_calls.fetch_add(1, Ordering::Relaxed);
             }
+            let ceiling = settings.log_bytes - settings.reserve();
+            if footprint(Path::new(engine.path()))? + bytes as u64 > ceiling {
+                // The engine's ordinary rewrite heuristic counts stale keys,
+                // not their bytes. A few replaced payloads can fill the budget
+                // while many tiny cold values keep that ratio below its gate.
+                // Spend the already reserved rewrite headroom once, then check
+                // actual physical bytes again. Never raise the budget or retry
+                // indefinitely; any failed maintenance fences this worker.
+                let start = Instant::now();
+                engine.purge_manager().reclaim_rewrite_queue()?;
+                counters
+                    .maintenance_ns
+                    .fetch_add(nanos(start), Ordering::Relaxed);
+                counters.maintenance_calls.fetch_add(1, Ordering::Relaxed);
+                counters.pressure_rewrites.fetch_add(1, Ordering::Relaxed);
+            }
             ensure!(
-                footprint(Path::new(engine.path()))? + bytes as u64
-                    <= settings.log_bytes - settings.reserve(),
+                footprint(Path::new(engine.path()))? + bytes as u64 <= ceiling,
                 "client journal disk admission closed"
             );
             let mut total: usize = lengths.iter().sum();

crates/lance-context-wal/src/client_journal/tests.rs

@@ -2,7 +2,9 @@ use super::*;
 use raft_engine::env::{DefaultFileSystem, Handle, Permission};
 use std::io;
 use std::sync::{mpsc as blocking, Mutex};
+mod pressure_faults;
 mod rewrite;
+mod skew;
 
 fn settings() -> Settings {
     Settings {
@@ -91,6 +93,10 @@ type RealHandle = <DefaultFileSystem as FileSystem>::Handle;
 #[derive(Default)]
 struct Control {
     fail: AtomicBool,
+    fail_rewrite: AtomicBool,
+    rewrite_failures: AtomicU64,
+    rewrite_skip_syncs: AtomicU64,
+    rewrite_sync_calls: AtomicU64,
     calls: AtomicU64,
     entered: Mutex<Option<blocking::SyncSender<()>>>,
     release: Mutex<Option<blocking::Receiver<()>>>,
@@ -110,16 +116,44 @@ impl Handle for FaultHandle {
     }
     fn sync(&self) -> io::Result<()> {
         if self.rewrite {
+            let call = self
+                .control
+                .rewrite_sync_calls
+                .fetch_add(1, Ordering::Relaxed)
+                + 1;
+            let mut skip = self.control.rewrite_skip_syncs.load(Ordering::Acquire);
+            while skip > 0 {
+                match self.control.rewrite_skip_syncs.compare_exchange_weak(
+                    skip,
+                    skip - 1,
+                    Ordering::AcqRel,
+                    Ordering::Acquire,
+                ) {
+                    Ok(_) => return self.inner.sync(),
+                    Err(actual) => skip = actual,
+                }
+            }
             let pause = self.control.rewrite_pause.lock().unwrap().take();
             if let Some((marker, after_sync)) = pause {
                 if after_sync {
                     self.inner.sync()?;
                 }
-                std::fs::write(marker, b"actual rewrite sync reached")?;
+                std::fs::write(
+                    marker,
+                    serde_json::to_vec(
+                        &serde_json::json!({"rewrite_sync":call,"file_bytes":self.inner.file_size()?}),
+                    )?,
+                )?;
                 loop {
                     std::thread::park();
                 }
             }
+            if self.control.fail_rewrite.load(Ordering::Acquire) {
+                self.control
+                    .rewrite_failures
+                    .fetch_add(1, Ordering::Relaxed);
+                return Err(io::Error::other("injected pressure rewrite sync failure"));
+            }
         }
         self.control.calls.fetch_add(1, Ordering::Relaxed);
         if let Some(tx) = self.control.entered.lock().unwrap().take() {

crates/lance-context-wal/src/client_journal/tests/pressure_faults.rs

@@ -0,0 +1,203 @@
+//! Exercise failures in the byte-pressure fallback, on the same live engine.
+use super::*;
+use std::process::{Child, Command, Stdio};
+
+const SCOPE: &[u8] = b"pressure-faults";
+type Seed = (Engine<FaultFs>, Vec<Bytes>, u32, Bytes);
+fn config() -> Settings {
+    Settings {
+        linger_ms: 0,
+        ..Settings::default()
+    }
+}
+fn put(engine: &Engine<FaultFs>, key: u32, value: &Bytes) -> Result<()> {
+    let mut batch = LogBatch::default();
+    batch.put(GROUP, key.to_be_bytes().to_vec(), value.to_vec())?;
+    engine.write(&mut batch, true)?;
+    Ok(())
+}
+
+fn seed(path: &Path, fs: Arc<FaultFs>) -> Result<Seed> {
+    let settings = config();
+    let (journal, _) = Journal::open(path, SCOPE, 128, settings.clone())?;
+    drop(journal);
+    let engine = Engine::open_with(settings.engine(path), fs, vec![])?;
+    let mut latest: Vec<_> = (0..128u32)
+        .map(|key| Bytes::from(vec![key as u8; 128]))
+        .collect();
+    for (key, value) in latest.iter().enumerate() {
+        put(&engine, key as u32, value)?;
+    }
+    for round in 0..40u8 {
+        for size in [4 << 20, 128] {
+            for key in 0..8u32 {
+                let mut value = vec![round; size];
+                value[0] = key as u8;
+                let value = Bytes::from(value);
+                if footprint(path)? >= 2 * settings.file_bytes as u64 {
+                    engine.purge_expired_files()?;
+                }
+                if footprint(path)? + (size + OVERHEAD) as u64
+                    > settings.log_bytes - settings.reserve()
+                {
+                    return Ok((engine, latest, key, value));
+                }
+                put(&engine, key, &value)?;
+                latest[key as usize] = value;
+            }
+        }
+    }
+    anyhow::bail!("fixture did not reach the byte-pressure fallback")
+}
+
+fn retained_dir(prefix: &str) -> Result<tempfile::TempDir> {
+    if let Some(base) = std::env::var_os("CONTEXT_CLIENT_JOURNAL_TEST_ARTIFACTS") {
+        std::fs::create_dir_all(&base)?;
+        let mut dir = tempfile::Builder::new().prefix(prefix).tempdir_in(base)?;
+        dir.disable_cleanup(true);
+        Ok(dir)
+    } else {
+        Ok(tempfile::tempdir()?)
+    }
+}
+
+#[tokio::test]
+async fn pressure_sync_error_fences_before_ack_and_preserves_all_prior_values() -> Result<()> {
+    let dir = retained_dir("pressure-error-")?;
+    let path = dir.path().join("journal");
+    let fs = Arc::new(FaultFs::default());
+    let control = fs.control.clone();
+    let (engine, expected, key, next) = seed(&path, fs)?;
+    let (journal, recovered) = Journal::with_engine(engine, SCOPE, 128, config())?;
+    assert_eq!(
+        recovered,
+        expected.iter().cloned().map(Some).collect::<Vec<_>>()
+    );
+    // Rotation first syncs the old file and the new header. Inject at the
+    // third sync, after the live values have actually been rewritten.
+    control.rewrite_sync_calls.store(0, Ordering::Release);
+    control.rewrite_skip_syncs.store(2, Ordering::Release);
+    control.fail_rewrite.store(true, Ordering::Release);
+    let error = journal.put(key, Bytes::new(), next).await.unwrap_err();
+    // The pinned engine panics on sync failure and syncs again during Drop.
+    // The journal catches that first panic and remains fenced. Clear only the
+    // injected filesystem failure before any assertion can unwind; never reuse
+    // the fenced journal. Separate SIGKILL tests cover recovery without Drop.
+    control.fail_rewrite.store(false, Ordering::Release);
+    assert!(
+        error
+            .to_string()
+            .contains("client journal mutation panicked; reopen required"),
+        "{error:#}"
+    );
+    assert_eq!(control.rewrite_failures.load(Ordering::Acquire), 1);
+    assert_eq!(control.rewrite_sync_calls.load(Ordering::Acquire), 3);
+    assert!(journal.statistics().fenced);
+    assert!(journal
+        .put(0, Bytes::new(), Bytes::from_static(b"later"))
+        .await
+        .is_err());
+    assert!(Journal::open(&path, SCOPE, 128, config()).is_err());
+    drop(journal);
+    let (_, recovered) = Journal::open(&path, SCOPE, 128, config())?;
+    assert_eq!(
+        recovered,
+        expected.into_iter().map(Some).collect::<Vec<_>>()
+    );
+    Ok(())
+}
+
+#[test]
+#[ignore = "invoked by pressure SIGKILL parent only"]
+fn pressure_child() -> Result<()> {
+    let root = std::path::PathBuf::from(std::env::var("CONTEXT_PRESSURE_ROOT")?);
+    let after = std::env::var("CONTEXT_PRESSURE_AFTER")? == "1";
+    let fs = Arc::new(FaultFs::default());
+    let control = fs.control.clone();
+    let (engine, expected, key, next) = seed(&root.join("journal"), fs)?;
+    // Every fixture value is regenerated byte-for-byte from this compact shape.
+    let shape: Vec<_> = expected.iter().map(|v| (v[0], v[1], v.len())).collect();
+    std::fs::write(root.join("expected.json"), serde_json::to_vec(&shape)?)?;
+    std::fs::File::open(root.join("expected.json"))?.sync_all()?;
+    std::fs::File::open(&root)?.sync_all()?;
+    let (journal, _) = Journal::with_engine(engine, SCOPE, 128, config())?;
+    control.rewrite_sync_calls.store(0, Ordering::Release);
+    control.rewrite_skip_syncs.store(2, Ordering::Release);
+    *control.rewrite_pause.lock().unwrap() = Some((root.join("paused"), after));
+    tokio::runtime::Runtime::new()?.block_on(journal.put(key, Bytes::new(), next))?;
+    anyhow::bail!("pressure request completed without reaching rewrite sync")
+}
+
+struct Stop(Child);
+impl Drop for Stop {
+    fn drop(&mut self) {
+        let _ = self.0.kill();
+        let _ = self.0.wait();
+    }
+}
+
+#[test]
+fn pressure_sigkill_before_and_after_sync_preserves_all_acknowledged_values() -> Result<()> {
+    for after in [false, true] {
+        let root = retained_dir("pressure-kill-")?;
+        let mut child = Stop(
+            Command::new(std::env::current_exe()?)
+                .args([
+                    "--ignored",
+                    "--exact",
+                    "client_journal::tests::pressure_faults::pressure_child",
+                    "--nocapture",
+                ])
+                .env("CONTEXT_PRESSURE_ROOT", root.path())
+                .env("CONTEXT_PRESSURE_AFTER", if after { "1" } else { "0" })
+                .stdin(Stdio::null())
+                .stdout(std::fs::File::create(root.path().join("stdout"))?)
+                .stderr(std::fs::File::create(root.path().join("stderr"))?)
+                .spawn()?,
+        );
+        let until = Instant::now() + Duration::from_secs(120);
+        while !root.path().join("paused").exists() {
+            ensure!(
+                child.0.try_wait()?.is_none(),
+                "pressure child exited: {}",
+                root.path().display()
+            );
+            ensure!(Instant::now() < until, "pressure child did not reach sync");
+            std::thread::sleep(Duration::from_millis(20));
+        }
+        let boundary: serde_json::Value =
+            serde_json::from_slice(&std::fs::read(root.path().join("paused"))?)?;
+        assert_eq!(boundary["rewrite_sync"], 3);
+        ensure!(
+            boundary["file_bytes"].as_u64().unwrap() > 1 << 20,
+            "paused before rewritten payloads existed"
+        );
+        ensure!(
+            footprint(&root.path().join("journal"))? <= config().log_bytes,
+            "rewrite exceeded reserved disk budget"
+        );
+        child.0.kill()?;
+        let status = child.0.wait()?;
+        #[cfg(unix)]
+        {
+            use std::os::unix::process::ExitStatusExt;
+            assert_eq!(status.signal(), Some(9));
+        }
+        let shape: Vec<(u8, u8, usize)> =
+            serde_json::from_slice(&std::fs::read(root.path().join("expected.json"))?)?;
+        let (_, recovered) = Journal::open(&root.path().join("journal"), SCOPE, 128, config())?;
+        assert_eq!(recovered.len(), shape.len());
+        for (value, (first, rest, length)) in recovered.into_iter().zip(shape) {
+            let mut expected = vec![rest; length];
+            expected[0] = first;
+            assert_eq!(value.unwrap().as_ref(), expected.as_slice());
+        }
+        std::fs::write(
+            root.path().join("verified.json"),
+            serde_json::to_vec(&serde_json::json!({
+                "after_sync":after,"signal":9,"all_128_latest_values_verified":true
+            }))?,
+        )?;
+    }
+    Ok(())
+}