Fence Optimize against late recovery intents (#347)

* Fence optimize against late recovery intents

* Preserve manifest compaction after table errors

crates/omnigraph/src/db/omnigraph/optimize.rs

@@ -195,6 +195,24 @@ pub async fn optimize_all_tables(db: &Omnigraph) -> Result<Vec<TableOptimizeStat
              recovery sweep before optimizing",
         ));
     }
+    // Deterministic race seam: the broad fast-path probe above has completed,
+    // but main's branch-writer gate is not held yet. A writer may arm recovery
+    // in this window; the load-bearing check below runs only after Optimize owns
+    // the branch authority every sidecar-enrolled main writer must cross.
+    crate::failpoints::maybe_fail(
+        crate::failpoints::names::OPTIMIZE_POST_RECOVERY_CHECK_PRE_MAIN_GATE,
+    )?;
+
+    // Optimize publishes each changed table pointer as a main graph commit, so
+    // its authority is branch-wide even though its physical effects are
+    // table-local. Hold main's process-local writer gate through every table
+    // effect/publish and the final physical-only __manifest compaction. This
+    // still lets this one Optimize call process distinct tables concurrently,
+    // but another main writer cannot arm an intent after our final check or
+    // have its fixed graph-head authority invalidated underneath recovery.
+    let _main_branch_guard = db.write_queue().acquire_branch(None).await;
+    ensure_no_pending_recovery_for_optimize_under_main_gate(db).await?;
+
     let snapshot = db.fresh_snapshot_for_branch(None).await?;
 
     // Compute per-table paths up front, in a scope that drops the catalog
@@ -603,6 +621,41 @@ async fn optimize_one_table(
     Ok(stat)
 }
 
+/// Final recovery-ownership check for main-branch Optimize.
+///
+/// The caller must hold main's branch-writer gate and retain it through all data
+/// effects and graph-head publishes. The top-level probe stays deliberately
+/// conservative and refuses any pre-existing sidecar; this final check rejects
+/// every late main-target intent plus graph-global SchemaApply. Table-disjoint
+/// intents still overlap because each fixed recovery authority includes the
+/// shared `graph_head:main` that Optimize advances.
+///
+/// This is a process-local gate proof, matching the repository's documented
+/// single-writer-process boundary. Separate processes remain governed by Lance
+/// OCC and recovery classification; this helper is not a distributed lock.
+async fn ensure_no_pending_recovery_for_optimize_under_main_gate(db: &Omnigraph) -> Result<()> {
+    let sidecars = crate::db::manifest::list_sidecars(db.root_uri(), db.storage_adapter()).await?;
+    let blocking = sidecars.iter().find(|sidecar| {
+        sidecar.writer_kind == crate::db::manifest::SidecarKind::SchemaApply
+            || sidecar
+                .branch
+                .as_deref()
+                .filter(|branch| *branch != "main")
+                .is_none()
+    });
+    if let Some(sidecar) = blocking {
+        return Err(OmniError::recovery_required(
+            sidecar.operation_id.clone(),
+            format!(
+                "pending {:?} recovery operation on branch '{}' blocks optimize",
+                sidecar.writer_kind,
+                sidecar.branch.as_deref().unwrap_or("main"),
+            ),
+        ));
+    }
+    Ok(())
+}
+
 /// Bound on the app-level retry of an internal-table compaction against a
 /// concurrent live writer (see [`is_retryable_lance_conflict`]).
 const COMPACTION_RETRY_BUDGET: u32 = 5;

crates/omnigraph/src/failpoints.rs

