From 73877adb339b58de16d01fb42e43e8a8b776341c Mon Sep 17 00:00:00 2001 From: Ruben Fiszel Date: Fri, 18 Sep 2026 20:54:01 +0200 Subject: [PATCH] perf: build the audit operation index in a background migration Co-Authored-By: Claude Opus 5 (1M context) --- ...itioned_workspace_operation_index.down.sql | 1 - ...rtitioned_workspace_operation_index.up.sql | 6 - backend/windmill-api/src/db.rs | 95 +++----------- backend/windmill-api/src/live_migrations.rs | 122 +++++++++++++++++- 4 files changed, 135 insertions(+), 89 deletions(-) delete mode 100644 backend/migrations/20260918173412_audit_partitioned_workspace_operation_index.down.sql delete mode 100644 backend/migrations/20260918173412_audit_partitioned_workspace_operation_index.up.sql diff --git a/backend/migrations/20260918173412_audit_partitioned_workspace_operation_index.down.sql b/backend/migrations/20260918173412_audit_partitioned_workspace_operation_index.down.sql deleted file mode 100644 index ebe897537f..0000000000 --- a/backend/migrations/20260918173412_audit_partitioned_workspace_operation_index.down.sql +++ /dev/null @@ -1 +0,0 @@ -DROP INDEX IF EXISTS ix_audit_partitioned_workspace_operation; diff --git a/backend/migrations/20260918173412_audit_partitioned_workspace_operation_index.up.sql b/backend/migrations/20260918173412_audit_partitioned_workspace_operation_index.up.sql deleted file mode 100644 index 4d07e66f5a..0000000000 --- a/backend/migrations/20260918173412_audit_partitioned_workspace_operation_index.up.sql +++ /dev/null @@ -1,6 +0,0 @@ --- The backend builds this per partition with CREATE INDEX CONCURRENTLY instead, --- see create_audit_operation_index_concurrently in windmill-api/src/db.rs. --- id matches list_audit's order and before_id cursor. timestamp trails it so a --- time window is checked in the index rather than on every row it reads. -CREATE INDEX IF NOT EXISTS ix_audit_partitioned_workspace_operation - ON audit_partitioned (workspace_id, operation, id DESC, "timestamp"); diff --git a/backend/windmill-api/src/db.rs b/backend/windmill-api/src/db.rs index 36dd4af1e1..9619e96ebc 100644 --- a/backend/windmill-api/src/db.rs +++ b/backend/windmill-api/src/db.rs @@ -30,79 +30,6 @@ async fn current_database(conn: &mut PgConnection) -> Result Result<(), MigrateError> { - conn.execute( - format!( - "CREATE INDEX IF NOT EXISTS ix_audit_partitioned_workspace_operation \ - ON ONLY audit_partitioned {AUDIT_OPERATION_INDEX_KEY}" - ) - .as_str(), - ) - .await?; - let partitions: Vec = 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 - AND NOT EXISTS ( - SELECT 1 FROM pg_inherits ip JOIN pg_index i ON i.indexrelid = ip.inhrelid - WHERE ip.inhparent = 'ix_audit_partitioned_workspace_operation'::regclass - AND i.indrelid = c.oid) - ORDER BY c.relname DESC", - ) - .fetch_all(&mut *conn) - .await?; - let quote = |name: &str| format!("\"{}\"", name.replace('"', "\"\"")); - for partition in partitions { - let index = quote(&format!( - "{partition}_workspace_id_operation_id_timestamp_idx" - )); - tracing::info!("Building ix_audit_partitioned_workspace_operation on {partition}"); - // An interrupted CONCURRENTLY build leaves an invalid index under this name. - conn.execute(format!("DROP INDEX CONCURRENTLY IF EXISTS {index}").as_str()) - .await?; - conn.execute( - format!( - "CREATE INDEX CONCURRENTLY {index} ON {} {AUDIT_OPERATION_INDEX_KEY}", - quote(&partition) - ) - .as_str(), - ) - .await?; - conn.execute( - format!( - "ALTER INDEX ix_audit_partitioned_workspace_operation ATTACH PARTITION {index}" - ) - .as_str(), - ) - .await?; - } - Ok(()) -} - -async fn record_overridden_migration( - conn: &mut PgConnection, - migration: &sqlx::migrate::Migration, -) -> Result<(), MigrateError> { - sqlx::query( - "INSERT INTO _sqlx_migrations (version, description, success, checksum, execution_time) - VALUES ($1, $2, TRUE, $3, -1) ON CONFLICT DO NOTHING", - ) - .bind(migration.version) - .bind(&*migration.description) - .bind(&*migration.checksum) - .execute(conn) - .await?; - Ok(()) -} - lazy_static::lazy_static! { pub static ref OVERRIDDEN_MIGRATIONS: std::collections::HashMap = vec![(20220123221903, include_str!( "../../migrations/20220123221903_first.up.sql" @@ -302,11 +229,6 @@ impl Migrate for CustomMigrator { migration.description ); - if migration.version == AUDIT_OPERATION_INDEX_MIGRATION { - create_audit_operation_index_concurrently(&mut self.inner).await?; - record_overridden_migration(&mut self.inner, migration).await?; - return Ok(std::time::Duration::from_secs(0)); - } if let Some(migration_sql) = OVERRIDDEN_MIGRATIONS.get(&migration.version) { tracing::info!("Using custom migration for version {}", migration.version); @@ -337,7 +259,17 @@ impl Migrate for CustomMigrator { } else if !migration_sql.is_empty() { self.inner.execute(&**migration_sql).await?; } - record_overridden_migration(&mut self.inner, migration).await?; + let _ = sqlx::query( + r#" + INSERT INTO _sqlx_migrations ( version, description, success, checksum, execution_time ) + VALUES ( $1, $2, TRUE, $3, -1 ) ON CONFLICT DO NOTHING + "#, + ) + .bind(migration.version) + .bind(&*migration.description) + .bind(&*migration.checksum) + .execute(&mut *self.inner) + .await?; return Ok(std::time::Duration::from_secs(0)); } else { let r = self.inner.apply(migration).await; @@ -440,7 +372,10 @@ pub async fn migrate( } crate::live_migrations::custom_migrations(&mut custom_migrator).await?; - Ok(None) + Ok(Some(crate::live_migrations::spawn_background_migrations( + db.clone(), + killpill_rx, + ))) } pub async fn wait_for_migrations( diff --git a/backend/windmill-api/src/live_migrations.rs b/backend/windmill-api/src/live_migrations.rs index 4c6f5e245c..4db42a24a7 100644 --- a/backend/windmill-api/src/live_migrations.rs +++ b/backend/windmill-api/src/live_migrations.rs @@ -6,8 +6,9 @@ * LICENSE-AGPL for a copy of the license. */ -use sqlx::Postgres; -use windmill_common::error::Error; +use sqlx::{PgConnection, Postgres}; +use tokio::task::JoinHandle; +use windmill_common::{db::DB, error::Error}; use crate::db::CustomMigrator; use sqlx::migrate::Migrate; @@ -121,3 +122,120 @@ async fn fix_flow_versioning_migration(migrator: &mut CustomMigrator) -> Result< } Ok(()) } + +// Held for the whole background run so only one server does it at a time. +const BACKGROUND_MIGRATIONS_LOCK_ID: i64 = 4_931_072_518_336_401; + +/// Schema changes too slow to hold server startup for, run by one server at a time after the +/// sqlx migrations. Each step is recorded in `windmill_migrations` once done; a step interrupted +/// by a restart or an error starts over on the next start and must resume safely. +pub fn spawn_background_migrations( + db: DB, + mut killpill_rx: tokio::sync::broadcast::Receiver<()>, +) -> JoinHandle<()> { + tokio::spawn(async move { + tokio::select! { + r = run_background_migrations(&db) => { + if let Err(err) = r { + tracing::error!("Background migrations stopped, retrying on the next start: {err:#}"); + } + } + _ = killpill_rx.recv() => { + tracing::info!("Killpill received, stopping background migrations"); + } + } + }) +} + +async fn run_background_migrations(db: &DB) -> Result<(), Error> { + // Detached so the pool's 5min statement_timeout, lifted here for the index builds, never + // comes back with this connection; closing it also releases the advisory lock. + let mut conn = db.acquire().await?.detach(); + conn.execute("SET statement_timeout = 0").await?; + let locked = sqlx::query_scalar::<_, bool>("SELECT pg_try_advisory_lock($1)") + .bind(BACKGROUND_MIGRATIONS_LOCK_ID) + .fetch_one(&mut conn) + .await?; + if !locked { + 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?; + } + Ok(()) +} + +async fn background_migration_done(conn: &mut PgConnection, name: &str) -> Result { + Ok(sqlx::query_scalar::<_, bool>( + "SELECT EXISTS(SELECT 1 FROM windmill_migrations WHERE name = $1)", + ) + .bind(name) + .fetch_one(conn) + .await?) +} + +async fn mark_background_migration_done(conn: &mut PgConnection, name: &str) -> Result<(), Error> { + sqlx::query("INSERT INTO windmill_migrations (name) VALUES ($1) ON CONFLICT DO NOTHING") + .bind(name) + .execute(conn) + .await?; + tracing::info!("Background migration {name} done"); + Ok(()) +} + +const AUDIT_OPERATION_INDEX_KEY: &str = r#"(workspace_id, operation, id DESC, "timestamp")"#; + +/// A plain `CREATE INDEX` on the partitioned table holds a SHARE lock on every partition until the +/// whole build ends, blocking the audit insert each job push makes in its own transaction. Each +/// 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}" + ) + .as_str(), + ) + .await?; + let partitions: Vec = 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 + AND NOT EXISTS ( + SELECT 1 FROM pg_inherits ip JOIN pg_index i ON i.indexrelid = ip.inhrelid + WHERE ip.inhparent = 'ix_audit_partitioned_workspace_operation'::regclass + AND i.indrelid = c.oid) + ORDER BY c.relname DESC", + ) + .fetch_all(&mut *conn) + .await?; + let quote = |name: &str| format!("\"{}\"", name.replace('"', "\"\"")); + for partition in partitions { + let index = quote(&format!( + "{partition}_workspace_id_operation_id_timestamp_idx" + )); + tracing::info!("Building ix_audit_partitioned_workspace_operation on {partition}"); + // An interrupted CONCURRENTLY build leaves an invalid index under this name. + conn.execute(format!("DROP INDEX CONCURRENTLY IF EXISTS {index}").as_str()) + .await?; + conn.execute( + format!( + "CREATE INDEX CONCURRENTLY {index} ON {} {AUDIT_OPERATION_INDEX_KEY}", + quote(&partition) + ) + .as_str(), + ) + .await?; + conn.execute( + format!( + "ALTER INDEX ix_audit_partitioned_workspace_operation ATTACH PARTITION {index}" + ) + .as_str(), + ) + .await?; + } + Ok(()) +}