Expose legacy merge progress and uncertain writes to the master

crates/lance-context-core/src/merge_write_scope.rs

@@ -106,6 +106,9 @@ pub struct MergeWriteScope {
     completed_steps: Arc<AtomicU64>,
     work: Mutex<MergeWork>,
     observe_file_io: bool,
+    /// Legacy diagnostics must not introduce detached storage tasks or alter
+    /// cancellation/slot lifetimes before durable ownership is available.
+    observation_only: bool,
     local_io: Mutex<Vec<(lance_io::utils::tracking_store::IOTracker, u64)>>,
     progress: Mutex<Progress>,
     changed: Notify,
@@ -117,7 +120,23 @@ pub struct MergeWriteScope {
     io_shim: Option<Arc<dyn lance_io::object_store::WrappingObjectStore>>,
 }
 
+/// Local write observations, not a durable fence or proof of remote quiescence.
+#[derive(Clone, Debug, serde::Serialize)]
+pub struct MergeWriteStatus {
+    pub active_commits: usize,
+    pub admission_closed: bool,
+    pub uncertain_commit: bool,
+}
+
 impl MergeWriteScope {
+    pub fn write_status(&self) -> MergeWriteStatus {
+        let progress = self.progress.lock().unwrap();
+        MergeWriteStatus {
+            active_commits: progress.active,
+            admission_closed: progress.closed,
+            uncertain_commit: progress.uncertain,
+        }
+    }
     pub fn work(&self) -> MergeWork {
         self.work.lock().unwrap().clone()
     }
@@ -139,6 +158,13 @@ impl MergeWriteScope {
         Arc::new(Self::default())
     }
 
+    pub fn for_legacy_observation() -> Arc<Self> {
+        Arc::new(Self {
+            observation_only: true,
+            ..Self::default()
+        })
+    }
+
     pub fn with_authorizer(authorizer: Arc<dyn CommitAuthorizer>) -> Arc<Self> {
         Arc::new(Self {
             authorizer: Some(authorizer),
@@ -241,6 +267,18 @@ pub(crate) fn wal_batch_read(rows: usize, bytes: usize) {
 }
 
 pub(crate) fn base_committed() {
+    // Legacy observation counts at the commit handler, including fallback
+    // paths which also invoke the owned executor's logical checkpoints.
+    if CURRENT
+        .try_with(|scope| scope.observation_only)
+        .unwrap_or(false)
+    {
+        return;
+    }
+    record_base_commit();
+}
+
+fn record_base_commit() {
     let _ = CURRENT.try_with(|scope| {
         let mut work = scope.work.lock().unwrap();
         work.base_commits += 1;
@@ -250,6 +288,16 @@ pub(crate) fn base_committed() {
 }
 
 pub(crate) fn generations_reclaimed(count: usize) {
+    if CURRENT
+        .try_with(|scope| scope.observation_only)
+        .unwrap_or(false)
+    {
+        return;
+    }
+    record_reclaimed(count);
+}
+
+fn record_reclaimed(count: usize) {
     if count == 0 {
         return;
     }
@@ -423,6 +471,18 @@ where
         scope,
         completed: false,
     };
+    if active.scope.observation_only {
+        let mut active = active;
+        let result = write.await;
+        if progressed(&result) {
+            active.scope.completed_steps.fetch_add(1, Ordering::Relaxed);
+        }
+        if result.is_err() {
+            active.scope.mark_uncertain();
+        }
+        active.completed = true;
+        return result;
+    }
     let (send, receive) = oneshot::channel();
     let leaf = tokio::spawn(async move {
         let mut active = active;
@@ -465,8 +525,10 @@ pub(crate) async fn drain_generations(
             .version
             .checked_add(1)
             .ok_or_else(|| Error::io("manifest version overflow"))?;
+        let before = next.sstables.len();
         next.sstables
             .retain(|g| !generations.contains(&g.generation));
+        let removed = before - next.sstables.len();
         authorize(&resource, next.version).await?;
         let writer = store.clone();
         let scope = CURRENT.try_with(Arc::clone).ok();
@@ -487,7 +549,15 @@ pub(crate) async fn drain_generations(
         )
         .await?;
         match written {
-            Ok(_) => return Ok(next),
+            Ok(_) => {
+                if CURRENT
+                    .try_with(|scope| scope.observation_only)
+                    .unwrap_or(false)
+                {
+                    record_reclaimed(removed);
+                }
+                return Ok(next);
+            }
             Err(Error::RetryableCommitConflict { .. }) => continue,
             Err(error) => return Err(error),
         }
@@ -718,6 +788,13 @@ impl CommitHandler for GuardedCommit {
             )
             .await?;
             *manifest = committed;
+            if location.is_ok()
+                && CURRENT
+                    .try_with(|scope| scope.observation_only)
+                    .unwrap_or(false)
+            {
+                record_base_commit();
+            }
             location
         };
         match &self.scope {
@@ -796,6 +873,45 @@ mod tests {
     use super::*;
     use std::sync::atomic::{AtomicBool, Ordering};
 
+    #[tokio::test]
+    async fn legacy_observation_retains_uncertainty_without_detaching_storage() {
+        struct LocalWrite(Arc<AtomicBool>);
+        impl Drop for LocalWrite {
+            fn drop(&mut self) {
+                self.0.store(false, Ordering::SeqCst);
+            }
+        }
+        let scope = MergeWriteScope::for_legacy_observation();
+        let writing = Arc::new(AtomicBool::new(false));
+        let flag = writing.clone();
+        let (entered, began) = oneshot::channel();
+        let run_scope = scope.clone();
+        let caller = tokio::spawn(async move {
+            run_scope
+                .run(shield(async move {
+                    flag.store(true, Ordering::SeqCst);
+                    let _local_write = LocalWrite(flag);
+                    entered.send(()).unwrap();
+                    std::future::pending::<Result<()>>().await
+                }))
+                .await
+        });
+        began.await.unwrap();
+        assert_eq!(scope.write_status().active_commits, 1);
+        caller.abort();
+        caller.await.unwrap_err();
+        assert!(
+            !writing.load(Ordering::SeqCst),
+            "diagnostics must not detach storage work"
+        );
+        assert_eq!(scope.write_status().active_commits, 0);
+        assert!(
+            scope.write_status().uncertain_commit,
+            "remote completion is still unknown"
+        );
+        assert_eq!(scope.completed_steps(), 0);
+    }
+
     #[tokio::test]
     async fn phase_changes_and_observation_do_not_count_as_completed_work() {
         let scope = MergeWriteScope::new();
@@ -939,6 +1055,8 @@ mod tests {
         began.await.unwrap();
         merge.abort();
         assert!(merge.await.unwrap_err().is_cancelled());
+        assert_eq!(scope.write_status().active_commits, 1);
+        assert!(!scope.write_status().uncertain_commit);
         let draining_scope = scope.clone();
         let drain = tokio::spawn(async move { draining_scope.drain().await });
         tokio::task::yield_now().await;
@@ -948,6 +1066,8 @@ mod tests {
         );
         release.send(()).unwrap();
         drain.await.unwrap();
+        assert_eq!(scope.write_status().active_commits, 0);
+        assert!(scope.write_status().admission_closed);
         assert!(committed.load(Ordering::SeqCst));
         assert!(!scope.has_uncertain_commit());
         assert!(

crates/lance-context-core/src/store_base.rs

@@ -1353,8 +1353,11 @@ impl StorageBase {
                             break 'generations;
                         }
                     }
+                    crate::merge_write_scope::wal_batch_read(
+                        batch.num_rows(),
+                        batch.get_array_memory_size(),
+                    );
                     current_batches.push(batch);
-                    crate::merge_write_scope::checkpoint();
                 }
             }
             generation_batches.push((flushed.generation, current_batches));
@@ -1429,6 +1432,7 @@ impl StorageBase {
         if let Some(watermark) = watermark {
             let params = WriteParams {
                 mode: WriteMode::Append,
+                write_progress: Some(crate::merge_write_scope::write_progress()),
                 ..Default::default()
             };
             let staged = lance::dataset::InsertBuilder::new(Arc::new(self.dataset.clone()))
@@ -1448,7 +1452,15 @@ impl StorageBase {
                 batches.into_iter().map(Ok::<RecordBatch, ArrowError>),
                 merge_schema,
             );
-            Box::pin(self.dataset.append(reader, None)).await?;
+            Box::pin(self.dataset.append(
+                reader,
+                Some(WriteParams {
+                    mode: WriteMode::Append,
+                    write_progress: Some(crate::merge_write_scope::write_progress()),
+                    ..Default::default()
+                }),
+            ))
+            .await?;
         }
         Ok(())
     }

crates/lance-context-master/src/merge_execution.rs

@@ -179,8 +179,11 @@ async fn run_legacy_workers(
             Err((false, error)) => {
                 // A response error, reset, or timeout after sending may leave
                 // an unfenced legacy writer alive. Do not advance the fan-out.
+                // Fetch evidence from the same worker without opening stores or
+                // querying etcd. This never permits retry, takeover or restart.
+                let progress = legacy_progress(http, endpoint, target).await;
                 return Err(format!(
-                    "merged {reclaimed} generations across {completed}/{} workers (legacy); stopped at {endpoint}: {error}",
+                    "merged {reclaimed} generations across {completed}/{} workers (legacy); stopped at {endpoint}: {error}; legacy progress: {progress}",
                     endpoints.len()
                 ));
             }
@@ -201,6 +204,67 @@ async fn run_legacy_workers(
     }
 }
 
+async fn legacy_progress(http: &reqwest::Client, endpoint: &str, target: &str) -> String {
+    let result =
+        tokio::time::timeout(Duration::from_secs(2), async {
+            let mut response = http
+                .get(format!("{endpoint}/api/v1/internal/legacy-merges"))
+                .query(&[("target", target)])
+                .send()
+                .await
+                .map_err(|e| e.to_string())?
+                .error_for_status()
+                .map_err(|e| e.to_string())?;
+            let mut bytes = Vec::new();
+            while let Some(chunk) = response.chunk().await.map_err(|e| e.to_string())? {
+                if bytes.len().saturating_add(chunk.len()) > 64 * 1024 {
+                    return Err("snapshot exceeds 64 KiB".to_string());
+                }
+                bytes.extend_from_slice(&chunk);
+            }
+            let value: serde_json::Value =
+                serde_json::from_slice(&bytes).map_err(|e| e.to_string())?;
+            if value["instance"].as_str().is_none()
+                || value["snapshot"]["protocol"] != 1
+                || value["snapshot"]["takeover_authorized"] != false
+            {
+                return Err("unsupported progress response".to_string());
+            }
+            let snapshot = &value["snapshot"];
+            let executions = snapshot["executions"]
+                .as_array()
+                .ok_or_else(|| "missing execution observations".to_string())?;
+            // Task errors are stored in etcd. Keep the durable summary small even
+            // if the endpoint returns a large history or an unexpected field.
+            let summaries: Vec<_> = executions.iter().take(8).map(|row| serde_json::json!({
+            "id": row["id"].as_str().map(|s| s.chars().take(128).collect::<String>()),
+            "phase": row["phase"].as_str().map(|s| s.chars().take(32).collect::<String>()),
+            "elapsed_ms": row["elapsed_ms"].as_u64(),
+            "completed_steps": row["completed_steps"].as_u64(),
+            "base_commits": row["work"]["base_commits"].as_u64(),
+            "generations_reclaimed": row["work"]["generations_reclaimed"].as_u64(),
+            "wal_decoded_bytes": row["work"]["wal_decoded_bytes"].as_u64(),
+            "encoded_bytes": row["work"]["encoded_bytes"].as_u64(),
+            "active_commits": row["writes"]["active_commits"].as_u64(),
+            "uncertain_commit": row["writes"]["uncertain_commit"].as_bool(),
+        })).collect();
+            Ok(serde_json::json!({
+            "instance": value["instance"].as_str().unwrap().chars().take(128).collect::<String>(),
+            "takeover_authorized": false,
+            "untracked_admissions": snapshot["untracked_admissions"].as_u64(),
+            "evicted_history": snapshot["evicted_history"].as_u64(),
+            "truncated": executions.len() > 8 || snapshot["truncated"].as_bool() != Some(false),
+            "executions": summaries,
+        }).to_string())
+        })
+        .await;
+    match result {
+        Ok(Ok(snapshot)) => snapshot,
+        Ok(Err(error)) => format!("unavailable ({error}); ownership unresolved"),
+        Err(_) => "unavailable (observation timed out); ownership unresolved".into(),
+    }
+}
+
 struct Deadlines {
     queue: tokio::time::Instant,
     idle_timeout: Duration,
@@ -1197,6 +1261,46 @@ mod legacy_connection_tests {
         .await
     }
 
+    #[tokio::test]
+    async fn progress_evidence_is_reported_but_never_allows_unsafe_fanout() {
+        let (bad, failed_server) = serve(Router::new()
+            .route("/api/v1/internal/legacy-merges", axum::routing::get(
+                |axum::extract::Query(query): axum::extract::Query<std::collections::HashMap<String, String>>| async move {
+                    assert_eq!(query.get("target").unwrap(), "table");
+                    Json(serde_json::json!({"instance":"boot-1", "snapshot":{
+                        "protocol":1, "takeover_authorized":false,
+                        "executions":[{"id":"exec-1","completed_steps":42,"writes":{"active_commits":1}}]
+                    }}))
+                }))
+            .fallback(post(|| async { StatusCode::INTERNAL_SERVER_ERROR }))).await;
+        let calls = Arc::new(AtomicUsize::new(0));
+        let (good, server) = healthy(calls.clone()).await;
+        let http = crate::state::worker_http_client().unwrap();
+        let error = run_legacy_workers(&http, &[bad, good], "table")
+            .await
+            .unwrap_err();
+        assert!(error.contains("exec-1"), "{error}");
+        assert!(error.contains("active_commits"), "{error}");
+        assert_eq!(calls.load(Ordering::SeqCst), 0);
+        failed_server.abort();
+        server.abort();
+    }
+
+    #[tokio::test]
+    async fn oversized_progress_is_bounded_and_not_accepted_as_handoff_evidence() {
+        let (endpoint, server) =
+            serve(Router::new().fallback(axum::routing::get(|| async { "x".repeat(65537) }))).await;
+        let evidence = super::legacy_progress(
+            &crate::state::worker_http_client().unwrap(),
+            &endpoint,
+            "table",
+        )
+        .await;
+        assert!(evidence.contains("exceeds 64 KiB"), "{evidence}");
+        assert!(evidence.contains("ownership unresolved"));
+        server.abort();
+    }
+
     #[tokio::test]
     async fn connection_refusal_preserves_progress_on_both_sides() {
         // Reserve a non-listening port so another test cannot take it.