store: Map other shards into a new shard on startup

Fixes #4719

Author: lutter

store/postgres/src/catalog.rs

@@ -415,14 +415,18 @@ pub fn drop_schema(conn: &PgConnection, nsp: &str) -> Result<(), StoreError> {
     Ok(conn.batch_execute(&query)?)
 }
 
-pub fn migration_count(conn: &PgConnection) -> Result<i64, StoreError> {
+pub fn migration_count(conn: &PgConnection) -> Result<usize, StoreError> {
     use __diesel_schema_migrations as m;
 
     if !table_exists(conn, NAMESPACE_PUBLIC, &MIGRATIONS_TABLE)? {
         return Ok(0);
     }
 
-    m::table.count().get_result(conn).map_err(StoreError::from)
+    m::table
+        .count()
+        .get_result(conn)
+        .map(|n: i64| n as usize)
+        .map_err(StoreError::from)
 }
 
 pub fn account_like(conn: &PgConnection, site: &Site) -> Result<HashSet<String>, StoreError> {

store/postgres/src/connection_pool.rs

@@ -1011,13 +1011,7 @@ impl PoolInner {
         let result = pool
             .configure_fdw(coord.servers.as_ref())
             .and_then(|()| migrate_schema(&pool.logger, &conn))
-            .and_then(|had_migrations| {
-                if had_migrations {
-                    coord.propagate_schema_change(&self.shard)
-                } else {
-                    Ok(())
-                }
-            });
+            .and_then(|count| coord.propagate(&pool, count));
         debug!(&pool.logger, "Release migration lock");
         advisory_lock::unlock_migration(&conn).unwrap_or_else(|err| {
             die(&pool.logger, "failed to release migration lock", &err);
@@ -1107,12 +1101,31 @@ impl PoolInner {
 
 embed_migrations!("./migrations");
 
+struct MigrationCount {
+    old: usize,
+    new: usize,
+}
+
+impl MigrationCount {
+    fn new(old: usize, new: usize) -> Self {
+        Self { old, new }
+    }
+
+    fn had_migrations(&self) -> bool {
+        self.old != self.new
+    }
+
+    fn is_new(&self) -> bool {
+        self.old == 0
+    }
+}
+
 /// Run all schema migrations.
 ///
 /// When multiple `graph-node` processes start up at the same time, we ensure
 /// that they do not run migrations in parallel by using `blocking_conn` to
 /// serialize them. The `conn` is used to run the actual migration.
-fn migrate_schema(logger: &Logger, conn: &PgConnection) -> Result<bool, StoreError> {
+fn migrate_schema(logger: &Logger, conn: &PgConnection) -> Result<MigrationCount, StoreError> {
     // Collect migration logging output
     let mut output = vec![];
 
@@ -1122,7 +1135,7 @@ fn migrate_schema(logger: &Logger, conn: &PgConnection) -> Result<bool, StoreErr
     let result = embedded_migrations::run_with_output(conn, &mut output);
     info!(logger, "Migrations finished");
 
-    let had_migrations = catalog::migration_count(conn)? != old_count;
+    let new_count = catalog::migration_count(conn)?;
 
     // If there was any migration output, log it now
     let msg = String::from_utf8(output).unwrap_or_else(|_| String::from("<unreadable>"));
@@ -1136,14 +1149,15 @@ fn migrate_schema(logger: &Logger, conn: &PgConnection) -> Result<bool, StoreErr
             debug!(logger, "Postgres migration output"; "output" => msg);
         }
     }
+    let count = MigrationCount::new(old_count, new_count);
 
-    if had_migrations {
+    if count.had_migrations() {
         // Reset the query statistics since a schema change makes them not
         // all that useful. An error here is not serious and can be ignored.
         conn.batch_execute("select pg_stat_statements_reset()").ok();
     }
 
-    Ok(had_migrations)
+    Ok(count)
 }
 
 /// Helper to coordinate propagating schema changes from the database that
@@ -1207,18 +1221,23 @@ impl PoolCoordinator {
 
     /// 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() {
-            if let Err(e) = pool.remap(server) {
-                error!(pool.logger, "Failed to map imports from {}", server.shard; "error" => e.to_string());
-                return Err(e);
+    /// `shard`. If `pool` is a new shard, we also map all other shards into
+    /// it.
+    fn propagate(&self, pool: &PoolInner, count: MigrationCount) -> Result<(), StoreError> {
+        // pool is a new shard, map all other shards into it
+        if count.is_new() {
+            for server in self.servers.iter() {
+                pool.remap(server)?;
+            }
+        }
+        // pool had schema changes, refresh the import from pool into all other shards
+        if count.had_migrations() {
+            let server = self.server(&pool.shard)?;
+            for pool in self.pools.lock().unwrap().values() {
+                if let Err(e) = pool.remap(server) {
+                    error!(pool.logger, "Failed to map imports from {}", server.shard; "error" => e.to_string());
+                    return Err(e);
+                }
             }
         }
         Ok(())
@@ -1231,4 +1250,11 @@ impl PoolCoordinator {
     pub fn servers(&self) -> Arc<Vec<ForeignServer>> {
         self.servers.clone()
     }
+
+    fn server(&self, shard: &Shard) -> Result<&ForeignServer, StoreError> {
+        self.servers
+            .iter()
+            .find(|server| &server.shard == shard)
+            .ok_or_else(|| constraint_violation!("unknown shard {shard}"))
+    }
 }