starred/windmill
1
0
Fork 0
mirror of https://github.com/windmill-labs/windmill.git synced 2026-09-21 00:02:30 +00:00
Code Issues Packages Projects Releases Wiki Activity

perf: build the audit operation index in a background migration

Browse Source 
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
This commit is contained in:
Ruben Fiszel
2026-09-18 20:54:01 +02:00
co-authored by Claude Opus 5
parent 70129b89f4
commit 73877adb33
4 changed files with 135 additions and 89 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/migrations/20260918173412_audit_partitioned_workspace_operation_index.down.sql
-1
View File
@@ -1 +0,0 @@
DROP INDEX IF EXISTS ix_audit_partitioned_workspace_operation;
backend/migrations/20260918173412_audit_partitioned_workspace_operation_index.up.sql
-6
View File
@@ -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");
backend/windmill-api/src/db.rs
+15 -80
View File
@@ -30,79 +30,6 @@ async fn current_database(conn: &mut PgConnection) -> Result<String, MigrateErro
.await?)
}
const AUDIT_OPERATION_INDEX_MIGRATION: i64 = 20260918173412;
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_concurrently(
conn: &mut PgConnection,
) -> 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<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
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<i64, String> = 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(
backend/windmill-api/src/live_migrations.rs
+120 -2
View File
@@ -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<bool, Error> {
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<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
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(())
}