Skip to content

refactor(api): fence Vault imports with work locks - #4419

Merged
chet merged 1 commit into
NVIDIA:mainfrom
chet:gh-issue-4393
Aug 1, 2026
Merged

refactor(api): fence Vault imports with work locks#4419
chet merged 1 commit into
NVIDIA:mainfrom
chet:gh-issue-4393

Conversation

@chet

@chet chet commented Jul 30, 2026

Copy link
Copy Markdown
Contributor

The one-time Vault import used to hold a dedicated PostgreSQL session from startup through every Vault and KMS call. Moving that coordination to WorkLockManager frees the checkout earlier, but its expiring leases mean the database write needs its own ownership fence.

So, this prepares the Vault/KMS data before opening a transaction, then commits every secret and the permanent marker behind the matching work_locks row. An expired owner cannot write after a replacement takes over, and a crash cannot leave a partially imported store.

Primary callouts are:

  • Use FOR KEY SHARE to block lease takeover and release while still allowing last_keepalive updates.
  • Keep the existing secrets:/_vault_import advisory lock as the mixed-version interlock while a legacy importer's detached lock session remains live.
  • Preserve advisory-first lock ordering for normal creates so mixed-version writers cannot deadlock the bulk table lock.
  • Use one secrets table lock, one existing-path query, and bind-limit-sized inserts so large batches avoid per-path lock and round-trip growth.
  • Make insert_if_missing acquire its table-write lock before checking the path, keeping normal creates ordered against the batch.
  • Treat a release failure after the import commits as lease cleanup trouble, not a failed import.

Tests added!

Related issues

This supports #4393

Type of Change

  • Add - New feature or capability
  • Change - Changes in existing functionality
  • Fix - Bug fixes
  • Remove - Removed features or deprecated functionality
  • Internal - Internal changes (refactoring, tests, docs, etc.)

Breaking Changes

  • This PR contains breaking changes

Testing

  • Unit tests added/updated
  • Integration tests added/updated
  • Manual testing performed
  • No testing required (docs, internal refactor, etc.)

The reviewed implementation passed:

cargo test --locked -p carbide-api-db work_lock_manager::tests --lib
cargo test --locked -p carbide-api-db batched_insert_splits_at_bind_limit --lib
cargo test --locked -p carbide-api-core vault_import_ --lib
cargo test --locked -p carbide-api-core bulk_secret_write_blocks_create_before_exists_check --lib
cargo test --locked -p carbide-api-core create_fails_when_credential_exists --lib
cargo test --locked -p carbide-api zero_database_pool_durations_are_rejected_at_startup --lib
cargo test --locked -p carbide-api-integration-tests --test secrets_import vault_to_postgres_import -- --nocapture
cargo make format-nightly
cargo make check-format-nightly
cargo make lint-error-messages
cargo make clippy
cargo make carbide-lints

The database regressions keep a fenced transaction open long enough to prove that keepalives continue while lease takeover waits, verify that the replacement acquires the lock and the stale owner can no longer fence or release it, and cross the SQL bind limit to exercise multi-statement batch insertion. The API-core coverage exercises takeover during KMS work, MissingOnly preservation, deletion during preparation, the permanent-marker no-op, and normal credential creates racing both the bulk table lock and a legacy path lock. The Vault/Postgres integration test imports and reads back 100 secrets through the batched path, then verifies the marker prevents another import.

Additional Notes

  • WorkLock::fence_transaction relies on the unique (work_key, worker_id) index: FOR KEY SHARE permits the non-key last_keepalive update while blocking the key-changing takeover and row deletion until commit.
  • The marker advisory lock is taken before the table lock. That keeps the new transaction coordinated with a legacy importer while its detached lock session remains live, without introducing a lock-order inversion.
  • An old importer cannot be fenced if its detached lock session dies while its pool-backed writes keep running, and an old credential writer does not know about the new bulk table-lock protocol. Deploy this code to every API replica before starting an import on a site without the marker, and keep credential rotation disabled until every replica runs a consistent writer configuration.
  • Documentation corrections remain isolated to a follow-up docs-only issue and PR, per repository policy.

