fix: bound the audit index parent's lock wait and retry background steps

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

backend/windmill-api/src/live_migrations.rs

@@ -160,14 +160,46 @@ async fn run_background_migrations(db: &DB) -> Result<(), Error> {
         return Ok(());
     }
 
-    const AUDIT_OPERATION_INDEX: &str = "audit_partitioned_workspace_operation_index";
-    if !background_migration_done(&mut conn, AUDIT_OPERATION_INDEX).await? {
-        create_audit_operation_index(&mut conn).await?;
-        mark_background_migration_done(&mut conn, AUDIT_OPERATION_INDEX).await?;
+    for step in [AUDIT_OPERATION_INDEX] {
+        if background_migration_done(&mut conn, step).await? {
+            continue;
+        }
+        // Steps bound the lock waits that would queue audit inserts behind them with
+        // lock_timeout, and resume where they stopped, so a lock timeout is retried here.
+        let mut attempt = 1;
+        loop {
+            let run = match step {
+                AUDIT_OPERATION_INDEX => create_audit_operation_index(&mut conn).await,
+                _ => unreachable!("background migration {step} has no step function"),
+            };
+            match run {
+                Ok(()) => break,
+                Err(err) if attempt < 10 && is_lock_timeout(&err) => {
+                    tracing::warn!(
+                        "Background migration {step} timed out on a lock, retrying in 30s: {err:#}"
+                    );
+                    attempt += 1;
+                    tokio::time::sleep(std::time::Duration::from_secs(30)).await;
+                }
+                Err(err) => return Err(err),
+            }
+        }
+        mark_background_migration_done(&mut conn, step).await?;
     }
     Ok(())
 }
 
+const AUDIT_OPERATION_INDEX: &str = "audit_partitioned_workspace_operation_index";
+
+// Short on purpose: a statement waiting for its lock is also a wait for every audit insert
+// queued behind it.
+const STEP_LOCK_TIMEOUT: &str = "SET lock_timeout = '1s'";
+
+fn is_lock_timeout(err: &Error) -> bool {
+    matches!(err, Error::SqlErr { error: sqlx::Error::Database(db_err), .. }
+        if db_err.code().as_deref() == Some("55P03"))
+}
+
 async fn background_migration_done(conn: &mut PgConnection, name: &str) -> Result<bool, Error> {
     Ok(sqlx::query_scalar::<_, bool>(
         "SELECT EXISTS(SELECT 1 FROM windmill_migrations WHERE name = $1)",
@@ -193,14 +225,20 @@ const AUDIT_OPERATION_INDEX_KEY: &str = r#"(workspace_id, operation, id DESC, "t
 /// partition is built CONCURRENTLY instead and attached to a parent created `ON ONLY`, which turns
 /// valid once all partitions are attached. Partitions created later get the index from the parent.
 async fn create_audit_operation_index(conn: &mut PgConnection) -> Result<(), Error> {
-    conn.execute(
-        format!(
-            "CREATE INDEX IF NOT EXISTS ix_audit_partitioned_workspace_operation \
-             ON ONLY audit_partitioned {AUDIT_OPERATION_INDEX_KEY}"
+    // Metadata only, but it waits for every open transaction that wrote an audit row, and every
+    // audit insert (so every job push) queues behind it meanwhile. Give up early instead.
+    conn.execute(STEP_LOCK_TIMEOUT).await?;
+    let created = conn
+        .execute(
+            format!(
+                "CREATE INDEX IF NOT EXISTS ix_audit_partitioned_workspace_operation \
+                 ON ONLY audit_partitioned {AUDIT_OPERATION_INDEX_KEY}"
+            )
+            .as_str(),
         )
-        .as_str(),
-    )
-    .await?;
+        .await;
+    conn.execute("RESET lock_timeout").await?;
+    created?;
     let partitions: Vec<String> = sqlx::query_scalar(
         "SELECT c.relname::text FROM pg_inherits p JOIN pg_class c ON c.oid = p.inhrelid
          WHERE p.inhparent = 'audit_partitioned'::regclass