Replica lag monitor (#613)

Author: levkk

integration/load_balancer/pgdog.toml

@@ -16,6 +16,7 @@ pooler_mode = "transaction"
 load_balancing_strategy = "round_robin"
 auth_type = "trust"
 read_write_split = "exclude_primary"
+lsn_check_delay = 5_000
 
 [rewrite]
 enabled = false

pgdog-config/src/database.rs

@@ -128,13 +128,15 @@ pub enum Role {
     #[default]
     Primary,
     Replica,
+    Auto,
 }
 
 impl std::fmt::Display for Role {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         match self {
             Self::Primary => write!(f, "primary"),
             Self::Replica => write!(f, "replica"),
+            Self::Auto => write!(f, "auto"),
         }
     }
 }
@@ -146,6 +148,7 @@ impl FromStr for Role {
         match s.to_lowercase().as_str() {
             "primary" => Ok(Self::Primary),
             "replica" => Ok(Self::Replica),
+            "auto" => Ok(Self::Auto),
             _ => Err(format!("Invalid role: {}", s)),
         }
     }

pgdog-config/src/general.rs

@@ -172,6 +172,15 @@ pub struct General {
     /// Connection cleanup algorithm.
     #[serde(default = "General::connection_recovery")]
     pub connection_recovery: ConnectionRecovery,
+    /// LSN check interval.
+    #[serde(default = "General::lsn_check_interval")]
+    pub lsn_check_interval: u64,
+    /// LSN check timeout.
+    #[serde(default = "General::lsn_check_timeout")]
+    pub lsn_check_timeout: u64,
+    /// LSN check delay.
+    #[serde(default = "General::lsn_check_delay")]
+    pub lsn_check_delay: u64,
 }
 
 impl Default for General {
@@ -233,6 +242,9 @@ impl Default for General {
             server_lifetime: Self::server_lifetime(),
             stats_period: Self::stats_period(),
             connection_recovery: Self::connection_recovery(),
+            lsn_check_interval: Self::lsn_check_interval(),
+            lsn_check_timeout: Self::lsn_check_timeout(),
+            lsn_check_delay: Self::lsn_check_delay(),
         }
     }
 }
@@ -398,6 +410,18 @@ impl General {
         Self::env_enum_or_default("PGDOG_POOLER_MODE")
     }
 
+    fn lsn_check_timeout() -> u64 {
+        Self::env_or_default("PGDOG_LSN_CHECK_TIMEOUT", 5_000)
+    }
+
+    fn lsn_check_interval() -> u64 {
+        Self::env_or_default("PGDOG_LSN_CHECK_INTERVAL", 5_000)
+    }
+
+    fn lsn_check_delay() -> u64 {
+        Self::env_or_default("PGDOG_LSN_CHECK_DELAY", Duration::MAX.as_millis() as u64)
+    }
+
     fn read_write_strategy() -> ReadWriteStrategy {
         Self::env_enum_or_default("PGDOG_READ_WRITE_STRATEGY")
     }

pgdog/src/admin/mod.rs

@@ -29,6 +29,7 @@ pub mod show_peers;
 pub mod show_pools;
 pub mod show_prepared_statements;
 pub mod show_query_cache;
+pub mod show_replication;
 pub mod show_server_memory;
 pub mod show_servers;
 pub mod show_stats;

pgdog/src/admin/parser.rs

@@ -7,9 +7,10 @@ use super::{
     show_client_memory::ShowClientMemory, show_clients::ShowClients, show_config::ShowConfig,
     show_instance_id::ShowInstanceId, show_lists::ShowLists, show_mirrors::ShowMirrors,
     show_peers::ShowPeers, show_pools::ShowPools, show_prepared_statements::ShowPreparedStatements,
-    show_query_cache::ShowQueryCache, show_server_memory::ShowServerMemory,
-    show_servers::ShowServers, show_stats::ShowStats, show_transactions::ShowTransactions,
-    show_version::ShowVersion, shutdown::Shutdown, Command, Error,
+    show_query_cache::ShowQueryCache, show_replication::ShowReplication,
+    show_server_memory::ShowServerMemory, show_servers::ShowServers, show_stats::ShowStats,
+    show_transactions::ShowTransactions, show_version::ShowVersion, shutdown::Shutdown, Command,
+    Error,
 };
 
 use tracing::debug;
