node, store: Properly propagate schema changes in a shard

Author: lutter

node/src/bin/manager.rs

@@ -21,6 +21,7 @@ use graph_node::{
     store_builder::StoreBuilder,
     MetricsContext,
 };
+use graph_store_postgres::connection_pool::PoolCoordinator;
 use graph_store_postgres::ChainStore;
 use graph_store_postgres::{
     connection_pool::ConnectionPool, BlockStore, NotificationSender, Shard, Store, SubgraphStore,
@@ -587,13 +588,14 @@ impl Context {
 
     fn primary_pool(self) -> ConnectionPool {
         let primary = self.config.primary_store();
+        let coord = Arc::new(PoolCoordinator::new(Arc::new(vec![])));
         let pool = StoreBuilder::main_pool(
             &self.logger,
             &self.node_id,
             PRIMARY_SHARD.as_str(),
             primary,
             self.metrics_registry(),
-            Arc::new(vec![]),
+            coord,
         );
         pool.skip_setup();
         pool

node/src/store_builder.rs

@@ -9,7 +9,9 @@ use graph::{
     prelude::{info, CheapClone, Logger},
     util::security::SafeDisplay,
 };
-use graph_store_postgres::connection_pool::{ConnectionPool, ForeignServer, PoolName};
+use graph_store_postgres::connection_pool::{
+    ConnectionPool, ForeignServer, PoolCoordinator, PoolName,
+};
 use graph_store_postgres::{
     BlockStore as DieselBlockStore, ChainHeadUpdateListener as PostgresChainHeadUpdateListener,
     NotificationSender, Shard as ShardName, Store as DieselStore, SubgraphStore,
@@ -102,6 +104,7 @@ impl StoreBuilder {
             .collect::<Result<Vec<_>, _>>()
             .expect("connection url's contain enough detail");
         let servers = Arc::new(servers);
+        let coord = Arc::new(PoolCoordinator::new(servers));
 
         let shards: Vec<_> = config
             .stores
@@ -114,7 +117,7 @@ impl StoreBuilder {
                     name,
                     shard,
                     registry.cheap_clone(),
-                    servers.clone(),
+                    coord.clone(),
                 );
 
                 let (read_only_conn_pools, weights) = Self::replica_pools(
@@ -123,7 +126,7 @@ impl StoreBuilder {
                     name,
                     shard,
                     registry.cheap_clone(),
-                    servers.clone(),
+                    coord.clone(),
                 );
 
                 let name =
@@ -191,7 +194,7 @@ impl StoreBuilder {
         name: &str,
         shard: &Shard,
         registry: Arc<dyn MetricsRegistry>,
-        servers: Arc<Vec<ForeignServer>>,
+        coord: Arc<PoolCoordinator>,
     ) -> ConnectionPool {
         let logger = logger.new(o!("pool" => "main"));
         let pool_size = shard.pool_size.size_for(node, name).expect(&format!(
@@ -209,15 +212,14 @@ impl StoreBuilder {
             "conn_pool_size" => pool_size,
             "weight" => shard.weight
         );
-        ConnectionPool::create(
+        coord.create_pool(
+            &logger,
             name,
             PoolName::Main,
             shard.connection.to_owned(),
             pool_size,
             Some(fdw_pool_size),
-            &logger,
             registry.cheap_clone(),
-            servers,
         )
     }
 
@@ -228,7 +230,7 @@ impl StoreBuilder {
         name: &str,
         shard: &Shard,
         registry: Arc<dyn MetricsRegistry>,
-        servers: Arc<Vec<ForeignServer>>,
+        coord: Arc<PoolCoordinator>,
     ) -> (Vec<ConnectionPool>, Vec<usize>) {
         let mut weights: Vec<_> = vec![shard.weight];
         (
@@ -250,15 +252,15 @@ impl StoreBuilder {
                         "we can determine the pool size for replica {}",
                         name
                     ));
-                    ConnectionPool::create(
+
+                    coord.clone().create_pool(
+                        &logger,
                         name,
                         PoolName::Replica(pool),
                         replica.connection.clone(),
                         pool_size,
                         None,
-                        &logger,
                         registry.cheap_clone(),
-                        servers.clone(),
                     )
                 })
                 .collect(),

store/postgres/src/connection_pool.rs

@@ -25,7 +25,7 @@ use graph::{
 
 use std::fmt::{self, Write};
 use std::sync::atomic::{AtomicBool, Ordering};
-use std::sync::Arc;
+use std::sync::{Arc, Mutex};
 use std::time::Duration;
 use std::{collections::HashMap, sync::RwLock};
 
@@ -158,9 +158,6 @@ impl ForeignServer {
 
     /// Map key tables from the primary into our local schema. If we are the
     /// primary, set them up as views.
-    ///
-    /// We recreate this mapping on every server start so that migrations that
-    /// change one of the mapped tables actually show up in the imported tables
     fn map_primary(conn: &PgConnection, shard: &Shard) -> Result<(), StoreError> {
         catalog::recreate_schema(conn, Self::PRIMARY_PUBLIC)?;
 
@@ -226,7 +223,7 @@ const FDW_IDLE_TIMEOUT: Duration = Duration::from_secs(60);
 enum PoolState {
     /// A connection pool, and all the servers for which we need to
     /// establish fdw mappings when we call `setup` on the pool
-    Created(Arc<PoolInner>, Arc<Vec<ForeignServer>>),
+    Created(Arc<PoolInner>, Arc<PoolCoordinator>),
     /// The pool has been successfully set up
     Ready(Arc<PoolInner>),
     /// The pool has been disabled by setting its size to 0
@@ -300,15 +297,15 @@ impl PoolStateTracker {
 }
 
 impl ConnectionPool {
-    pub fn create(
+    fn create(
         shard_name: &str,
         pool_name: PoolName,
         postgres_url: String,
         pool_size: u32,
         fdw_pool_size: Option<u32>,
         logger: &Logger,
         registry: Arc<dyn MetricsRegistry>,
-        servers: Arc<Vec<ForeignServer>>,
+        coord: Arc<PoolCoordinator>,
     ) -> ConnectionPool {
         let state_tracker = PoolStateTracker::new();
         let shard =
@@ -330,7 +327,7 @@ impl ConnectionPool {
                 if pool_name.is_replica() {
                     PoolState::Ready(Arc::new(pool))
                 } else {
-                    PoolState::Created(Arc::new(pool), servers)
+                    PoolState::Created(Arc::new(pool), coord)
                 }
             }
         };
@@ -968,7 +965,7 @@ impl PoolInner {
     /// # Panics
     ///
     /// If any errors happen during the migration, the process panics
-    pub fn setup(&self, servers: Arc<Vec<ForeignServer>>) -> Result<(), StoreError> {
+    fn setup(&self, coord: Arc<PoolCoordinator>) -> Result<(), StoreError> {
         fn die(logger: &Logger, msg: &'static str, err: &dyn std::fmt::Display) -> ! {
             crit!(logger, "{}", msg; "error" => format!("{:#}", err));
             panic!("{}: {}", msg, err);
@@ -980,11 +977,29 @@ impl PoolInner {
         let start = Instant::now();
         advisory_lock::lock_migration(&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
+        // pool B has at least finished `configure_fdw`, mapping tables will
+        // fail. In that case, the node must be restarted. The restart is
+        // guaranteed because this failure will lead to a panic in the setup
+        // for pool A
+        //
+        // This code can also leave the table mappings in a state where they
+        // have not been updated if the process is killed after migrating
+        // 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(servers.as_ref())
+            .configure_fdw(coord.servers.as_ref())
             .and_then(|()| migrate_schema(&pool.logger, &conn))
-            .and_then(|had_migrations| pool.map_primary())
-            .and_then(|()| pool.map_metadata(servers.as_ref()));
+            .and_then(|had_migrations| {
+                if had_migrations {
+                    coord.propagate_schema_change(&self.shard)
+                } else {
+                    Ok(())
+                }
+            });
         debug!(&pool.logger, "Release migration lock");
         advisory_lock::unlock_migration(&conn).unwrap_or_else(|err| {
             die(&pool.logger, "failed to release migration lock", &err);
@@ -1021,17 +1036,6 @@ impl PoolInner {
         })
     }
 
-    /// Map key tables from the primary into our local schema. If we are the
-    /// primary, set them up as views.
-    ///
-    /// We recreate this mapping on every server start so that migrations that
-    /// change one of the mapped tables actually show up in the imported tables
-    fn map_primary(&self) -> Result<(), StoreError> {
-        info!(&self.logger, "Mapping primary");
-        let conn = self.get()?;
-        conn.transaction(|| ForeignServer::map_primary(&conn, &self.shard))
-    }
-
     /// Copy the data from key tables in the primary into our local schema
     /// so it can be used as a fallback when the primary goes down
     pub async fn mirror_primary_tables(&self) -> Result<(), StoreError> {
@@ -1046,18 +1050,21 @@ impl PoolInner {
         .await
     }
 
-    // Map some tables from the `subgraphs` metadata schema from foreign
-    // servers to ourselves. The mapping is recreated on every server start
-    // so that we pick up possible schema changes in the mappings
-    fn map_metadata(&self, servers: &[ForeignServer]) -> Result<(), StoreError> {
-        info!(&self.logger, "Mapping metadata");
-        let conn = self.get()?;
-        conn.transaction(|| {
-            for server in servers.iter().filter(|server| server.shard != self.shard) {
-                server.map_metadata(&conn)?;
-            }
-            Ok(())
-        })
+    // The foreign server `server` had schema changes, and we therefore need
+    // to remap anything that we are importing via fdw to make sure we are
+    // using this updated schema
+    fn remap(&self, server: &ForeignServer) -> Result<(), StoreError> {
+        if &server.shard == &*PRIMARY_SHARD {
+            info!(&self.logger, "Mapping primary");
+            let conn = self.get()?;
+            conn.transaction(|| ForeignServer::map_primary(&conn, &self.shard))?;
+        }
+        if &server.shard != &self.shard {
+            info!(&self.logger, "Mapping metadata");
+            let conn = self.get()?;
+            conn.transaction(|| server.map_metadata(&conn))?;
+        }
+        Ok(())
     }
 }
 
@@ -1101,3 +1108,72 @@ fn migrate_schema(logger: &Logger, conn: &PgConnection) -> Result<bool, StoreErr
 
     Ok(had_migrations)
 }
+
+/// Helper to coordinate propagating schema changes from the database that
+/// changes schema to all other shards so they can update their fdw mappings
+/// of tables imported from that shard
+pub struct PoolCoordinator {
+    pools: Mutex<HashMap<Shard, Arc<PoolInner>>>,
+    servers: Arc<Vec<ForeignServer>>,
+}
+
+impl PoolCoordinator {
+    pub fn new(servers: Arc<Vec<ForeignServer>>) -> Self {
+        Self {
+            pools: Mutex::new(HashMap::new()),
+            servers,
+        }
+    }
+
+    pub fn create_pool(
+        self: Arc<Self>,
+        logger: &Logger,
+        name: &str,
+        pool_name: PoolName,
+        postgres_url: String,
+        pool_size: u32,
+        fdw_pool_size: Option<u32>,
+        registry: Arc<dyn MetricsRegistry>,
+    ) -> ConnectionPool {
+        let pool = ConnectionPool::create(
+            name,
+            pool_name,
+            postgres_url,
+            pool_size,
+            fdw_pool_size,
+            logger,
+            registry,
+            self.cheap_clone(),
+        );
+        // It is safe to take this lock here since nobody has seen the pool
+        // yet. We remember the `PoolInner` so that later, when we have to
+        // call `remap()`, we do not have to take this lock as that will be
+        // already held in `get_ready()`
+        match &*pool.inner.lock(logger) {
+            PoolState::Created(inner, _) | PoolState::Ready(inner) => {
+                self.pools
+                    .lock()
+                    .unwrap()
+                    .insert(pool.shard.clone(), inner.clone());
+            }
+            PoolState::Disabled => { /* nothing to do */ }
+        }
+        pool
+    }
+
+    /// Propagate changes to the schema in `shard` to all other pools. Those
+    /// other pools will then recreate any tables that they imported from
+    /// `shard`
+    fn propagate_schema_change(&self, shard: &Shard) -> Result<(), StoreError> {
+        let server = self
+            .servers
+            .iter()
+            .find(|server| &server.shard == shard)
+            .ok_or_else(|| constraint_violation!("unknown shard {shard}"))?;
+
+        for pool in self.pools.lock().unwrap().values() {
+            pool.remap(server)?;
+        }
+        Ok(())
+    }
+}