From 6320cc9ce5e0fffe4f48471154774af611cbc39b Mon Sep 17 00:00:00 2001 From: Ruben Fiszel Date: Thu, 14 Mar 2024 14:28:13 +0100 Subject: [PATCH] fix: use try_lock instead of lock to wait for global pg lock --- ...14774f829b37d69f87286bd192af9cf702ac4.json | 22 +++++++ backend/Cargo.lock | 1 + backend/Cargo.toml | 2 + backend/windmill-api/Cargo.toml | 1 + backend/windmill-api/src/db.rs | 62 +++++++++++++++++-- 5 files changed, 83 insertions(+), 5 deletions(-) create mode 100644 backend/.sqlx/query-96724ea1050e71438f7b892254514774f829b37d69f87286bd192af9cf702ac4.json diff --git a/backend/.sqlx/query-96724ea1050e71438f7b892254514774f829b37d69f87286bd192af9cf702ac4.json b/backend/.sqlx/query-96724ea1050e71438f7b892254514774f829b37d69f87286bd192af9cf702ac4.json new file mode 100644 index 0000000000..7f99c8770a --- /dev/null +++ b/backend/.sqlx/query-96724ea1050e71438f7b892254514774f829b37d69f87286bd192af9cf702ac4.json @@ -0,0 +1,22 @@ +{ + "db_name": "PostgreSQL", + "query": "SELECT pg_try_advisory_lock($1)", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "pg_try_advisory_lock", + "type_info": "Bool" + } + ], + "parameters": { + "Left": [ + "Int8" + ] + }, + "nullable": [ + null + ] + }, + "hash": "96724ea1050e71438f7b892254514774f829b37d69f87286bd192af9cf702ac4" +} diff --git a/backend/Cargo.lock b/backend/Cargo.lock index 77c33d14d4..b863685b3a 100644 --- a/backend/Cargo.lock +++ b/backend/Cargo.lock @@ -9462,6 +9462,7 @@ dependencies = [ "chrono", "chrono-tz", "cookie", + "crc", "cron", "futures", "git-version", diff --git a/backend/Cargo.toml b/backend/Cargo.toml index 71568f9300..21820139d3 100644 --- a/backend/Cargo.toml +++ b/backend/Cargo.toml @@ -231,3 +231,5 @@ openidconnect = { version = "3.4.0" } aws-config = "^1" aws-sdk-sts = "^1" + +crc = "^3" diff --git a/backend/windmill-api/Cargo.toml b/backend/windmill-api/Cargo.toml index 563c5ff6c2..805e15a3cd 100644 --- a/backend/windmill-api/Cargo.toml +++ b/backend/windmill-api/Cargo.toml @@ -87,4 +87,5 @@ polars-io = { workspace = true, optional = true} object_store = { workspace = true, optional = true} openidconnect = { workspace = true} pin-project.workspace = true +crc.workspace = true diff --git a/backend/windmill-api/src/db.rs b/backend/windmill-api/src/db.rs index eb8939e447..32f075129f 100644 --- a/backend/windmill-api/src/db.rs +++ b/backend/windmill-api/src/db.rs @@ -10,7 +10,11 @@ use futures::FutureExt; #[cfg(feature = "enterprise")] use sqlx::Executor; -use sqlx::{migrate::Migrate, pool::PoolConnection, Pool, Postgres}; +use sqlx::{ + migrate::{Migrate, MigrateError}, + pool::PoolConnection, + PgConnection, Pool, Postgres, +}; use windmill_common::{ db::{Authable, Authed}, error::Error, @@ -18,6 +22,20 @@ use windmill_common::{ pub type DB = Pool; +async fn current_database(conn: &mut PgConnection) -> Result { + // language=SQL + Ok(sqlx::query_scalar("SELECT current_database()") + .fetch_one(conn) + .await?) +} + +// inspired from rails: https://github.com/rails/rails/blob/6e49cc77ab3d16c06e12f93158eaf3e507d4120e/activerecord/lib/active_record/migration.rb#L1308 +fn generate_lock_id(database_name: &str) -> i64 { + const CRC_IEEE: crc::Crc = crc::Crc::::new(&crc::CRC_32_ISO_HDLC); + // 0x3d32ad9e chosen by fair dice roll + 0x3d32ad9e * (CRC_IEEE.checksum(database_name.as_bytes()) as i64) +} + struct CustomMigrator { inner: PoolConnection, } @@ -48,13 +66,37 @@ impl Migrate for CustomMigrator { &mut self, ) -> futures::prelude::future::BoxFuture<'_, Result<(), sqlx::migrate::MigrateError>> { async { + if std::env::var("SKIP_PG_LOCK").is_ok() { + tracing::info!("Skipping PG lock acquisition"); + return Ok(()); + } + let pid = sqlx::query_scalar!("SELECT pg_backend_pid()") .fetch_one(&mut *self.inner) .await?; tracing::info!("Acquiring global PG lock for potential migration with pid: {pid:?}"); - let r = self.inner.lock().await; + let database_name = current_database(&mut *self.inner).await?; + let lock_id = generate_lock_id(&database_name); + + let mut r = false; + + while !r { + r = sqlx::query_scalar!("SELECT pg_try_advisory_lock($1)", lock_id) + .fetch_one(&mut *self.inner) + .await + .map_err(|e| { + tracing::error!("Error acquiring lock: {e}"); + sqlx::migrate::MigrateError::Execute(e) + })? + .unwrap_or(false); + if !r { + tracing::info!("PG lock already acquired by another server or worker, retrying in 5s. (look for the advisory lock in pg_lock with granted = true)"); + tokio::time::sleep(std::time::Duration::from_secs(5)).await; + } + } tracing::info!("Acquired global PG lock"); - r + + return Ok(()); } .boxed() } @@ -63,10 +105,20 @@ impl Migrate for CustomMigrator { &mut self, ) -> futures::prelude::future::BoxFuture<'_, Result<(), sqlx::migrate::MigrateError>> { async { + if std::env::var("SKIP_PG_UNLOCK").is_ok() { + tracing::info!("Skipping PG lock release"); + return Ok(()); + } tracing::info!("Releasing PG lock"); - let r = self.inner.unlock().await; + let database_name = current_database(&mut *self.inner).await?; + let lock_id = generate_lock_id(&database_name); + let _ = sqlx::query("SELECT pg_advisory_unlock($1)") + .bind(lock_id) + .execute(&mut *self.inner) + .await?; + tracing::info!("Released PG lock"); - r + Ok(()) } .boxed() }