Clarify the code around locking migrations a bit

Author: lutter

graph/src/components/store/err.rs

@@ -141,7 +141,7 @@ impl Clone for StoreError {
 }
 
 impl StoreError {
-    fn from_diesel_error(e: &DieselError) -> Option<Self> {
+    pub fn from_diesel_error(e: &DieselError) -> Option<Self> {
         const CONN_CLOSE: &str = "server closed the connection unexpectedly";
         const STMT_TIMEOUT: &str = "canceling statement due to statement timeout";
         let DieselError::DatabaseError(_, info) = e else {

store/postgres/src/advisory_lock.rs

@@ -70,18 +70,43 @@ const WRITE: Scope = Scope { id: 2 };
 const PRUNE: Scope = Scope { id: 3 };
 
 /// Get a lock for running migrations. Blocks until we get the lock.
-pub(crate) fn lock_migration(conn: &mut PgConnection) -> Result<(), StoreError> {
-    sql_query("select pg_advisory_lock(1)").execute(conn)?;
+fn lock_migration(conn: &mut PgConnection) -> Result<(), StoreError> {
+    sql_query("select pg_advisory_lock(1)")
+        .execute(conn)
+        .map_err(|e| {
+            StoreError::from_diesel_error(&e).unwrap_or_else(|| {
+                StoreError::Unknown(anyhow::anyhow!("failed to acquire migration lock: {}", e))
+            })
+        })?;
 
     Ok(())
 }
 
 /// Release the migration lock.
-pub(crate) fn unlock_migration(conn: &mut PgConnection) -> Result<(), StoreError> {
-    sql_query("select pg_advisory_unlock(1)").execute(conn)?;
+fn unlock_migration(conn: &mut PgConnection) -> Result<(), StoreError> {
+    sql_query("select pg_advisory_unlock(1)")
+        .execute(conn)
+        .map_err(|e| {
+            StoreError::from_diesel_error(&e).unwrap_or_else(|| {
+                StoreError::Unknown(anyhow::anyhow!("failed to release migration lock: {}", e))
+            })
+        })?;
     Ok(())
 }
 
+/// Block until we can get the migration lock, then run `f` and unlock when
+/// it is done. This is used to make sure that only one node runs setup at a
+/// time.
+pub(crate) fn with_migration_lock<F, R>(conn: &mut PgConnection, f: F) -> Result<R, StoreError>
+where
+    F: FnOnce(&mut PgConnection) -> Result<R, StoreError>,
+{
+    lock_migration(conn)?;
+    let res = f(conn);
+    unlock_migration(conn)?;
+    res
+}
+
 /// Take the lock used to keep two copy operations to run simultaneously on
 /// the same deployment. Block until we can get the lock
 pub(crate) fn lock_copying(conn: &mut PgConnection, dst: &Site) -> Result<(), StoreError> {

store/postgres/src/connection_pool.rs

@@ -33,8 +33,9 @@ use std::{collections::HashMap, sync::RwLock};
 
 use postgres::config::{Config, Host};
 
+use crate::advisory_lock::with_migration_lock;
+use crate::catalog;
 use crate::primary::{self, NAMESPACE_PUBLIC};
-use crate::{advisory_lock, catalog};
 use crate::{Shard, PRIMARY_SHARD};
 
 /// Tables that we map from the primary into `primary_public` in each shard
@@ -1113,8 +1114,6 @@ impl PoolInner {
 
         let start = Instant::now();
 
-        advisory_lock::lock_migration(&mut conn)
-            .unwrap_or_else(|err| die(&pool.logger, "failed to get migration lock", &err));
         // This code can cause a race in database setup: if pool A has had
         // schema changes and pool B then tries to map tables from pool A,
         // but does so before the concurrent thread running this code for
@@ -1128,13 +1127,10 @@ impl PoolInner {
         // the schema but before finishing remapping in all shards.
         // Addressing that would require keeping track of the need to remap
         // in the database instead of just in memory
-        let result = pool
-            .configure_fdw(coord.servers.as_ref())
-            .and_then(|()| pool.drop_cross_shard_views())
-            .and_then(|()| migrate_schema(&pool.logger, &mut conn));
-        debug!(&pool.logger, "Release migration lock");
-        advisory_lock::unlock_migration(&mut conn).unwrap_or_else(|err| {
-            die(&pool.logger, "failed to release migration lock", &err);
+        let result = with_migration_lock(&mut conn, |conn| {
+            pool.configure_fdw(coord.servers.as_ref())
+                .and_then(|()| pool.drop_cross_shard_views())
+                .and_then(|()| migrate_schema(&pool.logger, conn))
         });
         let result = result
             .and_then(|count| coord.propagate(&pool, count))
@@ -1451,12 +1447,7 @@ impl PoolCoordinator {
             let server = self.server(&pool.shard)?;
             for pool in self.pools.lock().unwrap().values() {
                 let mut conn = pool.get()?;
-                let remap_res = {
-                    advisory_lock::lock_migration(&mut conn)?;
-                    let res = pool.remap(server);
-                    advisory_lock::unlock_migration(&mut conn)?;
-                    res
-                };
+                let remap_res = { with_migration_lock(&mut conn, |_| pool.remap(server)) };
                 if let Err(e) = remap_res {
                     error!(pool.logger, "Failed to map imports from {}", server.shard; "error" => e.to_string());
                     return Err(e);