@@ -35,6 +36,7 @@ pub enum ParseResult {
     Shutdown(Shutdown),
     ShowLists(ShowLists),
     ShowPrepared(ShowPreparedStatements),
+    ShowReplication(ShowReplication),
     ShowServerMemory(ShowServerMemory),
     ShowClientMemory(ShowClientMemory),
     Set(Set),
@@ -69,6 +71,7 @@ impl ParseResult {
             Shutdown(shutdown) => shutdown.execute().await,
             ShowLists(show_lists) => show_lists.execute().await,
             ShowPrepared(cmd) => cmd.execute().await,
+            ShowReplication(show_replication) => show_replication.execute().await,
             ShowServerMemory(show_server_memory) => show_server_memory.execute().await,
             ShowClientMemory(show_client_memory) => show_client_memory.execute().await,
             Set(set) => set.execute().await,
@@ -103,6 +106,7 @@ impl ParseResult {
             Shutdown(shutdown) => shutdown.name(),
             ShowLists(show_lists) => show_lists.name(),
             ShowPrepared(show) => show.name(),
+            ShowReplication(show_replication) => show_replication.name(),
             ShowServerMemory(show_server_memory) => show_server_memory.name(),
             ShowClientMemory(show_client_memory) => show_client_memory.name(),
             Set(set) => set.name(),
@@ -158,6 +162,7 @@ impl Parser {
                 "instance_id" => ParseResult::ShowInstanceId(ShowInstanceId::parse(&sql)?),
                 "lists" => ParseResult::ShowLists(ShowLists::parse(&sql)?),
                 "prepared" => ParseResult::ShowPrepared(ShowPreparedStatements::parse(&sql)?),
+                "replication" => ParseResult::ShowReplication(ShowReplication::parse(&sql)?),
                 command => {
                     debug!("unknown admin show command: '{}'", command);
                     return Err(Error::Syntax);

pgdog/src/admin/show_pools.rs

@@ -42,7 +42,6 @@ impl Command for ShowPools {
             Field::numeric("out_of_sync"),
             Field::numeric("force_closed"),
             Field::bool("online"),
-            Field::text("replica_lag"),
             Field::bool("schema_admin"),
         ]);
         let mut messages = vec![rd.message()?];
@@ -76,7 +75,6 @@ impl Command for ShowPools {
                         .add(state.out_of_sync)
                         .add(state.force_close)
                         .add(state.online)
-                        .add(state.replica_lag.simple_display())
                         .add(cluster.schema_admin());
 
                     messages.push(row.message()?);

pgdog/src/admin/show_replication.rs

@@ -0,0 +1,89 @@
+use tokio::time::Instant;
+
+use crate::{
+    backend::databases::databases,
+    net::{
+        data_row::Data,
+        messages::{DataRow, Field, Protocol, RowDescription},
+        ToDataRowColumn,
+    },
+};
+
+use super::prelude::*;
+
+pub struct ShowReplication;
+
+#[async_trait]
+impl Command for ShowReplication {
+    fn name(&self) -> String {
+        "SHOW REPLICATION".into()
+    }
+
+    fn parse(_sql: &str) -> Result<Self, Error> {
+        Ok(ShowReplication {})
+    }
+
+    async fn execute(&self) -> Result<Vec<Message>, Error> {
+        let rd = RowDescription::new(&[
+            Field::bigint("id"),
+            Field::text("database"),
+            Field::text("user"),
+            Field::text("addr"),
+            Field::numeric("port"),
+            Field::numeric("shard"),
+            Field::text("role"),
+            Field::text("pg_replica_lag"),
+            Field::text("pg_lsn"),
+            Field::text("lsn_age"),
+            Field::text("pg_is_in_recovery"),
+        ]);
+        let mut messages = vec![rd.message()?];
+        let now = Instant::now();
+        for (user, cluster) in databases().all() {
+            for (shard_num, shard) in cluster.shards().iter().enumerate() {
+                for (role, _ban, pool) in shard.pools_with_roles_and_bans() {
+                    let mut row = DataRow::new();
+                    let state = pool.state();
+
+                    let lsn_read = state.lsn_stats.lsn.lsn > 0;
+                    let lsn_age = now.duration_since(state.lsn_stats.fetched);
+
+                    row.add(pool.id() as i64)
+                        .add(user.database.as_str())
+                        .add(user.user.as_str())
+                        .add(pool.addr().host.as_str())
+                        .add(pool.addr().port as i64)
+                        .add(shard_num as i64)
+                        .add(role.to_string())
+                        .add(if lsn_read {
+                            state
+                                .replica_lag
+                                .as_millis()
+                                .to_string()
+                                .to_data_row_column()
+                        } else {
+                            Data::null()
+                        })
+                        .add(if lsn_read {
+                            state.lsn_stats.lsn.to_string().to_data_row_column()
+                        } else {
+                            Data::null()
+                        })
+                        .add(if lsn_read {
+                            lsn_age.as_millis().to_string().to_data_row_column()
+                        } else {
+                            Data::null()
+                        })
+                        .add(if lsn_read {
+                            state.lsn_stats.replica.to_data_row_column()
+                        } else {
+                            Data::null()
+                        });
+
+                    messages.push(row.message()?);
+                }
+            }
+        }
+        Ok(messages)
+    }
+}

pgdog/src/admin/tests/mod.rs

@@ -118,7 +118,6 @@ async fn show_pools_reports_schema_admin_flag() {
         "out_of_sync",
         "force_closed",
         "online",
-        "replica_lag",
         "schema_admin",
     ];
     assert_eq!(actual_names, expected_names);