From 87bbfc50a1b82d0670147a1e63350f9bcf9259a4 Mon Sep 17 00:00:00 2001 From: jif-oai Date: Wed, 11 Feb 2026 00:24:20 +0000 Subject: [PATCH] feat: prevent double backfill (#11377) ## Summary Add a DB-backed lease to prevent duplicate `.sqlite` backfill workers from running concurrently. ### What changed - Added StateRuntime::try_claim_backfill(lease_seconds) that atomically claims backfill only when: - backfill is not complete, and - no fresh running worker currently owns it. - Updated backfill_sessions to use the claim API and exit early when another worker already holds the lease. - Added runtime tests covering: - singleton claim behavior, - stale lease takeover, - claim blocked after complete. - Set backfill lease to 900s in production and 1s in tests. ### Why This avoids duplicate backfill work and reduces backfill status churn under concurrent startup, while preserving current best-effort fallback behavior. --- codex-rs/core/src/rollout/metadata.rs | 76 +++++++++++++++++++++---- codex-rs/state/src/runtime.rs | 80 +++++++++++++++++++++++++++ 2 files changed, 145 insertions(+), 11 deletions(-) diff --git a/codex-rs/core/src/rollout/metadata.rs b/codex-rs/core/src/rollout/metadata.rs index dc87c74da..5bfdc363a 100644 --- a/codex-rs/core/src/rollout/metadata.rs +++ b/codex-rs/core/src/rollout/metadata.rs @@ -30,6 +30,10 @@ use tracing::warn; const ROLLOUT_PREFIX: &str = "rollout-"; const ROLLOUT_SUFFIX: &str = ".jsonl"; const BACKFILL_BATCH_SIZE: usize = 200; +#[cfg(not(test))] +const BACKFILL_LEASE_SECONDS: i64 = 900; +#[cfg(test)] +const BACKFILL_LEASE_SECONDS: i64 = 1; pub(crate) fn builder_from_session_meta( session_meta: &SessionMetaLine, @@ -133,7 +137,7 @@ pub(crate) async fn backfill_sessions( otel: Option<&OtelManager>, ) { let timer = otel.and_then(|otel| otel.start_timer(DB_METRIC_BACKFILL_DURATION_MS, &[]).ok()); - let mut backfill_state = match runtime.get_backfill_state().await { + let backfill_state = match runtime.get_backfill_state().await { Ok(state) => state, Err(err) => { warn!( @@ -149,20 +153,66 @@ pub(crate) async fn backfill_sessions( if backfill_state.status == BackfillStatus::Complete { return; } - if let Err(err) = runtime.mark_backfill_running().await { - warn!( - "failed to mark backfill running at {}: {err}", + let claimed = match runtime.try_claim_backfill(BACKFILL_LEASE_SECONDS).await { + Ok(claimed) => claimed, + Err(err) => { + warn!( + "failed to claim backfill worker at {}: {err}", + config.codex_home.display() + ); + if let Some(otel) = otel { + otel.counter( + DB_ERROR_METRIC, + 1, + &[("stage", "backfill_state_claim_running")], + ); + } + return; + } + }; + if !claimed { + info!( + "state db backfill already running at {}; skipping duplicate worker", config.codex_home.display() ); - if let Some(otel) = otel { - otel.counter( - DB_ERROR_METRIC, - 1, - &[("stage", "backfill_state_mark_running")], + return; + } + let mut backfill_state = match runtime.get_backfill_state().await { + Ok(state) => state, + Err(err) => { + warn!( + "failed to read claimed backfill state at {}: {err}", + config.codex_home.display() ); + if let Some(otel) = otel { + otel.counter( + DB_ERROR_METRIC, + 1, + &[("stage", "backfill_state_read_claimed")], + ); + } + BackfillState { + status: BackfillStatus::Running, + ..Default::default() + } + } + }; + if backfill_state.status != BackfillStatus::Running { + if let Err(err) = runtime.mark_backfill_running().await { + warn!( + "failed to mark backfill running at {}: {err}", + config.codex_home.display() + ); + if let Some(otel) = otel { + otel.counter( + DB_ERROR_METRIC, + 1, + &[("stage", "backfill_state_mark_running")], + ); + } + } else { + backfill_state.status = BackfillStatus::Running; } - } else { - backfill_state.status = BackfillStatus::Running; } let sessions_root = config.codex_home.join(rollout::SESSIONS_SUBDIR); @@ -550,6 +600,10 @@ mod tests { .checkpoint_backfill(first_watermark.as_str()) .await .expect("checkpoint first watermark"); + tokio::time::sleep(std::time::Duration::from_secs( + (BACKFILL_LEASE_SECONDS + 1) as u64, + )) + .await; let mut config = crate::config::test_config(); config.codex_home = codex_home.clone(); diff --git a/codex-rs/state/src/runtime.rs b/codex-rs/state/src/runtime.rs index bd97979a5..541e02ad3 100644 --- a/codex-rs/state/src/runtime.rs +++ b/codex-rs/state/src/runtime.rs @@ -108,6 +108,34 @@ WHERE id = 1 crate::BackfillState::try_from_row(&row) } + /// Attempt to claim ownership of rollout metadata backfill. + /// + /// Returns `true` when this runtime claimed the backfill worker slot. + /// Returns `false` if backfill is already complete or currently owned by a + /// non-expired worker. + pub async fn try_claim_backfill(&self, lease_seconds: i64) -> anyhow::Result { + self.ensure_backfill_state_row().await?; + let now = Utc::now().timestamp(); + let lease_cutoff = now.saturating_sub(lease_seconds.max(0)); + let result = sqlx::query( + r#" +UPDATE backfill_state +SET status = ?, updated_at = ? +WHERE id = 1 + AND status != ? + AND (status != ? OR updated_at <= ?) + "#, + ) + .bind(crate::BackfillStatus::Running.as_str()) + .bind(now) + .bind(crate::BackfillStatus::Complete.as_str()) + .bind(crate::BackfillStatus::Running.as_str()) + .bind(lease_cutoff) + .execute(self.pool.as_ref()) + .await?; + Ok(result.rows_affected() == 1) + } + /// Mark rollout metadata backfill as running. pub async fn mark_backfill_running(&self) -> anyhow::Result<()> { self.ensure_backfill_state_row().await?; @@ -1028,6 +1056,58 @@ mod tests { let _ = tokio::fs::remove_dir_all(codex_home).await; } + #[tokio::test] + async fn backfill_claim_is_singleton_until_stale_and_blocked_when_complete() { + let codex_home = unique_temp_dir(); + let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string(), None) + .await + .expect("initialize runtime"); + + let claimed = runtime + .try_claim_backfill(3600) + .await + .expect("initial backfill claim"); + assert_eq!(claimed, true); + + let duplicate_claim = runtime + .try_claim_backfill(3600) + .await + .expect("duplicate backfill claim"); + assert_eq!(duplicate_claim, false); + + let stale_updated_at = Utc::now().timestamp().saturating_sub(10_000); + sqlx::query( + r#" +UPDATE backfill_state +SET status = ?, updated_at = ? +WHERE id = 1 + "#, + ) + .bind(crate::BackfillStatus::Running.as_str()) + .bind(stale_updated_at) + .execute(runtime.pool.as_ref()) + .await + .expect("force stale backfill lease"); + + let stale_claim = runtime + .try_claim_backfill(10) + .await + .expect("stale backfill claim"); + assert_eq!(stale_claim, true); + + runtime + .mark_backfill_complete(None) + .await + .expect("mark complete"); + let claim_after_complete = runtime + .try_claim_backfill(3600) + .await + .expect("claim after complete"); + assert_eq!(claim_after_complete, false); + + let _ = tokio::fs::remove_dir_all(codex_home).await; + } + #[tokio::test] async fn stage1_claim_skips_when_up_to_date() { let codex_home = unique_temp_dir();