fix: use try_lock instead of lock to wait for global pg lock

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"
+}

backend/Cargo.lock

@@ -9462,6 +9462,7 @@ dependencies = [
  "chrono",
  "chrono-tz",
  "cookie",
+ "crc",
  "cron",
  "futures",
  "git-version",

backend/Cargo.toml

@@ -231,3 +231,5 @@ openidconnect = { version = "3.4.0" }
 
 aws-config = "^1"
 aws-sdk-sts = "^1"
+
+crc = "^3"

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
 

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<Postgres>;
 
+async fn current_database(conn: &mut PgConnection) -> Result<String, MigrateError> {
+    // 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<u32> = crc::Crc::<u32>::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<Postgres>,
 }
@@ -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()
     }