Automatic database role detection (#616)

Author: levkk

integration/role_detector/failover.sh

@@ -0,0 +1,3 @@
+#!/bin/bash
+docker stop $(docker ps | grep 45000 | awk '{print $1}')
+PGPASSWORD=postgres psql -h 127.0.0.1 -p 45001 -U postgres postgres -c 'SELECT pg_promote()'

integration/role_detector/pgdog.toml

@@ -0,0 +1,25 @@
+[general]
+lsn_check_delay = 0
+lsn_check_interval = 5_000
+idle_timeout = 1_000
+
+[[databases]]
+name = "postgres"
+role = "auto"
+host = "127.0.0.1"
+port = 45000
+
+[[databases]]
+name = "postgres"
+role = "auto"
+host = "127.0.0.1"
+port = 45001
+
+[[databases]]
+name = "postgres"
+role = "auto"
+host = "127.0.0.1"
+port = 45002
+
+[admin]
+password = "pgdog"

integration/role_detector/users.toml

@@ -0,0 +1,4 @@
+[[users]]
+name = "postgres"
+database = "postgres"
+password = "postgres"

pgdog-config/src/core.rs

@@ -5,7 +5,7 @@ use std::path::PathBuf;
 use tracing::{info, warn};
 
 use crate::sharding::ShardedSchema;
-use crate::{Memory, PassthoughAuth, PreparedStatements, RewriteMode};
+use crate::{EnumeratedDatabase, Memory, PassthoughAuth, PreparedStatements, RewriteMode};
 
 use super::database::Database;
 use super::error::Error;
@@ -191,8 +191,9 @@ pub struct Config {
 
 impl Config {
     /// Organize all databases by name for quicker retrieval.
-    pub fn databases(&self) -> HashMap<String, Vec<Vec<Database>>> {
+    pub fn databases(&self) -> HashMap<String, Vec<Vec<EnumeratedDatabase>>> {
         let mut databases = HashMap::new();
+        let mut number = 0;
         for database in &self.databases {
             let entry = databases
                 .entry(database.name.clone())
@@ -203,7 +204,11 @@ impl Config {
             entry
                 .get_mut(database.shard)
                 .unwrap()
-                .push(database.clone());
+                .push(EnumeratedDatabase {
+                    number,
+                    database: database.clone(),
+                });
+            number += 1;
         }
         databases
     }

pgdog-config/src/database.rs

@@ -1,5 +1,8 @@
 use serde::{Deserialize, Serialize};
-use std::str::FromStr;
+use std::{
+    ops::{Deref, DerefMut},
+    str::FromStr,
+};
 
 use super::pooling::PoolerMode;
 
@@ -153,3 +156,25 @@ impl FromStr for Role {
         }
     }
 }
+
+/// Database with a unique number, identifying it
+/// in the config.
+#[derive(Debug, Clone)]
+pub struct EnumeratedDatabase {
+    pub number: usize,
+    pub database: Database,
+}
+
+impl Deref for EnumeratedDatabase {
+    type Target = Database;
+
+    fn deref(&self) -> &Self::Target {
+        &self.database
+    }
+}
+
+impl DerefMut for EnumeratedDatabase {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        &mut self.database
+    }
+}

pgdog-config/src/lib.rs

@@ -19,7 +19,9 @@ pub mod util;
 pub use auth::{AuthType, PassthoughAuth};
 pub use core::{Config, ConfigAndUsers};
 pub use data_types::*;
-pub use database::{Database, LoadBalancingStrategy, ReadWriteSplit, ReadWriteStrategy, Role};
+pub use database::{
+    Database, EnumeratedDatabase, LoadBalancingStrategy, ReadWriteSplit, ReadWriteStrategy, Role,
+};
 pub use error::Error;
 pub use general::General;
 pub use memory::*;

pgdog/src/admin/show_replication.rs

@@ -32,7 +32,7 @@ impl Command for ShowReplication {
             Field::numeric("port"),
             Field::numeric("shard"),
             Field::text("role"),
-            Field::text("pg_replica_lag"),
+            Field::text("replica_lag"),
             Field::text("pg_lsn"),
             Field::text("lsn_age"),
             Field::text("pg_is_in_recovery"),

pgdog/src/admin/show_stats.rs

@@ -36,6 +36,7 @@ impl Command for ShowStats {
                         Field::numeric(&format!("{}_received", prefix)),
                         Field::numeric(&format!("{}_sent", prefix)),
                         Field::numeric(&format!("{}_xact_time", prefix)),
+                        Field::numeric(&format!("{}_idle_xact_time", prefix)),
                         Field::numeric(&format!("{}_query_time", prefix)),
                         Field::numeric(&format!("{}_wait_time", prefix)),
                         // Field::numeric(&format!("{}_client_parse_count", prefix)),
@@ -83,6 +84,7 @@ impl Command for ShowStats {
                             .add(stat.received)
                             .add(stat.sent)
                             .add(stat.xact_time.as_millis() as u64)
+                            .add(stat.idle_xact_time.as_millis() as u64)
                             .add(stat.query_time.as_millis() as u64)
                             .add(stat.wait_time.as_millis() as u64)
                             .add(stat.parse_count)

pgdog/src/backend/databases.rs

@@ -71,6 +71,17 @@ pub fn reconnect() -> Result<(), Error> {
     Ok(())
 }
 
+/// Re-create databases from existing config,
+/// preserving connections.
+pub fn reload_from_existing() {
+    let _lock = lock();
+
+    let config = config();
+    let databases = from_config(&config);
+
+    replace_databases(databases, true);
+}
+
 /// Initialize the databases for the first time.
 pub fn init() {
     let config = config();
@@ -180,6 +191,15 @@ impl ToUser for (&str, Option<&str>) {
     }
 }
 
+impl ToUser for &pgdog_config::User {
+    fn to_user(&self) -> User {
+        User {
+            user: self.name.clone(),
+            database: self.database.clone(),
+        }
+    }
+}
+
 /// Databases.
 #[derive(Default, Clone)]
 pub struct Databases {
@@ -354,95 +374,123 @@ pub(crate) fn new_pool(
     user: &crate::config::User,
     config: &crate::config::Config,
 ) -> Option<(User, Cluster)> {
+    let existing_roles = databases()
+        .cluster(user)
+        .ok()
+        .map(|cluster| cluster.redetect_roles());
+
     let sharded_tables = config.sharded_tables();
     let omnisharded_tables = config.omnisharded_tables();
     let sharded_mappings = config.sharded_mappings();
     let sharded_schemas = config.sharded_schemas();
     let general = &config.general;
     let databases = config.databases();
-    let shards = databases.get(&user.database);
-
-    if let Some(shards) = shards {
-        let mut shard_configs = vec![];
-        for user_databases in shards {
-            let has_single_replica = user_databases.len() == 1;
-            let primary = user_databases
-                .iter()
-                .find(|d| d.role == Role::Primary)
-                .map(|primary| PoolConfig {
-                    address: Address::new(primary, user),
-                    config: Config::new(general, primary, user, has_single_replica),
-                });
-            let replicas = user_databases
-                .iter()
-                .filter(|d| d.role == Role::Replica)
-                .map(|replica| PoolConfig {
-                    address: Address::new(replica, user),
-                    config: Config::new(general, replica, user, has_single_replica),
-                })
-                .collect::<Vec<_>>();
-
-            shard_configs.push(ClusterShardConfig { primary, replicas });
-        }
-
-        let mut sharded_tables = sharded_tables
-            .get(&user.database)
-            .cloned()
-            .unwrap_or_default();
-        let sharded_schemas = sharded_schemas
-            .get(&user.database)
-            .cloned()
-            .unwrap_or_default();
-
-        for sharded_table in &mut sharded_tables {
-            let mappings = sharded_mappings.get(&(
-                sharded_table.database.clone(),
-                sharded_table.column.clone(),
-                sharded_table.name.clone(),
-            ));
 
-            if let Some(mappings) = mappings {
-                sharded_table.mapping = Mapping::new(mappings);
+    let mut shards = if let Some(shards) = databases.get(&user.database).cloned() {
+        shards
+    } else {
+        return None;
+    };
 
-                if let Some(ref mapping) = sharded_table.mapping {
-                    if !mapping_valid(mapping) {
-                        warn!(
-                            "sharded table name=\"{}\", column=\"{}\" has overlapping ranges",
-                            sharded_table.name.as_ref().unwrap_or(&String::from("")),
-                            sharded_table.column
-                        );
+    let mut shard_configs = vec![];
+    for (shard_number, user_databases) in shards.iter_mut().enumerate() {
+        let role_detector = user_databases.iter().any(|d| d.role == Role::Auto);
+
+        if let Some(ref shard_roles) = existing_roles
+            .as_ref()
+            .map(|existing_roles| existing_roles.get(shard_number).cloned())
+            .flatten()
+            .flatten()
+        {
+            for user_database in user_databases.iter_mut() {
+                // Override role with automatically detected one.
+                if let Some(role) = shard_roles.get(&user_database.number) {
+                    if user_database.role == Role::Auto {
+                        user_database.role = role.role;
                     }
                 }
             }
         }
 
-        let omnisharded_tables = omnisharded_tables
-            .get(&user.database)
-            .cloned()
-            .unwrap_or(vec![]);
-        let sharded_tables = ShardedTables::new(sharded_tables, omnisharded_tables);
-        let sharded_schemas = ShardedSchemas::new(sharded_schemas);
-
-        let cluster_config = ClusterConfig::new(
-            general,
-            user,
-            &shard_configs,
-            sharded_tables,
-            config.multi_tenant(),
-            sharded_schemas,
-            &config.rewrite,
-        );
+        let has_single_replica = user_databases.len() == 1;
+        let primary = user_databases
+            .iter()
+            .find(|d| d.role == Role::Primary)
+            .map(|primary| PoolConfig {
+                address: Address::new(primary, user, primary.number),
+                config: Config::new(general, primary, user, has_single_replica),
+            });
+        let replicas = user_databases
+            .iter()
+            .filter(|d| matches!(d.role, Role::Replica | Role::Auto))
+            .map(|replica| PoolConfig {
+                address: Address::new(replica, user, replica.number),
+                config: Config::new(general, replica, user, has_single_replica),
+            })
+            .collect::<Vec<_>>();
+
+        shard_configs.push(ClusterShardConfig {
+            primary,
+            replicas,
+            role_detector,
+        });
+    }
 
-        Some((
-            User {
-                user: user.name.clone(),
-                database: user.database.clone(),
-            },
-            Cluster::new(cluster_config),
-        ))
-    } else {
-        None
+    let mut sharded_tables = sharded_tables
+        .get(&user.database)
+        .cloned()
+        .unwrap_or_default();
+    let sharded_schemas = sharded_schemas
+        .get(&user.database)
+        .cloned()
+        .unwrap_or_default();
+
+    for sharded_table in &mut sharded_tables {
+        let mappings = sharded_mappings.get(&(
+            sharded_table.database.clone(),
+            sharded_table.column.clone(),
+            sharded_table.name.clone(),
+        ));
+
+        if let Some(mappings) = mappings {
+            sharded_table.mapping = Mapping::new(mappings);
+
+            if let Some(ref mapping) = sharded_table.mapping {
+                if !mapping_valid(mapping) {
+                    warn!(
+                        "sharded table name=\"{}\", column=\"{}\" has overlapping ranges",
+                        sharded_table.name.as_ref().unwrap_or(&String::from("")),
+                        sharded_table.column
+                    );
+                }
+            }
+        }
     }
+
+    let omnisharded_tables = omnisharded_tables
+        .get(&user.database)
+        .cloned()
+        .unwrap_or(vec![]);
+    let sharded_tables = ShardedTables::new(sharded_tables, omnisharded_tables);
+    let sharded_schemas = ShardedSchemas::new(sharded_schemas);
+
+    let cluster_config = ClusterConfig::new(
+        general,
+        user,
+        &shard_configs,
+        sharded_tables,
+        config.multi_tenant(),
+        sharded_schemas,
+        &config.rewrite,
+    );
+
+    Some((
+        User {
+            user: user.name.clone(),
+            database: user.database.clone(),
+        },
+        Cluster::new(cluster_config),
+    ))
 }
 
 /// Load databases from config.