Closes #4393

@chet
chet requested a review from a team as a code owner July 30, 2026 22:16
@coderabbitai

coderabbitai Bot commented Jul 30, 2026

Copy link
Copy Markdown
Contributor

Review Change Stack

Note

Reviews paused

It looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the reviews.auto_review.auto_pause_after_reviewed_commits setting.

Use the following commands to manage reviews:

  • @coderabbitai resume to resume automatic reviews.
  • @coderabbitai review to trigger a single review.

Use the checkboxes below for quick actions:

  • ▶️ Resume reviews
  • 🔍 Trigger review

Summary by CodeRabbit

  • New Features

    • Added coordinated one-time Vault secret imports across replicas using shared work locks.
    • Added bulk secret insertion and completion tracking for more efficient imports.
    • Added cancellation, retry, and confirmed lock-release handling.
  • Bug Fixes

    • Improved concurrent import handling by rechecking completion status and preventing stale writes.
  • Refactor

    • Replaced session-scoped advisory locking with shared, transaction-safe work-lock coordination.
  • Tests

    • Expanded coverage for contention, retries, cancellation, fencing, and bulk insertion.

Walkthrough

Vault imports now use shared WorkLockManager leases instead of session advisory locks. The flow adds cancellation-aware acquisition, transaction fencing, bulk writes, atomic completion marking, runtime handle wiring, and explicit release handling.

Changes

Vault import work-lock migration

Layer / File(s) Summary
Work-lock lifecycle and fencing
crates/api-db/src/work_lock_manager.rs
WorkLock now supports acknowledged release, conditional drop release, and stale-owner transaction fencing.
Bulk secret import and database coordination
crates/api-db/src/secrets.rs, crates/api-core/src/secrets/import.rs, crates/api-core/src/secrets/mod.rs, crates/api-core/src/secrets/tests.rs
Vault imports use bind-limited bulk insertion, existing-path lookup, table locking, atomic completion markers, and structured ImportResult outcomes.
Vault import lock acquisition
crates/api-core/src/bootstrap.rs
Acquisition uses a stable work key, completion-marker checks, contention retries, cancellation handling, and lock-release tests.
Runtime and import integration
crates/api/src/resources.rs, crates/api/src/run.rs, crates/api-core/src/setup.rs, crates/api-integration-tests/tests/secrets_import.rs
Resource setup starts the shared manager. The handle flows through runtime startup and import processing. Integration tests validate completed and already-completed imports.

Estimated code review effort: 4 (Complex) | ~60 minutes

Sequence Diagram(s)

sequenceDiagram
  participant setup_resources
  participant WorkLockManager
  participant import_vault_secrets_once
  participant Vault
  participant Database
  setup_resources->>WorkLockManager: start with shared database pool
  setup_resources->>import_vault_secrets_once: pass manager handle and cancellation token
  import_vault_secrets_once->>WorkLockManager: acquire Vault import work lock
  WorkLockManager-->>import_vault_secrets_once: lock acquired or completion detected
  import_vault_secrets_once->>Vault: enumerate and encrypt secrets
  import_vault_secrets_once->>Database: fence transaction and write secrets
  import_vault_secrets_once->>WorkLockManager: explicitly release lock
  WorkLockManager-->>import_vault_secrets_once: release result
Loading

Possibly related PRs

🚥 Pre-merge checks | ✅ 5
✅ Passed checks (5 passed)
Check name Status Explanation
Linked Issues check ✅ Passed The changes implement issue #4393 requirements, including manager propagation, stable locking, marker rechecks, fencing, retries, recovery, and tests.
Out of Scope Changes check ✅ Passed The code and test changes support Vault import coordination, work-lock fencing, bulk writes, and related compatibility requirements.
Docstring Coverage ✅ Passed Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.
Title check ✅ Passed The title clearly summarizes the main refactor: using work locks to fence Vault imports.
Description check ✅ Passed The description directly explains the work-lock refactor, fencing behavior, compatibility constraints, and associated tests.
✨ Finishing Touches
🧪 Generate unit tests (beta)
  • Create PR with unit tests