@@ -123,6 +123,11 @@ pub mod names {
     pub const OPEN_BEFORE_SCHEMA_CONTRACT_READ: &str = "open.before_schema_contract_read";
     pub const OPTIMIZE_BEFORE_COMPACT: &str = "optimize.before_compact";
     pub const OPTIMIZE_INJECT_REINDEX_CONFLICT: &str = "optimize.inject_reindex_conflict";
+    /// After Optimize's broad recovery fast-path check, before the main-branch
+    /// writer gate is acquired. Tests arm a late recovery intent in this window
+    /// and prove the under-branch-gate check refuses to advance around it.
+    pub const OPTIMIZE_POST_RECOVERY_CHECK_PRE_MAIN_GATE: &str =
+        "optimize.post_recovery_check_pre_main_gate";
     pub const OPTIMIZE_POST_PHASE_B_PRE_MANIFEST_COMMIT: &str =
         "optimize.post_phase_b_pre_manifest_commit";
     pub const RECOVERY_BEFORE_ROLL_FORWARD_PUBLISH: &str = "recovery.before_roll_forward_publish";

crates/omnigraph/tests/failpoints.rs

@@ -5425,13 +5425,281 @@ async fn optimize_phase_b_failure_recovered_on_next_open() {
         .expect("schema apply after optimize recovery must succeed");
 }
 
-/// Separately-opened handles for one root share the table/recovery gates. A
-/// served insert that starts while optimize is paused before compaction must
-/// wait, then commit after optimize releases its sidecar lifetime; it must not
-/// deadlock in the synchronous recovery barrier or lose either result.
+async fn seed_optimize_late_sidecar_race(dir: &tempfile::TempDir) {
+    let mut seed = helpers::init_and_load(dir).await;
+    // Leave real compaction work behind so a missing barrier advances Person
+    // instead of accidentally passing because Optimize was a no-op.
+    helpers::commit_many(&mut seed, 4).await;
+}
+
+/// Optimize's entry recovery probe is only a fast path. A graph-global
+/// SchemaApply can arm its durable sidecar after that probe while Optimize is
+/// waiting for the main-branch gate, then fail before staging any schema/table
+/// effect. Once Optimize owns that gate it must relist recovery state and return
+/// the older operation's exact RecoveryRequired id before opening HEAD, writing
+/// an Optimize sidecar, or moving either the table or manifest.
+#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
+#[serial]
+#[serial(optimize)]
+async fn optimize_rechecks_late_schema_apply_sidecar_after_main_gate() {
+    let _scenario = FailScenario::setup();
+    let dir = tempfile::tempdir().unwrap();
+    let uri = dir.path().to_str().unwrap().to_string();
+
+    seed_optimize_late_sidecar_race(&dir).await;
+
+    // Both handles must exist before the late sidecar: a new read-write open
+    // would recover it instead of exercising the long-lived-handle race.
+    let optimize_db = std::sync::Arc::new(Omnigraph::open(&uri).await.unwrap());
+    let schema_db = Omnigraph::open(&uri).await.unwrap();
+    let person_uri = node_table_uri(&uri, "Person");
+    let manifest_uri = format!("{uri}/__manifest");
+    let manifest_before = version_main(optimize_db.as_ref()).await.unwrap();
+    let physical_manifest_before = lance::Dataset::open(&manifest_uri)
+        .await
+        .unwrap()
+        .version()
+        .version;
+    let graph_head_before = branch_head_commit_id(dir.path(), "main").await.unwrap();
+    let person_head_before = lance::Dataset::open(&person_uri)
+        .await
+        .unwrap()
+        .version()
+        .version;
+
+    let rendezvous = helpers::failpoint::Rendezvous::park_first(
+        names::OPTIMIZE_POST_RECOVERY_CHECK_PRE_MAIN_GATE,
+    );
+    let optimize_task_db = std::sync::Arc::clone(&optimize_db);
+    let optimize = tokio::spawn(async move { optimize_task_db.optimize().await });
+    rendezvous.wait_until_reached().await;
+
+    // Adding @index changes only accepted schema metadata. The failpoint fires
+    // after SchemaApply's zero-table v5 sidecar is durable, but before schema
+    // staging or any physical table effect.
+    let indexed_schema = helpers::TEST_SCHEMA.replace("age: I32?", "age: I32? @index");
+    {
+        let _failpoint = ScopedFailPoint::new(names::SCHEMA_APPLY_BEFORE_STAGING_WRITE, "return");
+        let error = schema_db.apply_schema(&indexed_schema).await.unwrap_err();
+        assert!(
+            error
+                .to_string()
+                .contains("injected failpoint triggered: schema_apply.before_staging_write"),
+            "unexpected SchemaApply failure: {error}",
+        );
+    }
+    let schema_operation_id = single_sidecar_operation_id(dir.path());
+
+    rendezvous.release();
+    let optimize_error = tokio::time::timeout(std::time::Duration::from_secs(20), optimize)
+        .await
+        .expect("Optimize task hung after releasing the recovery-check rendezvous")
+        .unwrap()
+        .expect_err("Optimize must refuse the late graph-global recovery intent");
+    match optimize_error {
+        OmniError::RecoveryRequired { operation_id, .. } => {
+            assert_eq!(operation_id, schema_operation_id)
+        }
+        other => panic!("expected exact RecoveryRequired attribution, got {other}"),
+    }
+    drop(rendezvous);
+
+    assert_eq!(
+        single_sidecar_operation_id(dir.path()),
+        schema_operation_id,
+        "Optimize must not add its own sidecar beside the older SchemaApply intent",
+    );
+    assert_eq!(
+        version_main(optimize_db.as_ref()).await.unwrap(),
+        manifest_before,
+        "late-barrier refusal must happen before any data or internal manifest effect",
+    );
+    assert_eq!(
+        lance::Dataset::open(&manifest_uri)
+            .await
+            .unwrap()
+            .version()
+            .version,
+        physical_manifest_before,
+        "late-barrier refusal must not compact the physical __manifest afterward",
+    );
+    assert_eq!(
+        branch_head_commit_id(dir.path(), "main").await.unwrap(),
+        graph_head_before,
+        "late-barrier refusal must not manufacture graph lineage",
+    );
+    assert_eq!(
+        lance::Dataset::open(&person_uri)
+            .await
+            .unwrap()
+            .version()
+            .version,
+        person_head_before,
+        "Optimize must not move the target table HEAD around the older sidecar",
+    );
+
+    drop(optimize_db);
+    drop(schema_db);
+
+    // Full recovery owns the abandoned SchemaApply intent. Once it rolls the
+    // metadata-only attempt back, the same Optimize work is safe and succeeds.
+    let recovered = Omnigraph::open(&uri)
+        .await
+        .expect("read-write open must recover the abandoned SchemaApply intent");
+    assert_post_recovery_invariants(
+        dir.path(),
+        &schema_operation_id,
+        RecoveryExpectation::RolledBack { tables: vec![] },
+    )
+    .await
+    .unwrap();
+    recovered
+        .optimize()
+        .await
+        .expect("Optimize must succeed after the older recovery intent is resolved");
+}
+
+/// A main-branch recovery intent overlaps Optimize even when it touched a
+/// different table: both operations publish through the shared
+/// `graph_head:main`. This pins the branch-wide half of the final barrier. A
+/// confirmed Company mutation fails before visibility while Optimize is parked;
+/// Optimize must not compact Person or advance main around that fixed authority.
+#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
+#[serial]
+#[serial(optimize)]
+async fn optimize_rechecks_late_disjoint_main_sidecar_after_main_gate() {
+    let _scenario = FailScenario::setup();
+    let dir = tempfile::tempdir().unwrap();
+    let uri = dir.path().to_str().unwrap().to_string();
+    seed_optimize_late_sidecar_race(&dir).await;
+
+    let optimize_db = std::sync::Arc::new(Omnigraph::open(&uri).await.unwrap());
+    let writer_db = Omnigraph::open(&uri).await.unwrap();
+    let person_uri = node_table_uri(&uri, "Person");
+    let manifest_uri = format!("{uri}/__manifest");
+
+    let rendezvous = helpers::failpoint::Rendezvous::park_first(
+        names::OPTIMIZE_POST_RECOVERY_CHECK_PRE_MAIN_GATE,
+    );
+    let optimize_task_db = std::sync::Arc::clone(&optimize_db);
+    let optimize = tokio::spawn(async move { optimize_task_db.optimize().await });
+    rendezvous.wait_until_reached().await;
+
+    // The interrupted writer affects Company, while the productive Optimize
+    // work this test protects is Person compaction. Its v3 sidecar is confirmed
+    // and durable; only the main graph-head authority overlaps.
+    {
+        let _failpoint =
+            ScopedFailPoint::new(names::MUTATION_POST_FINALIZE_PRE_PUBLISHER, "return");
+        let error = writer_db
+            .mutate(
+                "main",
+                OCC_DISJOINT_MUTATIONS,
+                "insert_company",
+                &params(&[("$name", "InterruptedCo")]),
+            )
+            .await
+            .expect_err("Company mutation must stop after confirming its physical effect");
+        assert!(
+            error
+                .to_string()
+                .contains("mutation.post_finalize_pre_publisher"),
+            "unexpected mutation failure: {error}",
+        );
+    }
+    let writer_operation_id = single_sidecar_operation_id(dir.path());
+
+    // Baseline after the older writer's owned Company effect. Nothing below may
+    // add Optimize movement before recovery resolves that intent.
+    let manifest_before_optimize = version_main(optimize_db.as_ref()).await.unwrap();
+    let physical_manifest_before_optimize = lance::Dataset::open(&manifest_uri)
+        .await
+        .unwrap()
+        .version()
+        .version;
+    let graph_head_before_optimize = branch_head_commit_id(dir.path(), "main").await.unwrap();
+    let person_head_before_optimize = lance::Dataset::open(&person_uri)
+        .await
+        .unwrap()
+        .version()
+        .version;
+
+    rendezvous.release();
+    let optimize_error = tokio::time::timeout(std::time::Duration::from_secs(20), optimize)
+        .await
+        .expect("Optimize task hung after releasing the recovery-check rendezvous")
+        .unwrap()
+        .expect_err("Optimize must refuse a late main-branch recovery intent");
+    match optimize_error {
+        OmniError::RecoveryRequired { operation_id, .. } => {
+            assert_eq!(operation_id, writer_operation_id)
+        }
+        other => panic!("expected exact RecoveryRequired attribution, got {other}"),
+    }
+    drop(rendezvous);
+
+    assert_eq!(
+        single_sidecar_operation_id(dir.path()),
+        writer_operation_id,
+        "Optimize must not add a sidecar beside the disjoint main writer",
+    );
+    assert_eq!(
+        version_main(optimize_db.as_ref()).await.unwrap(),
+        manifest_before_optimize,
+        "Optimize must not publish a main-table pointer around the older intent",
+    );
+    assert_eq!(
+        lance::Dataset::open(&manifest_uri)
+            .await
+            .unwrap()
+            .version()
+            .version,
+        physical_manifest_before_optimize,
+        "Optimize must not compact the physical __manifest after refusing",
+    );
+    assert_eq!(
+        branch_head_commit_id(dir.path(), "main").await.unwrap(),
+        graph_head_before_optimize,
+        "Optimize must not invalidate the older intent's fixed graph-head authority",
+    );
+    assert_eq!(
+        lance::Dataset::open(&person_uri)
+            .await
+            .unwrap()
+            .version()
+            .version,
+        person_head_before_optimize,
+        "disjoint Company recovery must still block Person compaction",
+    );
+
+    drop(optimize_db);
+    drop(writer_db);
+
+    let recovered = Omnigraph::open(&uri)
+        .await
+        .expect("read-write open must recover the confirmed Company mutation");
+    assert_post_recovery_invariants(
+        dir.path(),
+        &writer_operation_id,
+        RecoveryExpectation::RolledForwardOriginalLineage {
+            tables: vec![TableExpectation::main("node:Company")],
+        },
+    )
+    .await
+    .unwrap();
+    recovered
+        .optimize()
+        .await
+        .expect("Optimize must succeed after the main recovery intent is resolved");
+}
+
+/// Optimize retains main's branch gate after its final recovery relist and
+/// through its effects. A Company insert started while productive Person
+/// compaction is paused must therefore wait despite sharing no table gate, then
+/// commit after Optimize releases the branch-wide legacy-adapter envelope.
 #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
 #[serial(optimize)]
-async fn optimize_serializes_concurrent_insert_across_handles() {
+async fn optimize_holds_main_gate_through_disjoint_table_effects() {
     let _scenario = FailScenario::setup();
     let dir = tempfile::tempdir().unwrap();
     let uri = dir.path().to_str().unwrap().to_string();
@@ -5452,8 +5720,8 @@ async fn optimize_serializes_concurrent_insert_across_handles() {
 
     let db_b = std::sync::Arc::new(Omnigraph::open(&uri).await.unwrap());
 
-    // Pause optimize before compaction while it owns the Person table's
-    // sidecar/queue lifetime.
+    // Pause productive Person optimize after the under-main-gate recovery
+    // relist and after its Person sidecar is durable.
     let failpoint = ScopedFailPoint::new(names::OPTIMIZE_BEFORE_COMPACT, "pause");
 
     let uri_opt = uri.clone();
@@ -5473,16 +5741,16 @@ async fn optimize_serializes_concurrent_insert_across_handles() {
         writer_db
             .mutate(
                 "main",
-                MUTATION_QUERIES,
-                "insert_person",
-                &mixed_params(&[("$name", "eve")], &[("$age", 34)]),
+                OCC_DISJOINT_MUTATIONS,
+                "insert_company",
+                &params(&[("$name", "QueuedCo")]),
             )
             .await
     });
     tokio::time::sleep(std::time::Duration::from_millis(50)).await;
     assert!(
         !writer.is_finished(),
-        "same-root writer must wait for optimize's guarded sidecar lifetime"
+        "table-disjoint main writer must wait for Optimize's branch-gate lifetime"
     );
 
     drop(failpoint); // release optimize
@@ -5490,18 +5758,23 @@ async fn optimize_serializes_concurrent_insert_across_handles() {
         .await
         .expect("optimize task hung")
         .unwrap();
-    result.expect("optimize must finish before the queued same-table writer");
+    result.expect("optimize must finish before the queued disjoint main writer");
     writer
         .await
         .expect("writer task panicked")
-        .expect("queued insert must resume after optimize");
+        .expect("queued Company insert must resume after Optimize");
 
-    // No lost write: 4 seed + eve all present; graph remains re-optimizable.
+    // No lost work on either table; graph remains re-optimizable.
     let db = Omnigraph::open(&uri).await.unwrap();
     assert_eq!(
         helpers::count_rows(&db, "node:Person").await,
-        5,
-        "concurrent insert must not be lost",
+        4,
+        "Person compaction must preserve every seed row",
+    );
+    assert_eq!(
+        helpers::count_rows(&db, "node:Company").await,
+        1,
+        "queued table-disjoint Company insert must persist",
     );
     db.optimize()
         .await