From d26139417ef6d9a3f0b35565e0a4ab2fa6d238d4 Mon Sep 17 00:00:00 2001 From: Daniel Rammer Date: Mon, 27 Jul 2026 19:38:33 -0500 Subject: [PATCH] feat(table): checkpoint_lsm, flush_lsm, compact_lsm, get_lsm_stats MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Converge a table's LSM write path into its base table, and inspect it. `checkpoint_lsm` is `flush` then `compact`, repeated until the fresh tier is empty — and the loop runs **client-side**. Putting it on the server would mean a background task, which means a single-flight intent, an intent that leaks on panic, a bounded-iteration policy, an "is it done" observable, and a story for every way a client can vanish mid-operation. None of that exists in this shape: each request does a bounded unit of work and reports what is left, so completion is *carried in the responses* rather than inferred from a shared counter that cannot distinguish "converged" from "hasn't started yet". Best-effort by construction. Nothing is frozen, so `converged` means L0 was empty as of the last pass. It is idempotent, abandonable at any point with zero consequence, and safe to run on a cadence — an already- converged table costs one round trip and zero compaction passes, because `flush` reports `generations_remaining` and the loop is never entered. The failure taxonomy is the load-bearing part. Five distinct conditions used to arrive at a client as one 503, so `Error::LsmRoute` carries a classification read from the response body's namespace error code *at the point of receipt* — before any generic helper folds the body into a string and keeps only the status. Contention (429) retries; a draining node (503 / `InvalidTableState`) stops, because the drain gate is a one-way latch and retrying it spins until the deadline to report a failure that was knowable on the first response. Transport retry is disabled on these routes for the same reason: it treats every 503 alike. `get_lsm_stats` returns `Option`, matching `get_lsm_write_spec` — `None` only when the table has no LSM write path, since a struct of zeros would read as measurements. Python bindings mirror all four, preserving the per-bucket detail rather than flattening to a table-level summary. Co-Authored-By: Claude Opus 5 (1M context) --- python/python/lancedb/_lancedb.pyi | 9 + python/python/lancedb/table.py | 115 ++++++++ python/src/table.rs | 223 ++++++++++++++- rust/lancedb/src/error.rs | 18 ++ rust/lancedb/src/remote/table.rs | 328 +++++++++++++++++++++ rust/lancedb/src/table.rs | 187 ++++++++++++ rust/lancedb/src/table/checkpoint.rs | 409 +++++++++++++++++++++++++++ 7 files changed, 1288 insertions(+), 1 deletion(-) create mode 100644 rust/lancedb/src/table/checkpoint.rs diff --git a/python/python/lancedb/_lancedb.pyi b/python/python/lancedb/_lancedb.pyi index c922aa69b..5f8fc928c 100644 --- a/python/python/lancedb/_lancedb.pyi +++ b/python/python/lancedb/_lancedb.pyi @@ -284,6 +284,15 @@ class Table: async def set_lsm_write_spec(self, spec: LsmWriteSpec) -> None: ... async def unset_lsm_write_spec(self) -> None: ... async def get_lsm_write_spec(self) -> Optional[LsmWriteSpec]: ... + async def checkpoint_lsm( + self, + deadline_secs: Optional[float], + max_generations_per_bucket: Optional[int], + include_stats: bool, + ) -> dict: ... + async def flush_lsm(self) -> dict: ... + async def compact_lsm(self, max_generations_per_bucket: Optional[int]) -> dict: ... + async def get_lsm_stats(self, include_generation_rows: bool) -> Optional[dict]: ... async def close_lsm_writers(self) -> None: ... @property def tags(self) -> Tags: ... diff --git a/python/python/lancedb/table.py b/python/python/lancedb/table.py index 5f89f6308..32cad9f2f 100644 --- a/python/python/lancedb/table.py +++ b/python/python/lancedb/table.py @@ -3872,6 +3872,44 @@ class LanceTable(Table): [`AsyncTable.get_lsm_write_spec`][lancedb.AsyncTable.get_lsm_write_spec].""" return LOOP.run(self._table.get_lsm_write_spec()) + def checkpoint_lsm( + self, + *, + deadline_secs: Optional[float] = 300.0, + max_generations_per_bucket: Optional[int] = None, + include_stats: bool = False, + ) -> dict: + """Synchronous version of + [`AsyncTable.checkpoint_lsm`][lancedb.AsyncTable.checkpoint_lsm].""" + return LOOP.run( + self._table.checkpoint_lsm( + deadline_secs=deadline_secs, + max_generations_per_bucket=max_generations_per_bucket, + include_stats=include_stats, + ) + ) + + def flush_lsm(self) -> dict: + """Synchronous version of + [`AsyncTable.flush_lsm`][lancedb.AsyncTable.flush_lsm].""" + return LOOP.run(self._table.flush_lsm()) + + def compact_lsm(self, *, max_generations_per_bucket: Optional[int] = None) -> dict: + """Synchronous version of + [`AsyncTable.compact_lsm`][lancedb.AsyncTable.compact_lsm].""" + return LOOP.run( + self._table.compact_lsm( + max_generations_per_bucket=max_generations_per_bucket + ) + ) + + def get_lsm_stats(self, *, include_generation_rows: bool = False) -> Optional[dict]: + """Synchronous version of + [`AsyncTable.get_lsm_stats`][lancedb.AsyncTable.get_lsm_stats].""" + return LOOP.run( + self._table.get_lsm_stats(include_generation_rows=include_generation_rows) + ) + def close_lsm_writers(self) -> None: """Close cached MemWAL shard writers. See [`AsyncTable.close_lsm_writers`][lancedb.AsyncTable.close_lsm_writers].""" @@ -4582,6 +4620,83 @@ class AsyncTable: """ return await self._inner.get_lsm_write_spec() + async def checkpoint_lsm( + self, + *, + deadline_secs: Optional[float] = 300.0, + max_generations_per_bucket: Optional[int] = None, + include_stats: bool = False, + ) -> dict: + """Converge this table's LSM write path into its base table. + + One flush (sealing every memtable into L0), then bounded compaction + passes until L0 is empty. The loop runs client-side, so each request + does a bounded unit of work and reports what is left. + + Best-effort: nothing is frozen, so with writes flowing there may be + new rows by the time it returns — ``converged`` means L0 was empty as + of the last pass. Idempotent and safe on a cadence; an + already-converged table costs one round trip and zero compaction + passes. + + Parameters + ---------- + deadline_secs + Stop and report ``converged=False`` after this much wall clock. + ``None`` runs until the table converges — safe only when you know + it is not under write load. + max_generations_per_bucket + Bound on generations merged **per bucket** per request. An + N-bucket table admits ``N * max_generations_per_bucket`` merges + per call. + include_stats + Attach a ``get_lsm_stats`` block to the report. + """ + return await self._inner.checkpoint_lsm( + deadline_secs, max_generations_per_bucket, include_stats + ) + + async def flush_lsm(self) -> dict: + """Seal every bucket's active memtable into L0. + + Does not touch the base table — moving L0 into base is + `compact_lsm`. On a node that has not claimed this table, this claims + it and replays its WAL log first. + """ + return await self._inner.flush_lsm() + + async def compact_lsm( + self, *, max_generations_per_bucket: Optional[int] = None + ) -> dict: + """Run one bounded L0 to base compaction pass per bucket. + + One pass, not convergence: that is what bounds each request's cost + and gives a caller driving its own cadence a progress signal per + round trip. Use `checkpoint_lsm` to loop to convergence. + """ + return await self._inner.compact_lsm(max_generations_per_bucket) + + async def get_lsm_stats( + self, *, include_generation_rows: bool = False + ) -> Optional[dict]: + """Read live LSM state: per-bucket detail plus a table-level aggregate. + + Answers "how far behind is my fresh tier", "which bucket is hot", and + "why is my fresh-tier vector search brute-force". Mutates no table + state, though on a node that has not claimed this table it claims it, + exactly as a read would. + + Returns ``None`` — and only — when the LSM write path is not enabled + for the table. + + Parameters + ---------- + include_generation_rows + Report row counts per L0 generation. Off by default: each costs + one Lance manifest read. + """ + return await self._inner.get_lsm_stats(include_generation_rows) + async def close_lsm_writers(self) -> None: """Drain and close any cached MemWAL shard writers for this table. diff --git a/python/src/table.rs b/python/src/table.rs index 5d6910685..160e95705 100644 --- a/python/src/table.rs +++ b/python/src/table.rs @@ -28,11 +28,103 @@ use pyo3::{ Bound, FromPyObject, Py, PyAny, PyRef, PyResult, Python, exceptions::{PyRuntimeError, PyValueError}, pyclass, pyfunction, pymethods, - types::{IntoPyDict, PyAnyMethods, PyBytes, PyDict, PyDictMethods}, + types::{IntoPyDict, PyAnyMethods, PyBytes, PyDict, PyDictMethods, PyList, PyListMethods}, }; mod scannable; +/// Convert `LsmStats` to a Python dict, preserving the per-bucket list. +/// +/// Deliberately not flattened to a table-level summary: a table is N +/// buckets on one node, and the per-bucket detail is the reason the +/// endpoint exists — flattening hides the single hot bucket someone opened +/// it to find. +fn lsm_stats_to_py(py: Python<'_>, stats: &lancedb::table::LsmStats) -> PyResult> { + let out = PyDict::new(py); + let buckets = PyList::empty(py); + for b in &stats.buckets { + let e = PyDict::new(py); + e.set_item("shard_id", &b.shard_id)?; + e.set_item("status", &b.status)?; + e.set_item("writer_epoch", b.writer_epoch)?; + e.set_item("manifest_version", b.manifest_version)?; + e.set_item("current_generation", b.current_generation)?; + e.set_item( + "replay_after_wal_entry_position", + b.replay_after_wal_entry_position, + )?; + e.set_item( + "wal_entry_position_last_seen", + b.wal_entry_position_last_seen, + )?; + e.set_item("wal_lag", b.wal_lag)?; + e.set_item("l0_bytes", b.l0_bytes)?; + e.set_item("fenced", b.fenced)?; + e.set_item("compaction_in_progress", b.compaction_in_progress)?; + + let generations = PyList::empty(py); + for g in &b.generations { + let ge = PyDict::new(py); + ge.set_item("generation", g.generation)?; + ge.set_item("path", &g.path)?; + ge.set_item("bytes", g.bytes)?; + ge.set_item("rows", g.rows)?; + generations.append(ge)?; + } + e.set_item("generations", generations)?; + + let memtable = |m: &lancedb::table::MemtableStats| -> PyResult> { + let d = PyDict::new(py); + d.set_item("generation", m.generation)?; + d.set_item("rows", m.rows)?; + d.set_item("bytes", m.bytes)?; + d.set_item("batches", m.batches)?; + Ok(d.unbind()) + }; + e.set_item( + "active_memtable", + b.active_memtable.as_ref().map(memtable).transpose()?, + )?; + e.set_item( + "frozen_memtables", + b.frozen_memtables + .as_ref() + .map(|fs| { + let l = PyList::empty(py); + for m in fs { + l.append(memtable(m)?)?; + } + PyResult::Ok(l.unbind()) + }) + .transpose()?, + )?; + e.set_item( + "memtable_indexes", + b.memtable_indexes + .as_ref() + .map(|ixs| { + let l = PyList::empty(py); + for i in ixs { + let d = PyDict::new(py); + d.set_item("name", &i.name)?; + d.set_item("kind", &i.kind)?; + d.set_item("column", &i.column)?; + l.append(d)?; + } + PyResult::Ok(l.unbind()) + }) + .transpose()?, + )?; + buckets.append(e)?; + } + out.set_item("buckets", buckets)?; + out.set_item("bucket_count", stats.bucket_count)?; + out.set_item("generations_total", stats.generations_total)?; + out.set_item("l0_bytes_total", stats.l0_bytes_total)?; + out.set_item("memtable_rows_total", stats.memtable_rows_total)?; + Ok(out.unbind()) +} + #[derive(FromPyObject)] enum PredicateArg { Expr(PyExpr), @@ -1267,6 +1359,135 @@ impl Table { }) } + /// Converge the table's LSM write path into its base table. + /// + /// Returns a dict mirroring `CheckpointReport`. Best-effort: with + /// writes flowing, new rows may land after the last pass. + #[pyo3(signature = (deadline_secs=None, max_generations_per_bucket=None, include_stats=false))] + pub fn checkpoint_lsm( + self_: PyRef<'_, Self>, + deadline_secs: Option, + max_generations_per_bucket: Option, + include_stats: bool, + ) -> PyResult> { + let inner = self_.inner_ref()?.clone(); + let opts = lancedb::table::CheckpointOptions { + deadline: deadline_secs.map(std::time::Duration::from_secs_f64), + max_generations_per_bucket, + include_stats, + }; + future_into_py(self_.py(), async move { + let report = inner.checkpoint_lsm(opts).await.infer_error()?; + Python::attach(|py| { + let d = PyDict::new(py); + d.set_item("converged", report.converged)?; + d.set_item( + "stop_reason", + report.stop_reason.map(|r| match r { + lancedb::table::CheckpointStopReason::DeadlineExceeded => { + "deadline_exceeded" + } + lancedb::table::CheckpointStopReason::NodeDraining => "node_draining", + lancedb::table::CheckpointStopReason::RepeatedlyUnclaimed => { + "repeatedly_unclaimed" + } + }), + )?; + d.set_item("rows_sealed", report.rows_sealed)?; + d.set_item("generations_compacted", report.generations_compacted)?; + d.set_item("rows_merged", report.rows_merged)?; + d.set_item("generations_remaining", report.generations_remaining)?; + d.set_item("compact_calls", report.compact_calls)?; + d.set_item( + "stats", + report.stats.map(|s| lsm_stats_to_py(py, &s)).transpose()?, + )?; + Ok(d.unbind()) + }) + }) + } + + /// Seal every bucket's active memtable into L0. Returns a dict + /// mirroring `FlushReport`. + pub fn flush_lsm(self_: PyRef<'_, Self>) -> PyResult> { + let inner = self_.inner_ref()?.clone(); + future_into_py(self_.py(), async move { + let report = inner.flush_lsm().await.infer_error()?; + Python::attach(|py| { + let d = PyDict::new(py); + let buckets = PyList::empty(py); + // Per-bucket detail is preserved rather than flattened: a + // table is N buckets on one node, and the one hot bucket is + // usually why someone is looking. + for b in &report.buckets { + let e = PyDict::new(py); + e.set_item("shard_id", &b.shard_id)?; + e.set_item("sealed_generation", b.sealed_generation)?; + e.set_item("rows_sealed", b.rows_sealed)?; + e.set_item("generations_remaining", b.generations_remaining)?; + buckets.append(e)?; + } + d.set_item("buckets", buckets)?; + d.set_item("rows_sealed", report.rows_sealed)?; + d.set_item("generations_remaining", report.generations_remaining)?; + Ok(d.unbind()) + }) + }) + } + + /// Run one bounded L0 → base pass per bucket. Returns a dict mirroring + /// `CompactReport`. + #[pyo3(signature = (max_generations_per_bucket=None))] + pub fn compact_lsm( + self_: PyRef<'_, Self>, + max_generations_per_bucket: Option, + ) -> PyResult> { + let inner = self_.inner_ref()?.clone(); + let opts = lancedb::table::CompactOptions { + max_generations_per_bucket, + }; + future_into_py(self_.py(), async move { + let report = inner.compact_lsm(opts).await.infer_error()?; + Python::attach(|py| { + let d = PyDict::new(py); + let buckets = PyList::empty(py); + for b in &report.buckets { + let e = PyDict::new(py); + e.set_item("shard_id", &b.shard_id)?; + e.set_item("generations_consumed", b.generations_consumed)?; + e.set_item("rows_merged", b.rows_merged)?; + e.set_item("generations_remaining", b.generations_remaining)?; + buckets.append(e)?; + } + d.set_item("buckets", buckets)?; + d.set_item("generations_consumed", report.generations_consumed)?; + d.set_item("rows_merged", report.rows_merged)?; + d.set_item("generations_remaining", report.generations_remaining)?; + d.set_item( + "highest_wal_entry_position", + report.highest_wal_entry_position, + )?; + Ok(d.unbind()) + }) + }) + } + + /// Live LSM state, or `None` when the LSM write path is not enabled. + #[pyo3(signature = (include_generation_rows=false))] + pub fn get_lsm_stats( + self_: PyRef<'_, Self>, + include_generation_rows: bool, + ) -> PyResult> { + let inner = self_.inner_ref()?.clone(); + let opts = lancedb::table::LsmStatsOptions { + include_generation_rows, + }; + future_into_py(self_.py(), async move { + let stats = inner.get_lsm_stats(opts).await.infer_error()?; + Python::attach(|py| stats.map(|s| lsm_stats_to_py(py, &s)).transpose()) + }) + } + pub fn close_lsm_writers(self_: PyRef<'_, Self>) -> PyResult> { let inner = self_.inner_ref()?.clone(); future_into_py(self_.py(), async move { diff --git a/rust/lancedb/src/error.rs b/rust/lancedb/src/error.rs index 7bfbc6f2e..69c81b0c8 100644 --- a/rust/lancedb/src/error.rs +++ b/rust/lancedb/src/error.rs @@ -77,6 +77,24 @@ pub enum Error { source: Box, status_code: Option, }, + /// An LSM operator route (`flush_lsm` / `compact_lsm` / + /// `get_lsm_stats`) failed, carrying the classification the checkpoint + /// loop decides on. + /// + /// It is its own variant because the decision — retry, re-issue from + /// `flush`, or stop — is made from the *response body's* namespace + /// error code on a 503, and every generic HTTP path in this crate folds + /// the body into a string and keeps only the status. Classifying at the + /// point of receipt is what keeps a draining node distinguishable from + /// a fenced writer, a busy compactor, and a transport blip, all of + /// which are otherwise the same 503. + #[cfg(feature = "remote")] + #[snafu(display("LSM route error ({fault:?}, status={status}): {message}"))] + LsmRoute { + fault: crate::table::checkpoint::LsmFault, + status: u16, + message: String, + }, #[snafu(display("Arrow error: {source}"))] Arrow { source: ArrowError }, #[snafu(display("LanceDBError: not supported: {message}"))] diff --git a/rust/lancedb/src/remote/table.rs b/rust/lancedb/src/remote/table.rs index 92c4ad296..0e447d27e 100644 --- a/rust/lancedb/src/remote/table.rs +++ b/rust/lancedb/src/remote/table.rs @@ -26,10 +26,12 @@ use crate::table::MergeBranchResult; use crate::table::MergeResult; use crate::table::Tags; use crate::table::UpdateResult; +use crate::table::checkpoint::{GetLsmStatsResponse, classify_lsm_fault}; use crate::table::merge::MergeFilter; use crate::table::query::create_multi_vector_plan; use crate::table::{AlterColumnsResult, FieldMetadataUpdate, UpdateFieldMetadataResult}; use crate::table::{AnyQuery, Filter, Predicate, PreprocessingOutput, TableStatistics}; +use crate::table::{CompactOptions, CompactReport, FlushReport, LsmStats, LsmStatsOptions}; use crate::utils::background_cache::BackgroundCache; use crate::utils::{ resolve_arrow_field_path, supported_btree_data_type, supported_vector_data_type, @@ -907,6 +909,44 @@ impl RemoteTable { } } + /// Send an LSM operator request and decode its JSON body, classifying + /// a failure into an [`Error::LsmRoute`] the checkpoint loop can act on. + /// + /// Deliberately **not** routed through `check_response` / + /// `check_table_response`: those fold the body into a string and keep + /// only the status, and the body's `code` is the only thing that tells + /// a draining node apart from a fenced writer or a transport blip on a + /// 503. Classifying here — before any generic helper touches the + /// response — is what preserves a distinction threaded through three + /// crates and two error enums. + async fn send_lsm_route( + &self, + request: RequestBuilder, + route: &str, + ) -> Result { + // `with_retry = false` on purpose. The transport retry layer treats + // every 503 alike, so it would burn its whole budget re-asking a + // node that already said, terminally, that it is draining — and it + // would do so *before* this classifier ever saw the body that says + // which 503 it was. Retry policy on these routes belongs to the + // checkpoint loop, which can tell the cases apart. + let (request_id, response) = self.send(request, false).await?; + let status = response.status(); + let body = response.text().await.err_to_http(request_id.clone())?; + if !status.is_success() { + return Err(Error::LsmRoute { + fault: classify_lsm_fault(status.as_u16(), &body), + status: status.as_u16(), + message: format!("{route}: {status}: {body}"), + }); + } + serde_json::from_str(&body).map_err(|e| Error::Http { + source: format!("Failed to parse {route} response: {e}").into(), + request_id, + status_code: None, + }) + } + /// Build a POST request and attach the read-freshness headers /// (`x-lancedb-min-version`, `x-lancedb-min-timestamp`). fn post_read(&self, uri: &str) -> RequestBuilder { @@ -2361,6 +2401,35 @@ impl BaseTable for RemoteTable { }) } + async fn flush_lsm(&self) -> Result { + let request = self + .client + .post(&format!("/v1/table/{}/flush_lsm/", self.identifier)); + self.send_lsm_route(request, "flush_lsm").await + } + + async fn compact_lsm(&self, opts: CompactOptions) -> Result { + let request = self + .client + .post(&format!("/v1/table/{}/compact_lsm/", self.identifier)) + .json(&serde_json::json!({ + "max_generations_per_bucket": opts.max_generations_per_bucket, + })); + self.send_lsm_route(request, "compact_lsm").await + } + + async fn get_lsm_stats(&self, opts: LsmStatsOptions) -> Result> { + // Read-semantics POST, like `get_lsm_write_spec`. + let request = self + .post_read(&format!("/v1/table/{}/get_lsm_stats/", self.identifier)) + .json(&serde_json::json!({ + "include_generation_rows": opts.include_generation_rows, + })); + let parsed: GetLsmStatsResponse = self.send_lsm_route(request, "get_lsm_stats").await?; + // `null` — and only — when the table has no LSM write path. + Ok(parsed.lsm_stats) + } + async fn set_lsm_write_spec(&self, spec: LsmWriteSpec) -> Result<()> { self.check_mutable().await?; @@ -5765,6 +5834,265 @@ mod tests { assert!(table.get_lsm_write_spec().await.unwrap().is_none()); } + /// A flush that lands in an empty L0 converges on the flush response + /// alone — **zero** `compact` calls issued. Asserting the call count is + /// the point: `generations_consumed: 0` is also true of a loop that ran + /// a pointless pass. + #[tokio::test] + async fn test_checkpoint_short_circuits_on_empty_l0() { + let calls = Arc::new(std::sync::atomic::AtomicUsize::new(0)); + let seen = calls.clone(); + let table = Table::new_with_handler("my_table", move |request| { + let path = request.url().path().to_string(); + if path.contains("compact_lsm") { + seen.fetch_add(1, std::sync::atomic::Ordering::SeqCst); + panic!("an already-converged table must issue no compact calls"); + } + assert_eq!(path, "/v1/table/my_table/flush_lsm/"); + let response = serde_json::json!({ + "buckets": [{ + "shard_id": "b0", + "sealed_generation": 7, + "rows_sealed": 12, + "generations_remaining": 0, + }], + "rows_sealed": 12, + "generations_remaining": 0, + }); + http::Response::builder() + .status(200) + .body(response.to_string()) + .unwrap() + }); + + let report = table + .checkpoint_lsm(crate::table::CheckpointOptions::default()) + .await + .unwrap(); + assert!(report.converged); + assert_eq!(report.rows_sealed, 12); + assert_eq!(report.compact_calls, 0); + assert_eq!(calls.load(std::sync::atomic::Ordering::SeqCst), 0); + } + + /// The loop drives `compact` until the *server* reports zero remaining, + /// accumulating what each pass merged. + #[tokio::test] + async fn test_checkpoint_loops_compact_until_converged() { + let remaining = Arc::new(std::sync::atomic::AtomicUsize::new(3)); + let table = Table::new_with_handler("my_table", move |request| { + let path = request.url().path().to_string(); + let response = if path.contains("flush_lsm") { + serde_json::json!({ + "buckets": [{ "shard_id": "b0", "sealed_generation": 4, + "rows_sealed": 30, "generations_remaining": 3 }], + "rows_sealed": 30, + "generations_remaining": 3, + }) + } else { + let left = remaining.fetch_sub(1, std::sync::atomic::Ordering::SeqCst) - 1; + serde_json::json!({ + "buckets": [{ "shard_id": "b0", "generations_consumed": 1, + "rows_merged": 10, "generations_remaining": left }], + "generations_consumed": 1, + "rows_merged": 10, + "generations_remaining": left, + "highest_wal_entry_position": 991823, + }) + }; + http::Response::builder() + .status(200) + .body(response.to_string()) + .unwrap() + }); + + let report = table + .checkpoint_lsm(crate::table::CheckpointOptions::default()) + .await + .unwrap(); + assert!(report.converged); + assert_eq!(report.stop_reason, None); + assert_eq!(report.compact_calls, 3); + assert_eq!(report.generations_compacted, 3); + assert_eq!(report.rows_merged, 30); + assert_eq!(report.generations_remaining, 0); + } + + /// **Contention is not draining.** A 429 must be retried, not read as + /// the one terminal condition. This is the regression that would pass + /// silently otherwise: every layer says "unavailable" either way, and + /// the checkpoint just stops early blaming a healthy node. + #[tokio::test] + async fn test_checkpoint_retries_contention_and_does_not_report_draining() { + let compacts = Arc::new(std::sync::atomic::AtomicUsize::new(0)); + let seen = compacts.clone(); + let table = Table::new_with_handler("my_table", move |request| { + if request.url().path().contains("flush_lsm") { + let response = serde_json::json!({ + "buckets": [{ "shard_id": "b0", "sealed_generation": 1, + "rows_sealed": 5, "generations_remaining": 1 }], + "rows_sealed": 5, + "generations_remaining": 1, + }); + return http::Response::builder() + .status(200) + .body(response.to_string()) + .unwrap(); + } + // First two compacts: latch contention. Then it clears. + if seen.fetch_add(1, std::sync::atomic::Ordering::SeqCst) < 2 { + return http::Response::builder() + .status(429) + .body(r#"{"code":21,"error":"Too many concurrent writes"}"#.to_string()) + .unwrap(); + } + let response = serde_json::json!({ + "buckets": [{ "shard_id": "b0", "generations_consumed": 1, + "rows_merged": 5, "generations_remaining": 0 }], + "generations_consumed": 1, + "rows_merged": 5, + "generations_remaining": 0, + "highest_wal_entry_position": 1, + }); + http::Response::builder() + .status(200) + .body(response.to_string()) + .unwrap() + }); + + let report = table + .checkpoint_lsm(crate::table::CheckpointOptions::default()) + .await + .unwrap(); + assert!(report.converged, "contention must not abort the checkpoint"); + assert_ne!( + report.stop_reason, + Some(crate::table::CheckpointStopReason::NodeDraining), + "a busy compactor is not a draining node" + ); + assert_eq!( + compacts.load(std::sync::atomic::Ordering::SeqCst), + 3, + "assert the retry count, not just the outcome" + ); + } + + /// A draining node stops the loop immediately — **without retrying**. + /// Asserting the call count matters: a regression to retry-until- + /// deadline passes any test that only checks the outcome. + #[tokio::test] + async fn test_checkpoint_stops_without_retrying_on_draining() { + let calls = Arc::new(std::sync::atomic::AtomicUsize::new(0)); + let seen = calls.clone(); + let table = Table::new_with_handler("my_table", move |_request| { + seen.fetch_add(1, std::sync::atomic::Ordering::SeqCst); + http::Response::builder() + .status(503) + .body(r#"{"code":19,"error":"WAL node draining"}"#.to_string()) + .unwrap() + }); + + let report = table + .checkpoint_lsm(crate::table::CheckpointOptions::default()) + .await + .unwrap(); + assert!(!report.converged); + assert_eq!( + report.stop_reason, + Some(crate::table::CheckpointStopReason::NodeDraining) + ); + assert_eq!( + calls.load(std::sync::atomic::Ordering::SeqCst), + 1, + "draining is terminal: exactly one request, no retries" + ); + } + + /// WAL off ⇒ `None`; WAL on ⇒ a fully populated `Some` with no field + /// defaulting to a zero it did not measure. `include_generation_rows` + /// is off by default. + #[tokio::test] + async fn test_get_lsm_stats_round_trip() { + let table = Table::new_with_handler("my_table", |request| { + assert_eq!(request.url().path(), "/v1/table/my_table/get_lsm_stats/"); + let body = request.body().unwrap().as_bytes().unwrap(); + let body: serde_json::Value = serde_json::from_slice(body).unwrap(); + assert_eq!( + body["include_generation_rows"], false, + "generation rows cost a manifest read each and must be opt-in" + ); + let response = serde_json::json!({ + "lsm_stats": { + "buckets": [{ + "shard_id": "b0", + "status": "Active", + "writer_epoch": 3, + "manifest_version": 11, + "current_generation": 9, + "replay_after_wal_entry_position": 100, + "wal_entry_position_last_seen": 140, + "wal_lag": 40, + "generations": [{ "generation": 8, "path": "ab_gen_8", "bytes": 4096 }], + "l0_bytes": 4096, + "fenced": false, + "compaction_in_progress": false, + "active_memtable": { "generation": 9, "rows": 12, "bytes": 900, "batches": 2 }, + "frozen_memtables": [], + "memtable_indexes": [ + { "name": "vec_idx", "kind": "hnsw", "column": "vector" } + ], + }], + "bucket_count": 1, + "generations_total": 1, + "l0_bytes_total": 4096, + "memtable_rows_total": 12, + } + }); + http::Response::builder() + .status(200) + .body(response.to_string()) + .unwrap() + }); + + let stats = table + .get_lsm_stats(crate::table::LsmStatsOptions::default()) + .await + .unwrap() + .expect("a WAL-backed table reports Some"); + assert_eq!(stats.bucket_count, 1); + assert_eq!(stats.l0_bytes_total, 4096); + assert_eq!(stats.memtable_rows_total, 12); + let bucket = &stats.buckets[0]; + assert_eq!(bucket.wal_lag, 40); + assert_eq!(bucket.generations[0].generation, 8); + assert!( + bucket.generations[0].rows.is_none(), + "rows are absent unless asked for" + ); + // The row that answers "why is my fresh-tier vector search + // brute-force" — an absent hnsw entry is the whole explanation. + let indexes = bucket.memtable_indexes.as_ref().unwrap(); + assert_eq!(indexes[0].kind, "hnsw"); + assert_eq!(indexes[0].column, "vector"); + } + + #[tokio::test] + async fn test_get_lsm_stats_absent_when_wal_off() { + let table = Table::new_with_handler("my_table", |_request| { + http::Response::builder() + .status(200) + .body(serde_json::json!({ "lsm_stats": null }).to_string()) + .unwrap() + }); + assert!( + table + .get_lsm_stats(crate::table::LsmStatsOptions::default()) + .await + .unwrap() + .is_none() + ); + } + #[tokio::test] async fn test_wait_for_index() { let table = _make_table_with_indices(0); diff --git a/rust/lancedb/src/table.rs b/rust/lancedb/src/table.rs index 1f7b2f75b..5fe44541b 100644 --- a/rust/lancedb/src/table.rs +++ b/rust/lancedb/src/table.rs @@ -63,6 +63,7 @@ use self::merge::MergeInsertBuilder; mod add_data; pub mod branch_merge; +pub mod checkpoint; mod create_index; pub mod datafusion; pub(crate) mod dataset; @@ -82,6 +83,12 @@ pub use branch_merge::{ BranchDiff, ColumnChange, ColumnSummary, IndexSummary, MergeBlocker, MergeBlockerCode, MergeBranchResult, MergeBranchStatus, MergePreview, RowCountSummary, }; +pub use checkpoint::{ + BucketStats, CheckpointOptions, CheckpointReport, CheckpointStopReason, CompactBucketReport, + CompactOptions, CompactReport, FlushBucketReport, FlushReport, GenerationStats, LsmFault, + LsmStats, LsmStatsOptions, MemIndexStats, MemtableStats, +}; +use checkpoint::{CheckpointControl, checkpoint_fault}; pub use chrono::Duration; pub use delete::DeleteResult; use futures::future::join_all; @@ -627,6 +634,31 @@ pub trait BaseTable: std::fmt::Display + std::fmt::Debug + Send + Sync { message: "get_lsm_write_spec is not supported on this table type".into(), }) } + /// Seal every bucket's active memtable into L0. + /// + /// The default implementation returns `NotSupported`. + async fn flush_lsm(&self) -> Result { + Err(Error::NotSupported { + message: "flush_lsm is not supported on this table type".into(), + }) + } + /// Run one bounded L0 → base compaction pass per bucket. + /// + /// The default implementation returns `NotSupported`. + async fn compact_lsm(&self, _opts: CompactOptions) -> Result { + Err(Error::NotSupported { + message: "compact_lsm is not supported on this table type".into(), + }) + } + /// Read live LSM state, or `None` when the LSM write path is not + /// enabled for this table. + /// + /// The default implementation returns `NotSupported`. + async fn get_lsm_stats(&self, _opts: LsmStatsOptions) -> Result> { + Err(Error::NotSupported { + message: "get_lsm_stats is not supported on this table type".into(), + }) + } /// Drain and close any cached MemWAL shard writers for this table. /// /// The default implementation is a no-op; table types that maintain @@ -1620,6 +1652,161 @@ impl Table { self.inner.get_lsm_write_spec().await } + /// Converge this table's LSM write path into its base table. + /// + /// One `flush` (seal every memtable into L0), then `compact` until L0 + /// is empty. The loop runs here, in the client: each request does a + /// bounded unit of work and reports what is left, so there is no held + /// socket, no server-side task, and nothing to reconcile if you drop + /// this future partway through. + /// + /// **Best-effort.** Nothing is frozen, so with writes flowing there may + /// be new rows by the time it returns — `converged` means L0 was empty + /// as of the last pass. Idempotent and safe to run on a cadence: an + /// already-converged table costs one round trip and zero compaction + /// passes, because the opening flush reports an empty L0 and the loop + /// is never entered. + /// + /// Two independent bounds end the loop: `generations_remaining == 0` + /// and [`CheckpointOptions::deadline`]. A table under sustained write + /// load exits on the deadline with `converged: false` — a *result*, not + /// an error. + /// + /// # Example + /// + /// ```no_run + /// # use lancedb::table::{Table, CheckpointOptions}; + /// # async fn example(table: &Table) -> Result<(), Box> { + /// let report = table.checkpoint_lsm(CheckpointOptions::default()).await?; + /// if !report.converged { + /// println!("stopped early: {:?}", report.stop_reason); + /// } + /// # Ok(()) + /// # } + /// ``` + pub async fn checkpoint_lsm(&self, opts: CheckpointOptions) -> Result { + let started = std::time::Instant::now(); + let expired = |elapsed: std::time::Duration| { + opts.deadline.is_some_and(|deadline| elapsed >= deadline) + }; + + // Cap re-issues from `flush` after a 404. A node in a crash loop + // otherwise turns flush → compact → 404 → flush into a spin that + // burns the whole deadline making no progress and reports nothing + // useful about why. + const MAX_REISSUES: usize = 3; + + let mut report = CheckpointReport { + converged: false, + stop_reason: None, + rows_sealed: 0, + generations_compacted: 0, + rows_merged: 0, + generations_remaining: 0, + compact_calls: 0, + stats: None, + }; + let compact_opts = CompactOptions { + max_generations_per_bucket: opts.max_generations_per_bucket, + }; + + 'reissue: for attempt in 0..=MAX_REISSUES { + let flushed = match self.flush_lsm().await { + Ok(f) => f, + Err(e) => match checkpoint_fault(e, &mut report, attempt, MAX_REISSUES).await? { + // A retryable flush is just re-issued; the seal is + // idempotent (sealing an empty active memtable is a + // no-op, so repeats do not churn empty generations). + CheckpointControl::Retry | CheckpointControl::ReissueFromFlush => { + continue 'reissue; + } + CheckpointControl::Stop => break 'reissue, + }, + }; + // Only the first flush's seal counts toward `rows_sealed`; a + // re-issue after a 404 re-seals rows a previous attempt already + // reported. + if attempt == 0 { + report.rows_sealed = flushed.rows_sealed; + } + report.generations_remaining = flushed.generations_remaining; + report.converged = flushed.generations_remaining == 0; + + while !report.converged { + if expired(started.elapsed()) { + report.stop_reason = Some(CheckpointStopReason::DeadlineExceeded); + break 'reissue; + } + let compacted = match self.compact_lsm(compact_opts.clone()).await { + Ok(c) => c, + // A retryable fault re-enters this loop; a 404 restarts + // from `flush`, which is the call that re-claims and + // replays. + Err(e) => { + match checkpoint_fault(e, &mut report, attempt, MAX_REISSUES).await? { + CheckpointControl::Retry => continue, + CheckpointControl::ReissueFromFlush => continue 'reissue, + CheckpointControl::Stop => break 'reissue, + } + } + }; + report.compact_calls += 1; + report.generations_compacted += compacted.generations_consumed; + report.rows_merged += compacted.rows_merged; + report.generations_remaining = compacted.generations_remaining; + report.converged = compacted.generations_remaining == 0; + } + break; + } + + if opts.include_stats { + report.stats = self.get_lsm_stats(LsmStatsOptions::default()).await?; + } + Ok(report) + } + + /// Seal every bucket's active memtable into L0 without touching the + /// base table. + /// + /// Independently useful: flushing makes memtable rows readable from L0 + /// at a lower per-query cost. On a node that has not claimed this table + /// (after a restart, say) this claims it and replays its WAL log first — + /// reporting "nothing to flush" without replaying would be lying about + /// durable data. + pub async fn flush_lsm(&self) -> Result { + self.inner.flush_lsm().await + } + + /// Run one bounded L0 → base compaction pass per bucket, reporting what + /// it merged and what is left. + /// + /// One pass, not convergence: that is what makes each request's cost + /// bounded and gives a caller driving its own cadence a progress signal + /// per round trip. + pub async fn compact_lsm(&self, opts: CompactOptions) -> Result { + self.inner.compact_lsm(opts).await + } + + /// Read live LSM state: per-bucket fresh-tier detail plus a table-level + /// aggregate. + /// + /// Answers "how far behind is my fresh tier", "which bucket is hot", + /// and "why is my fresh-tier vector search brute-force". Mutates no + /// table state, though on a node that has not claimed this table it + /// claims it — exactly as a read would. + /// + /// `Ok(None)` — and only — when the LSM write path is not enabled for + /// the table, matching [`Table::get_lsm_write_spec`]. Stats is + /// fresh-tier only, so with the WAL off there is no bucket count, no + /// manifest, and no watermark; a struct of zeros would read as + /// measurements. + /// + /// Do not build a checkpoint's termination on this: the completion + /// predicate lives in the `flush` and `compact` responses. + pub async fn get_lsm_stats(&self, opts: LsmStatsOptions) -> Result> { + self.inner.get_lsm_stats(opts).await + } + /// Drain and close any cached MemWAL shard writers held for this table. /// /// When an [`LsmWriteSpec`] is installed, `merge_insert` opens MemWAL shard diff --git a/rust/lancedb/src/table/checkpoint.rs b/rust/lancedb/src/table/checkpoint.rs new file mode 100644 index 000000000..240860c41 --- /dev/null +++ b/rust/lancedb/src/table/checkpoint.rs @@ -0,0 +1,409 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The LanceDB Authors + +//! Converging a table's LSM write path into its base table. +//! +//! `checkpoint_lsm` is `flush` then `compact`, repeated until the fresh +//! tier is empty — and the loop runs **here, in the client**, not on the +//! server. Putting it server-side would mean a background task, which +//! means a single-flight intent, an intent that leaks on panic, a +//! bounded-iteration policy, an "is it done" observable, and a story for +//! every way a client can vanish mid-operation. None of that exists in +//! this shape. +//! +//! Three properties follow: +//! +//! * **No held socket.** Each request does a bounded unit of work and +//! returns. Nothing sits on a connection for minutes waiting on a merge. +//! * **Termination is structural.** Each `compact` drains at most one +//! prefix per bucket and reports what is left. A caller that wants to +//! stop, stops. +//! * **Completion is carried in the responses, not inferred.** A poll-a +//! -shared-counter design cannot distinguish "converged" from "hasn't +//! started yet"; being *told* `generations_remaining == 0` can. +//! +//! **Best-effort by construction.** Nothing is frozen, so a checkpoint +//! converges the fresh tier *as of some instant*; with writes flowing +//! there may be new rows by the time it returns. That is the correct +//! contract for a checkpoint and exactly Postgres's. It is idempotent, +//! abandonable at any point with zero consequence, and safe to run on a +//! cadence. + +use std::collections::HashMap; +use std::time::Duration; + +use serde::Deserialize; + +/// Options for [`crate::Table::checkpoint_lsm`]. +#[derive(Debug, Clone)] +pub struct CheckpointOptions { + /// Stop and report `converged: false` once this much wall-clock has + /// elapsed. `None` runs until the table converges or a terminal error + /// arrives — safe only on a table you know is not under write load. + pub deadline: Option, + /// Per-request, per-bucket bound on generations merged. `None` uses the + /// server's `compact_prefix_max`. Note this is **per bucket**: an + /// N-bucket table admits `N × max_generations_per_bucket` merges per + /// request. [`LsmStats::bucket_count`] reports N. + pub max_generations_per_bucket: Option, + /// Collect a [`LsmStats`] block into the report when the loop ends. + pub include_stats: bool, +} + +impl Default for CheckpointOptions { + fn default() -> Self { + Self { + deadline: Some(Duration::from_secs(300)), + max_generations_per_bucket: None, + include_stats: false, + } + } +} + +/// Why a checkpoint stopped short of convergence. Absent ⇒ it converged. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum CheckpointStopReason { + /// The caller's deadline elapsed with generations still in L0. A + /// *result*, not an error: a table under sustained write load may + /// legitimately never reach zero. Callers on a cadence ignore it; + /// callers about to disable the WAL check it. + DeadlineExceeded, + /// The owning node is draining. **Terminal** — the drain gate is a + /// one-way latch that never releases until the process restarts, so + /// retrying against that node spins until the deadline and then reports + /// failure when the truthful answer was available on the first + /// response. It resolves only when the ring changes. + /// + /// Stopping is also the *correct* answer: a node drain is a checkpoint + /// of every claimed table on the pod, driven by something that will + /// finish it. + NodeDraining, + /// The registry entry vanished repeatedly under pod restarts. A pod in + /// a crash loop would otherwise turn flush → compact → 404 → flush into + /// a spin that burns the whole deadline making no progress. + RepeatedlyUnclaimed, +} + +/// Outcome of a [`crate::Table::checkpoint_lsm`] run. +#[derive(Debug, Clone)] +pub struct CheckpointReport { + /// L0 was empty across every bucket as of the last pass. + /// + /// Precisely: the loop seals once at the top, so rows written *during* + /// the loop sit in the active memtable and are not sealed again. + /// Writes are never blocked to make this look better — a caller who + /// wants a quiescent table quiesces it. + pub converged: bool, + /// Set when `converged` is false. + pub stop_reason: Option, + /// Rows sealed out of the memtables by the single opening flush. + pub rows_sealed: u64, + /// L0 generations merged into base across every pass. + pub generations_compacted: usize, + /// Rows merged into base across every pass. + pub rows_merged: u64, + /// Generations still in L0 as of the last response. + pub generations_remaining: usize, + /// Number of `compact` round trips issued. Zero when the opening flush + /// landed in an empty L0 — the common case on a cadence. + pub compact_calls: usize, + /// Present when [`CheckpointOptions::include_stats`] was set. + pub stats: Option, +} + +/// Per-bucket result of one flush. +#[derive(Debug, Clone, Deserialize)] +pub struct FlushBucketReport { + pub shard_id: String, + /// `None` when the bucket's active memtable was already empty. + #[serde(default)] + pub sealed_generation: Option, + pub rows_sealed: u64, + pub generations_remaining: usize, +} + +/// Result of [`crate::Table::flush_lsm`]: everything written before the +/// seal is now in L0. +#[derive(Debug, Clone, Deserialize)] +pub struct FlushReport { + pub buckets: Vec, + pub rows_sealed: u64, + /// Summed across buckets, so `== 0` means every bucket is empty. This + /// is what lets an already-converged table cost one round trip and zero + /// compaction passes. + pub generations_remaining: usize, +} + +/// Options for [`crate::Table::compact_lsm`]. +#[derive(Debug, Clone, Default)] +pub struct CompactOptions { + /// See [`CheckpointOptions::max_generations_per_bucket`]. + pub max_generations_per_bucket: Option, +} + +/// Per-bucket result of one compaction pass. +#[derive(Debug, Clone, Deserialize)] +pub struct CompactBucketReport { + pub shard_id: String, + pub generations_consumed: usize, + pub rows_merged: u64, + pub generations_remaining: usize, +} + +/// Result of [`crate::Table::compact_lsm`]: one bounded L0 → base pass per +/// bucket. +/// +/// The per-bucket list is authoritative; the table-level fields are a +/// convenience. Consumed/merged/remaining are **sums**; +/// `highest_wal_entry_position` is a **max**, since it is a watermark. +#[derive(Debug, Clone, Deserialize)] +pub struct CompactReport { + pub buckets: Vec, + pub generations_consumed: usize, + pub rows_merged: u64, + pub generations_remaining: usize, + pub highest_wal_entry_position: u64, +} + +/// Options for [`crate::Table::get_lsm_stats`]. +#[derive(Debug, Clone, Default)] +pub struct LsmStatsOptions { + /// Report `rows` per L0 generation. Off by default: the shard manifest + /// stores only `{generation, path}`, so rows cost one Lance manifest + /// read per generation. + pub include_generation_rows: bool, +} + +/// One flushed L0 generation. +#[derive(Debug, Clone, Deserialize)] +pub struct GenerationStats { + pub generation: u64, + pub path: String, + pub bytes: u64, + /// Present only when [`LsmStatsOptions::include_generation_rows`] was set. + #[serde(default)] + pub rows: Option, +} + +/// One in-memory memtable. +#[derive(Debug, Clone, Deserialize)] +pub struct MemtableStats { + pub generation: u64, + pub rows: u64, + pub bytes: u64, + pub batches: u64, +} + +/// An index the memtable carries. An absent `hnsw` entry on a vector +/// column is the whole answer to "why is my fresh-tier vector search +/// brute-force". +#[derive(Debug, Clone, Deserialize)] +pub struct MemIndexStats { + pub name: String, + /// `btree` | `hnsw` | `fts`. + pub kind: String, + pub column: String, +} + +/// Live state of one bucket. A table is N buckets on one node; flattening +/// to a single number hides the one hot bucket that is usually why someone +/// opened this endpoint. +#[derive(Debug, Clone, Deserialize)] +pub struct BucketStats { + pub shard_id: String, + /// `Active` | `Sealed` (drop-table 2PC in flight). + pub status: String, + pub writer_epoch: u64, + pub manifest_version: u64, + pub current_generation: u64, + pub replay_after_wal_entry_position: u64, + pub wal_entry_position_last_seen: u64, + /// Accepted WAL entries not yet covered by a flush. + pub wal_lag: u64, + pub generations: Vec, + pub l0_bytes: u64, + pub fenced: bool, + pub compaction_in_progress: bool, + /// Absent for a `Sealed` bucket, whose in-memory state is torn down. + #[serde(default)] + pub active_memtable: Option, + #[serde(default)] + pub frozen_memtables: Option>, + #[serde(default)] + pub memtable_indexes: Option>, +} + +/// Live LSM state: per-bucket detail plus a table-level aggregate. +/// +/// Every field is measured. There is no "WAL is off" shape here — that +/// case is `None` from [`crate::Table::get_lsm_stats`], because a struct of +/// zeros would read as measurements. +#[derive(Debug, Clone, Deserialize)] +pub struct LsmStats { + pub buckets: Vec, + pub bucket_count: usize, + pub generations_total: usize, + pub l0_bytes_total: u64, + pub memtable_rows_total: u64, +} + +/// Server-side JSON envelope for `get_lsm_stats`. `lsm_stats` is null when +/// the table has no LSM write path. +#[derive(Debug, Deserialize)] +pub(crate) struct GetLsmStatsResponse { + #[serde(default)] + pub lsm_stats: Option, +} + +/// How a failed LSM request should be handled by the checkpoint loop. +/// +/// Five distinct conditions used to arrive at a client as one 503. This is +/// the classification that keeps them apart, and it reads the body's `code` +/// on 503 only — every other status is already unambiguous. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum LsmFault { + /// 429, or 503 with `ServiceUnavailable`: contention, a fenced writer, + /// no ready slot, or transport. Retry with backoff. + /// + /// Fenced and slot-unavailable are deliberately not distinguished from + /// contention here: the client action is identical. + Retry, + /// 404: the registry entry vanished under a node restart mid-loop. + /// `flush` is the call that re-claims and replays, so re-issue from + /// there — with a cap, or a crash-looping node turns the loop into a + /// spin. + ReissueFromFlush, + /// 503 with `InvalidTableState`: the owning node is draining. Terminal. + Draining, + /// 409 (dropping), 400 (not WAL-backed), or anything else. Terminal. + Fatal, +} + +/// Lance-namespace `ErrorCode::InvalidTableState`. Nothing else in the +/// server maps to it, which is what makes it usable as the draining signal: +/// a `WalDraining` variant that mapped to the generic service-unavailable +/// code would be byte-identical to every other 503 on the wire, leaving +/// this classifier matching message strings — exactly what a string `code` +/// field was rejected to avoid. +pub(crate) const CODE_INVALID_TABLE_STATE: u64 = 19; + +/// Classify an LSM-route failure from its status and response body. +/// +/// The body must be inspected **before** any helper that folds it into a +/// string and keeps only the status, or the discrimination designed across +/// three crates is lost at the last hop. +pub(crate) fn classify_lsm_fault(status: u16, body: &str) -> LsmFault { + match status { + 429 => LsmFault::Retry, + 404 => LsmFault::ReissueFromFlush, + 503 => { + let code = serde_json::from_str::>(body) + .ok() + .and_then(|m| m.get("code").and_then(|c| c.as_u64())); + if code == Some(CODE_INVALID_TABLE_STATE) { + LsmFault::Draining + } else { + LsmFault::Retry + } + } + _ => LsmFault::Fatal, + } +} + +/// What the checkpoint loop should do next after a failed request. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum CheckpointControl { + /// Re-issue the same call after a backoff. + Retry, + /// Restart from `flush`, which is what re-claims and replays. + ReissueFromFlush, + /// Stop; `report.stop_reason` has been set. + Stop, +} + +/// Base backoff between retries, doubled per consecutive retry up to +/// [`RETRY_BACKOFF_MAX`]. Latch contention clears in about the time one +/// compaction pass takes, so starting small is right; a saturated pool +/// wants the ceiling. +const RETRY_BACKOFF_BASE: Duration = Duration::from_millis(100); +const RETRY_BACKOFF_MAX: Duration = Duration::from_secs(5); + +/// Decide what a failed LSM request means for the loop, recording the stop +/// reason on `report` when it is terminal. Returns `Err` only for genuinely +/// fatal faults (dropping table, not WAL-backed, anything unrecognized) — +/// those are errors the caller should see, not outcomes. +pub(crate) async fn checkpoint_fault( + e: crate::Error, + report: &mut CheckpointReport, + attempt: usize, + max_reissues: usize, +) -> crate::Result { + #[cfg(feature = "remote")] + let fault = match &e { + crate::Error::LsmRoute { fault, .. } => Some(*fault), + _ => None, + }; + #[cfg(not(feature = "remote"))] + let fault: Option = None; + + match fault { + Some(LsmFault::Retry) => { + let backoff = RETRY_BACKOFF_BASE + .saturating_mul(1u32 << attempt.min(8) as u32) + .min(RETRY_BACKOFF_MAX); + tokio::time::sleep(backoff).await; + Ok(CheckpointControl::Retry) + } + Some(LsmFault::ReissueFromFlush) if attempt < max_reissues => { + Ok(CheckpointControl::ReissueFromFlush) + } + Some(LsmFault::ReissueFromFlush) => { + report.stop_reason = Some(CheckpointStopReason::RepeatedlyUnclaimed); + Ok(CheckpointControl::Stop) + } + Some(LsmFault::Draining) => { + report.stop_reason = Some(CheckpointStopReason::NodeDraining); + Ok(CheckpointControl::Stop) + } + Some(LsmFault::Fatal) | None => Err(e), + } +} + +#[cfg(test)] +mod tests { + use super::*; + + /// Every row of the taxonomy, asserted on the `(status, code)` pair the + /// client actually receives. A `_ =>` arm added later passes any test + /// that only asserts "it errored". + #[test] + fn taxonomy_round_trips() { + assert_eq!(classify_lsm_fault(429, r#"{"code":21}"#), LsmFault::Retry); + assert_eq!( + classify_lsm_fault(503, r#"{"code":19}"#), + LsmFault::Draining, + "a draining node must be told apart from every other 503" + ); + assert_eq!( + classify_lsm_fault(503, r#"{"code":17}"#), + LsmFault::Retry, + "fenced / no-slot / transport stay retryable" + ); + assert_eq!( + classify_lsm_fault(404, r#"{"code":4}"#), + LsmFault::ReissueFromFlush + ); + assert_eq!(classify_lsm_fault(409, r#"{"code":14}"#), LsmFault::Fatal); + assert_eq!(classify_lsm_fault(400, r#"{"code":13}"#), LsmFault::Fatal); + } + + /// A 503 whose body is missing, truncated, or not the expected JSON must + /// fall back to *retryable*, never to terminal: mistaking a healthy node + /// for a draining one aborts the checkpoint on a lie. + #[test] + fn unparseable_503_body_is_retryable() { + assert_eq!(classify_lsm_fault(503, ""), LsmFault::Retry); + assert_eq!(classify_lsm_fault(503, "gateway timeout"), LsmFault::Retry); + assert_eq!(classify_lsm_fault(503, r#"{"code":"19"}"#), LsmFault::Retry); + } +}