Comment @coderabbitai help to get the list of available commands.

chet commented Jul 30, 2026

Copy link
Copy Markdown
Contributor Author

@coderabbitai full_review, thanks!

@coderabbitai

coderabbitai Bot commented Jul 30, 2026

Copy link
Copy Markdown
Contributor

@chet: I’ll perform a complete review of the current PR changes.

🐇🔍

✅ Action performed

Full review finished.

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 1

🧹 Nitpick comments (1)
crates/api-core/src/bootstrap.rs (1)

185-217: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Consider collapsing the cfg-split wrapper into a single body.

The two branches differ only by a trailing None. Making contention_signal: Option<&Notify> unconditional (with Notify imported unconditionally) removes the duplicated call and keeps one signature to maintain, at the cost of one always-None argument in production. Optional, but it removes a seam that must be edited twice on every future parameter change.

♻️ Sketch
-pub async fn acquire_vault_import_work_lock(
-    db_pool: &PgPool,
-    work_lock_manager: &WorkLockManagerHandle,
-    cancel_token: &CancellationToken,
-) -> eyre::Result<Option<WorkLock>> {
-    #[cfg(not(test))]
-    {
-        acquire_vault_import_work_lock_with_retry(
-            db_pool,
-            work_lock_manager,
-            cancel_token,
-            VAULT_IMPORT_LOCK_RETRY_INTERVAL,
-        )
-        .await
-    }
-    #[cfg(test)]
-    {
-        acquire_vault_import_work_lock_with_retry(
-            db_pool,
-            work_lock_manager,
-            cancel_token,
-            VAULT_IMPORT_LOCK_RETRY_INTERVAL,
-            None,
-        )
-        .await
-    }
-}
+pub async fn acquire_vault_import_work_lock(
+    db_pool: &PgPool,
+    work_lock_manager: &WorkLockManagerHandle,
+    cancel_token: &CancellationToken,
+) -> eyre::Result<Option<WorkLock>> {
+    acquire_vault_import_work_lock_with_retry(
+        db_pool,
+        work_lock_manager,
+        cancel_token,
+        VAULT_IMPORT_LOCK_RETRY_INTERVAL,
+        None,
+    )
+    .await
+}

As per coding guidelines, "Prefer simple, explicit, readable Rust over clever or unnecessarily abstract code".

🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@crates/api-core/src/bootstrap.rs` around lines 185 - 217, Collapse the
cfg-split body of acquire_vault_import_work_lock into one call by importing
Notify unconditionally and passing a single contention_signal: Option<&Notify>
argument, using None for production and the existing test-specific signal
behavior. Keep the retry semantics unchanged while removing the duplicated
function invocation and maintaining one signature to update.

Source: Coding guidelines

🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@crates/api-db/src/work_lock_manager.rs`:
- Around line 1297-1303: Update the failure message in the row-count assertion
within the work-lock release check so it states that the release acknowledgment
occurred before deletion when row_count is nonzero. Keep the assertion condition
and query behavior unchanged.

---

Nitpick comments:
In `@crates/api-core/src/bootstrap.rs`:
- Around line 185-217: Collapse the cfg-split body of
acquire_vault_import_work_lock into one call by importing Notify unconditionally
and passing a single contention_signal: Option<&Notify> argument, using None for
production and the existing test-specific signal behavior. Keep the retry
semantics unchanged while removing the duplicated function invocation and
maintaining one signature to update.
🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

  • Push a commit to this branch (recommended)
  • Create a new PR with the fixes

ℹ️ Review info
⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 6b38a780-8514-4631-875a-c6e4c6066b16

📥 Commits

Reviewing files that changed from the base of the PR and between e4a4af9 and e558c3a.

📒 Files selected for processing (6)
  • crates/api-core/src/bootstrap.rs
  • crates/api-core/src/setup.rs
  • crates/api-db/src/secrets.rs
  • crates/api-db/src/work_lock_manager.rs
  • crates/api/src/resources.rs
  • crates/api/src/run.rs
💤 Files with no reviewable changes (1)
  • crates/api-db/src/secrets.rs

Comment thread crates/api-db/src/work_lock_manager.rs Outdated

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 2

🧹 Nitpick comments (3)
crates/api-core/src/bootstrap.rs (2)

185-225: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Collapse the cfg-split wrapper by making the contention hook unconditional.

The duplicated call body exists solely to satisfy the #[cfg(test)] parameter, and both branches independently repeat VAULT_IMPORT_LOCK_RETRY_INTERVAL — a future change to one branch will not be caught by a non-test build. Option<&Notify> is zero-cost when it is always None, so dropping the gate removes the seam entirely.

♻️ Proposed simplification
 pub async fn acquire_vault_import_work_lock(
     db_pool: &PgPool,
     work_lock_manager: &WorkLockManagerHandle,
     cancel_token: &CancellationToken,
 ) -> eyre::Result<Option<WorkLock>> {
-    #[cfg(not(test))]
-    {
-        acquire_vault_import_work_lock_with_retry(
-            db_pool,
-            work_lock_manager,
-            cancel_token,
-            VAULT_IMPORT_LOCK_RETRY_INTERVAL,
-        )
-        .await
-    }
-    #[cfg(test)]
-    {
-        acquire_vault_import_work_lock_with_retry(
-            db_pool,
-            work_lock_manager,
-            cancel_token,
-            VAULT_IMPORT_LOCK_RETRY_INTERVAL,
-            None,
-        )
-        .await
-    }
+    acquire_vault_import_work_lock_with_retry(
+        db_pool,
+        work_lock_manager,
+        cancel_token,
+        VAULT_IMPORT_LOCK_RETRY_INTERVAL,
+        None,
+    )
+    .await
 }
 
 async fn acquire_vault_import_work_lock_with_retry(
     db_pool: &PgPool,
     work_lock_manager: &WorkLockManagerHandle,
     cancel_token: &CancellationToken,
     retry_interval: Duration,
-    #[cfg(test)] contention_signal: Option<&Notify>,
+    contention_signal: Option<&Notify>,
 ) -> eyre::Result<Option<WorkLock>> {

The #[cfg(test)] guard on the Notify import and on the notify_one() call site would go away with it.

As per path instructions, "prefer simple explicit code, designs that are hard to misuse, justified abstractions".

🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@crates/api-core/src/bootstrap.rs` around lines 185 - 225, Remove the
cfg-split in acquire_vault_import_work_lock by making the contention_signal:
Option<&Notify> parameter unconditional and passing None from the single wrapper
call. Remove the corresponding cfg(test) guards from the Notify import and
notify_one() usage, while preserving the existing retry and contention behavior.

Source: Path instructions


279-300: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick win

Consider a periodic progress signal while waiting for the work key.

logged_contention limits the notice to a single info! on first contention. Since the wait is unbounded — the peer may legitimately hold the key for the whole Vault import — an operator investigating a replica that appears stuck at boot sees one line and no subsequent evidence of liveness. A re-log every N attempts (or an elapsed-time field) would make the wait diagnosable without adding log volume in the common fast path.

🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@crates/api-core/src/bootstrap.rs` around lines 279 - 300, The work-lock wait
loop around logged_contention currently logs contention only once; add periodic
progress logging during prolonged waits. Track retry attempts or elapsed wait
time and emit a throttled tracing::info message every N attempts or interval,
while preserving the existing first-contention log and avoiding additional logs
on fast-path waits.
crates/api/src/resources.rs (1)

486-529: 🩺 Stability & Availability | 🔵 Trivial

Lease loss during a long import is undetected by this critical section.

Per the caveat newly documented in crates/api-db/src/work_lock_manager.rs (Lines 163-167), the work key is a lease rather than a fencing token: if keepalives stall past the 60-second timeout while Vault enumeration or KMS wrapping is in flight, a peer may take the key and begin its own import while this block continues to completion. The keepalive loop emits work_lock_failed with failure=lock_lost, so the signal exists — it simply is not consumed here.

An alert on carbide_work_lock_failures_total{operation="keep_alive",failure="lock_lost"} for the import key, or observing the keepalive stop channel to abort the import before the marker write, would close the observability gap without changing the primitive.

🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@crates/api/src/resources.rs` around lines 486 - 529, The vault import
critical section does not react when its work-lock lease is lost during
enumeration or secret import. Connect the keepalive lock-loss signal to the
import flow around get_secrets_strict and import_secrets, aborting before
mark_vault_import_complete when the lease is lost; alternatively, emit an
operation-specific alert for the lock_lost metric. Preserve the existing
successful import and marker behavior when the lease remains valid.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@crates/api-db/src/work_lock_manager.rs`:
- Around line 778-781: Remove the “BUG:” prefix from the WorkLockManagerReply
error message, since run_loop can intentionally drop the reply sender after
failed connection reacquisition. Keep the existing error context and RecvError
conversion unchanged.

In `@crates/api/src/resources.rs`:
- Around line 454-460: Guard the initial import path around import_from with a
mechanical rollout precondition: when no completion marker exists, query the
latest observed deployed versions via
db::carbide_version::observe_as_latest_version and refuse to begin if any
recently observed replica is below the WorkLockManager version. Preserve normal
imports once all replicas support the new lock and allow already-completed
imports to proceed.

---

Nitpick comments:
In `@crates/api-core/src/bootstrap.rs`:
- Around line 185-225: Remove the cfg-split in acquire_vault_import_work_lock by
making the contention_signal: Option<&Notify> parameter unconditional and
passing None from the single wrapper call. Remove the corresponding cfg(test)
guards from the Notify import and notify_one() usage, while preserving the
existing retry and contention behavior.
- Around line 279-300: The work-lock wait loop around logged_contention
currently logs contention only once; add periodic progress logging during
prolonged waits. Track retry attempts or elapsed wait time and emit a throttled
tracing::info message every N attempts or interval, while preserving the
existing first-contention log and avoiding additional logs on fast-path waits.

In `@crates/api/src/resources.rs`:
- Around line 486-529: The vault import critical section does not react when its
work-lock lease is lost during enumeration or secret import. Connect the
keepalive lock-loss signal to the import flow around get_secrets_strict and
import_secrets, aborting before mark_vault_import_complete when the lease is
lost; alternatively, emit an operation-specific alert for the lock_lost metric.
Preserve the existing successful import and marker behavior when the lease
remains valid.
🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

  • Push a commit to this branch (recommended)
  • Create a new PR with the fixes

ℹ️ Review info
⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: d8b24bd0-56a1-4d3c-a615-38eef492bdde

📥 Commits

Reviewing files that changed from the base of the PR and between e4a4af9 and e558c3a.

📒 Files selected for processing (6)
  • crates/api-core/src/bootstrap.rs
  • crates/api-core/src/setup.rs
  • crates/api-db/src/secrets.rs
  • crates/api-db/src/work_lock_manager.rs
  • crates/api/src/resources.rs
  • crates/api/src/run.rs
💤 Files with no reviewable changes (1)
  • crates/api-db/src/secrets.rs

Comment thread crates/api-db/src/work_lock_manager.rs
Comment thread crates/api/src/resources.rs Outdated

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🧹 Nitpick comments (1)
crates/api-core/src/bootstrap.rs (1)

245-247: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Use the same wrap_err spelling here as in bootstrap.rs. PgSecretsError already fits wrap_err(...) directly, so map_err(eyre::Report::new) is redundant.

🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@crates/api-core/src/bootstrap.rs` around lines 245 - 247, Update the vault
import status handling around marker_status to use wrap_err directly on
is_vault_import_complete(db_pool).await, removing any redundant
map_err(eyre::Report::new) conversion and matching the existing bootstrap.rs
error-wrapping style.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Nitpick comments:
In `@crates/api-core/src/bootstrap.rs`:
- Around line 245-247: Update the vault import status handling around
marker_status to use wrap_err directly on
is_vault_import_complete(db_pool).await, removing any redundant
map_err(eyre::Report::new) conversion and matching the existing bootstrap.rs
error-wrapping style.

ℹ️ Review info
⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 0dae21fb-9713-4262-ae2a-ec3f1138cbed

📥 Commits

Reviewing files that changed from the base of the PR and between e558c3a and 6ad7ee4.

📒 Files selected for processing (6)
  • crates/api-core/src/bootstrap.rs
  • crates/api-core/src/setup.rs
  • crates/api-db/src/secrets.rs
  • crates/api-db/src/work_lock_manager.rs
  • crates/api/src/resources.rs
  • crates/api/src/run.rs
💤 Files with no reviewable changes (1)
  • crates/api-db/src/secrets.rs

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🧹 Nitpick comments (2)
crates/api-core/src/bootstrap.rs (2)

196-224: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Consider dropping the cfg-varying parameter arity in favour of an unconditional Option<&Notify>.

acquire_vault_import_work_lock_with_retry changes its signature under cfg(test), which forces the wrapper to duplicate its entire body across two cfg blocks. tokio::sync::Notify is not test-only, so the parameter can simply always be present and None in production — one call site, no divergent signatures, and no risk of the two branches drifting.

♻️ Suggested simplification
-use sqlx::PgPool;
-#[cfg(test)]
 use tokio::sync::Notify;
 pub async fn acquire_vault_import_work_lock(
     db_pool: &PgPool,
     work_lock_manager: &WorkLockManagerHandle,
     cancel_token: &CancellationToken,
 ) -> eyre::Result<Option<WorkLock>> {
-    #[cfg(not(test))]
-    {
-        acquire_vault_import_work_lock_with_retry(
-            db_pool,
-            work_lock_manager,
-            cancel_token,
-            VAULT_IMPORT_LOCK_RETRY_INTERVAL,
-        )
-        .await
-    }
-    #[cfg(test)]
-    {
-        acquire_vault_import_work_lock_with_retry(
-            db_pool,
-            work_lock_manager,
-            cancel_token,
-            VAULT_IMPORT_LOCK_RETRY_INTERVAL,
-            None,
-        )
-        .await
-    }
+    acquire_vault_import_work_lock_with_retry(
+        db_pool,
+        work_lock_manager,
+        cancel_token,
+        VAULT_IMPORT_LOCK_RETRY_INTERVAL,
+        None,
+    )
+    .await
 }
 
 async fn acquire_vault_import_work_lock_with_retry(
     db_pool: &PgPool,
     work_lock_manager: &WorkLockManagerHandle,
     cancel_token: &CancellationToken,
     retry_interval: Duration,
-    #[cfg(test)] contention_signal: Option<&Notify>,
+    contention_signal: Option<&Notify>,
 ) -> eyre::Result<Option<WorkLock>> {

The #[cfg(test)] guard on the notify_one() call inside the loop can then go as well.

As per coding guidelines, "Prefer simple, explicit, readable Rust over clever or unnecessarily abstract code; design APIs to be difficult to misuse".

🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@crates/api-core/src/bootstrap.rs` around lines 196 - 224, Make
acquire_vault_import_work_lock_with_retry always accept an Option<&Notify>,
removing the cfg(test) parameter guard. Simplify its wrapper to one
unconditional call, passing None in production and the existing test
notification in tests, and remove the cfg(test) guard around the loop’s
notify_one() call while preserving current behavior.

Source: Coding guidelines


227-306: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

The cancellation message is repeated four times and is inaccurate on two of those paths.

"vault import canceled while waiting for its work lock" appears at Lines 229, 242, 275 and 290. On Lines 242 and 275 the lock has already been acquired, so the operator is told the process was cancelled while waiting when it was in fact cancelled while holding the lease. Extracting a constant plus a distinct post-acquisition wording would keep the diagnostics honest without changing control flow (the existing test asserts only on the "vault import canceled" substring, so both variants remain covered).

🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@crates/api-core/src/bootstrap.rs` around lines 227 - 306, Update the
cancellation diagnostics in the lock-acquisition flow around try_acquire_lock:
define a shared message for cancellation while waiting, use it only before lock
acquisition, and add distinct wording for cancellation detected after acquiring
the work lock at the post-acquisition checks. Preserve the existing cancellation
control flow and the “vault import canceled” substring.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Nitpick comments:
In `@crates/api-core/src/bootstrap.rs`:
- Around line 196-224: Make acquire_vault_import_work_lock_with_retry always
accept an Option<&Notify>, removing the cfg(test) parameter guard. Simplify its
wrapper to one unconditional call, passing None in production and the existing
test notification in tests, and remove the cfg(test) guard around the loop’s
notify_one() call while preserving current behavior.
- Around line 227-306: Update the cancellation diagnostics in the
lock-acquisition flow around try_acquire_lock: define a shared message for
cancellation while waiting, use it only before lock acquisition, and add
distinct wording for cancellation detected after acquiring the work lock at the
post-acquisition checks. Preserve the existing cancellation control flow and the
“vault import canceled” substring.

ℹ️ Review info
⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 6ba7f4b1-fa4c-46d9-b543-529abfde80e3

📥 Commits

Reviewing files that changed from the base of the PR and between 6ad7ee4 and 76a78c6.

📒 Files selected for processing (6)
  • crates/api-core/src/bootstrap.rs
  • crates/api-core/src/setup.rs
  • crates/api-db/src/secrets.rs
  • crates/api-db/src/work_lock_manager.rs
  • crates/api/src/resources.rs
  • crates/api/src/run.rs
💤 Files with no reviewable changes (1)
  • crates/api-db/src/secrets.rs

@github-actions

Copy link
Copy Markdown

@chet chet changed the title refactor(api): use work locks for Vault import refactor(api): fence Vault imports with work locks Jul 31, 2026

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🧹 Nitpick comments (2)
crates/api-integration-tests/tests/secrets_import.rs (1)

344-346: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Bound the manager shutdown wait with a timeout.

The release-then-drop ordering is correct: release() consumes work_lock, so the keepalive task releases its cloned handle before the manager shuts down. The concern is the unbounded join_all(). If the run loop ever fails to terminate, this test hangs until the CI job timeout instead of failing with an actionable message. The equivalent unit test in crates/api-core/src/secrets/tests.rs (lines 613-616) already wraps the same call in a three-second tokio::time::timeout with an explicit expectation message. Align this test with that pattern.

♻️ Proposed bounded shutdown
     work_lock.release().await?;
     drop(work_lock_manager);
-    work_lock_tasks.join_all().await;
+    tokio::time::timeout(
+        std::time::Duration::from_secs(3),
+        work_lock_tasks.join_all(),
+    )
+    .await
+    .map_err(|_| eyre::eyre!("WorkLockManager did not shut down in a timely manner"))?;

Based on path instructions: "Review test-support crates for deterministic behavior, clear fixtures, reusable helpers, and avoiding sleeps or shared mutable global state that can make CI flaky."

🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@crates/api-integration-tests/tests/secrets_import.rs` around lines 344 - 346,
Bound the manager shutdown wait in the test cleanup after work_lock.release()
and drop(work_lock_manager) by wrapping work_lock_tasks.join_all().await in a
three-second tokio::time::timeout, and assert it completes with an explicit
failure message matching the existing unit-test pattern.

Source: Path instructions

crates/api-db/src/secrets.rs (1)

85-117: 🗄️ Data Integrity & Integration | 🔵 Trivial | ⚡ Quick win

Document that multi-chunk batches require a caller-supplied transaction.

insert_many accepts &mut PgConnection, so a caller can pass a pooled connection, as the new test does. In that case each chunk commits independently, and a mid-batch failure leaves a partial batch behind. For the Vault import this matters: the completion marker is one entry in the same batch, so a partial commit could publish a permanent marker without all secrets. The current production caller passes a transaction, so this is a documentation and misuse-prevention concern rather than an active defect.

♻️ Suggested documentation addition
 /// Append a batch of journal entries with as few INSERT statements as the
 /// PostgreSQL bind limit permits. Each row binds seven values, so oversized
 /// batches are split at [`BIND_LIMIT`] / 7 entries per statement.
+///
+/// A batch larger than one chunk spans several statements, so pass a
+/// transaction when the batch must be atomic.
 pub async fn insert_many(
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@crates/api-db/src/secrets.rs` around lines 85 - 117, Update the documentation
for insert_many to state that callers must provide a transaction when entries
may span multiple chunks, ensuring all chunks commit atomically; clarify that
passing a pooled PgConnection can commit chunks independently and leave a
partial batch after failure. Keep the existing implementation unchanged.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Nitpick comments:
In `@crates/api-db/src/secrets.rs`:
- Around line 85-117: Update the documentation for insert_many to state that
callers must provide a transaction when entries may span multiple chunks,
ensuring all chunks commit atomically; clarify that passing a pooled
PgConnection can commit chunks independently and leave a partial batch after
failure. Keep the existing implementation unchanged.

In `@crates/api-integration-tests/tests/secrets_import.rs`:
- Around line 344-346: Bound the manager shutdown wait in the test cleanup after
work_lock.release() and drop(work_lock_manager) by wrapping
work_lock_tasks.join_all().await in a three-second tokio::time::timeout, and
assert it completes with an explicit failure message matching the existing
unit-test pattern.

ℹ️ Review info
⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 8fe45dbf-cd53-4964-9d82-efb1b345a104

📥 Commits

Reviewing files that changed from the base of the PR and between e245f88 and dfe198f.

📒 Files selected for processing (10)
  • crates/api-core/src/bootstrap.rs
  • crates/api-core/src/secrets/import.rs
  • crates/api-core/src/secrets/mod.rs
  • crates/api-core/src/secrets/tests.rs
  • crates/api-core/src/setup.rs
  • crates/api-db/src/secrets.rs
  • crates/api-db/src/work_lock_manager.rs
  • crates/api-integration-tests/tests/secrets_import.rs
  • crates/api/src/resources.rs
  • crates/api/src/run.rs
🚧 Files skipped from review as they are similar to previous changes (3)
  • crates/api/src/run.rs
  • crates/api-core/src/setup.rs
  • crates/api-core/src/bootstrap.rs

The one-time Vault import used to hold a dedicated PostgreSQL session from startup through every Vault and KMS call. Moving that coordination to `WorkLockManager` frees the checkout earlier, but its expiring leases mean the database write needs its own ownership fence.

So, this prepares every Vault/KMS value before opening a transaction, then commits every secret and the permanent marker behind the matching `work_locks` row. An expired owner cannot write after a replacement takes over, and a crash cannot leave a partially imported store.

Primary callouts are:
- Use `FOR KEY SHARE` to block lease takeover and release while still allowing `last_keepalive` updates.
- Keep the existing `secrets:/_vault_import` advisory lock as the mixed-version interlock while a legacy importer's detached lock session remains live.
- Preserve advisory-first lock ordering for normal creates so mixed-version writers cannot deadlock the bulk table lock.
- Use one `secrets` table lock, one existing-path query, and bind-limit-sized inserts so large batches avoid per-path lock and round-trip growth.
- Make `insert_if_missing` acquire its table-write lock before checking the path, keeping normal creates ordered against the batch.
- Treat a release failure after the import commits as lease cleanup trouble, not a failed import.

Tests added!

This supports NVIDIA#4393

Signed-off-by: Chet Nichols III <chetn@nvidia.com>
@chet
chet enabled auto-merge (squash) August 1, 2026 00:12
@chet
chet merged commit 8bb1f2c into NVIDIA:main Aug 1, 2026
64 checks passed
@chet
chet deleted the gh-issue-4393 branch August 1, 2026 06:43
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Run Vault Imports Without Their Own Database Connection

3 participants