feat: add admin FORCE_RELOAD command

integration/rust/tests/integration/admin_reload.rs

@@ -0,0 +1,111 @@
+use sqlx::Row;
+use std::time::Duration;
+
+use tokio::time::{Instant, sleep};
+
+use crate::setup::{admin_sqlx, connections_sqlx};
+
+/// <https://github.com/pgdogdev/pgdog/issues/1472>
+/// Test the implementation of the command FORCE_RELOAD, which is a normal RELOAD + terminates all in-flight transactions.
+#[tokio::test]
+async fn admin_force_reload_test() {
+    let admin = admin_sqlx().await;
+    let connections = connections_sqlx().await; // [pgdog, pgdog_sharded]
+
+    // Test with an idle transaction.
+    {
+        let mut transaction = connections.get(1).unwrap().begin().await.unwrap();
+
+        // Isn't strictly needed for the functionaltiy of the test; but why not?
+        sqlx::raw_sql("SELECT * FROM sharded")
+            .fetch_all(&mut *transaction)
+            .await
+            .unwrap();
+
+        // After we force reload, existing transactions (i.e. this one) are terminated.
+        sqlx::raw_sql("FORCE_RELOAD").execute(&admin).await.unwrap();
+
+        let err = sqlx::raw_sql("SELECT * FROM sharded")
+            .fetch_all(&mut *transaction)
+            .await
+            .err()
+            .unwrap();
+
+        // Identical to Postgres error pertaining to normal backend termination.
+        // <https://www.postgresql.org/docs/current/functions-admin.html>
+        assert!(
+            err.as_database_error()
+                .unwrap()
+                .message()
+                .contains("terminating connection due to administrator command")
+        );
+
+        // Should be overwritten from the generic error code to the same one as Postgres.
+        assert_eq!(err.as_database_error().unwrap().code().unwrap(), "57P01");
+
+        // The transaction drops (allowing another connection in sqlx `Pool`)
+    }
+
+    // Does it work with a new Pool?
+    let test = connections_sqlx().await;
+    let test = test.get(1).unwrap();
+    sqlx::raw_sql("SELECT 1234").fetch_all(test).await.unwrap();
+
+    // Does it still with the old Pool?
+    let conn = connections.get(1).unwrap();
+    sqlx::raw_sql("SELECT 1000").fetch_all(conn).await.unwrap();
+
+    // Additionally try when we're in an active query (as opposed to an idle transaction)
+    {
+        let query = async {
+            let connections = connections_sqlx().await;
+            let mut transaction = connections.get(1).unwrap().begin().await.unwrap();
+
+            let start_ts = Instant::now();
+
+            // Arbitrary amount of time; it should be cancelled long before; lets us be sure that
+            // FORCE_RELOAD is actually cancelling it mid-query
+            let err = sqlx::raw_sql("SELECT pg_sleep(5)")
+                .execute(&mut *transaction)
+                .await
+                .err()
+                .unwrap();
+
+            assert!(
+                err.as_database_error()
+                    .unwrap()
+                    .message()
+                    .contains("terminating connection due to administrator command")
+            );
+
+            // We're sleeping for 5s. Waiting for 500ms on `reload` async block.
+            // So it should be <1s until we get the admin termination error.
+            assert!(start_ts.elapsed() < Duration::from_secs(1));
+        };
+
+        let reload = async {
+            sleep(Duration::from_millis(500)).await;
+
+            // After we force reload, existing transactions (i.e. this one) are terminated.
+            sqlx::raw_sql("FORCE_RELOAD").execute(&admin).await.unwrap();
+        };
+
+        // Runs both the query and reload concurrently
+        // Returns after both have finished.
+        tokio::join!(query, reload);
+
+        // The transaction was terminated (through PgDog), however, is Postgres still running it?
+        let rows = sqlx::raw_sql(
+            "SELECT * FROM pg_stat_activity WHERE state = 'active' AND query NOT LIKE '%pg_stat_activity%'",
+        )
+        .fetch_all(conn)
+        .await
+        .unwrap();
+
+        // I did it this way to prevent flaky tests if the health-check were to run in parallel
+        // Usually there's no rows.
+        for row in rows {
+            assert!(!row.get::<&str, &str>("query").eq("SELECT pg_sleep(5)"));
+        }
+    }
+}

integration/rust/tests/integration/mod.rs

@@ -1,4 +1,5 @@
 pub mod admin;
+pub mod admin_reload;
 pub mod admin_termination;
 pub mod auth;
 pub mod auto_id;

pgdog/src/admin/force_reload.rs

@@ -0,0 +1,24 @@
+//! FORCE RELOAD command.
+
+use super::prelude::*;
+use crate::backend::databases::reload;
+
+pub(crate) struct ForceReload;
+
+#[async_trait]
+impl Command for ForceReload {
+    fn name(&self) -> String {
+        "FORCE_RELOAD".into()
+    }
+
+    fn parse(_sql: &str) -> Result<Self, Error> {
+        Ok(ForceReload)
+    }
+
+    async fn execute(&self) -> Result<Vec<Message>, Error> {
+        // true = force reload
+        reload(true)?;
+
+        Ok(vec![])
+    }
+}

pgdog/src/admin/mod.rs

@@ -8,6 +8,7 @@ pub(crate) mod ban;
 pub(crate) mod copy_data;
 pub(crate) mod cutover;
 pub(crate) mod error;
+pub(crate) mod force_reload;
 pub(crate) mod healthcheck;
 pub(crate) mod maintenance_mode;
 pub(crate) mod named_row;
@@ -55,6 +56,7 @@ pub(crate) use ban::*;
 pub(crate) use copy_data::*;
 pub(crate) use cutover::*;
 pub(crate) use error::Error;
+pub(crate) use force_reload::*;
 pub(crate) use healthcheck::*;
 pub(crate) use maintenance_mode::*;
 pub(crate) use pause::*;

pgdog/src/admin/parser.rs

@@ -1,8 +1,7 @@
 //! Admin command parser.
 
-use crate::admin::show_guc::get_show_variable;
-
 use super::*;
+use crate::admin::show_guc::get_show_variable;
 
 use tracing::debug;
 
@@ -12,6 +11,7 @@ pub(crate) enum ParseResult {
     Reconnect(Reconnect),
     ShowClients(ShowClients),
     Reload(Reload),
+    ForceReload(ForceReload),
     ShowPools(ShowPools),
     ShowBans(ShowBans),
     ShowConfig(ShowConfig),
@@ -61,6 +61,7 @@ impl ParseResult {
             Reconnect(reconnect) => reconnect.execute().await,
             ShowClients(show_clients) => show_clients.execute().await,
             Reload(reload) => reload.execute().await,
+            ForceReload(force_reload) => force_reload.execute().await,
             ShowPools(show_pools) => show_pools.execute().await,
             ShowBans(show_bans) => show_bans.execute().await,
             ShowConfig(show_config) => show_config.execute().await,
@@ -110,6 +111,7 @@ impl ParseResult {
             Reconnect(reconnect) => reconnect.name(),
             ShowClients(show_clients) => show_clients.name(),
             Reload(reload) => reload.name(),
+            ForceReload(force_reload) => force_reload.name(),
             ShowPools(show_pools) => show_pools.name(),
             ShowBans(show_bans) => show_bans.name(),
             ShowConfig(show_config) => show_config.name(),
@@ -206,6 +208,7 @@ impl Parser {
             "shutdown" => ParseResult::Shutdown(Shutdown::parse(&sql)?),
             "reconnect" => ParseResult::Reconnect(Reconnect::parse(&sql)?),
             "reload" => ParseResult::Reload(Reload::parse(&sql)?),
+            "force_reload" => ParseResult::ForceReload(ForceReload::parse(&sql)?),
             "ban" | "unban" => ParseResult::Ban(Ban::parse(&sql)?),
             "healthcheck" => ParseResult::Healthcheck(Healthcheck::parse(&sql)?),
             // These are not covered by the show handler above
@@ -286,6 +289,7 @@ mod tests {
         assert_parses!("RESUME", ParseResult::Pause(_));
         assert_parses!("RECONNECT", ParseResult::Reconnect(_));
         assert_parses!("RELOAD", ParseResult::Reload(_));
+        assert_parses!("FORCE_RELOAD", ParseResult::ForceReload(_));
         assert_parses!("SHUTDOWN", ParseResult::Shutdown(_));
         assert_parses!("BAN", ParseResult::Ban(_));
         assert_parses!("UNBAN", ParseResult::Ban(_));

pgdog/src/admin/reload.rs

@@ -16,7 +16,8 @@ impl Command for Reload {
     }
 
     async fn execute(&self) -> Result<Vec<Message>, Error> {
-        reload()?;
+        // false = NOT a force reload
+        reload(false)?;
         Ok(vec![])
     }
 }

pgdog/src/backend/databases.rs

@@ -1,9 +1,5 @@
 //! Databases behind pgDog.
 
-use std::collections::HashMap;
-use std::ops::Deref;
-use std::sync::Arc;
-
 use arc_swap::ArcSwap;
 use futures::future::try_join_all;
 use indexmap::IndexMap;
@@ -17,6 +13,9 @@ use pgdog_config::{
     EnumeratedDatabase, QueryParser, ShardedMappingConfig, ShardedMappingKey, ShardedMappingKeyRef,
     ShardedMappingKindDeprecated, ShardedMappingList, ShardedMappingRange, ShardedTableConfig,
 };
+use std::collections::HashMap;
+use std::ops::Deref;
+use std::sync::Arc;
 use tracing::{debug, error, info, warn};
 
 use crate::auth::AuthResult;
@@ -133,15 +132,32 @@ pub(crate) async fn cancel_all(database: &str) -> Result<(), Error> {
     Ok(())
 }
 
+/// Terminates all active connections on all `Pool`s.
+pub(crate) fn terminate_active_connections() {
+    databases()
+        .all()
+        .values()
+        .for_each(Cluster::terminate_active_connections);
+}
+
 /// Re-create pools from config.
-pub(crate) fn reload() -> Result<(), Error> {
-    info!("reloading configuration");
+pub(crate) fn reload(force: bool) -> Result<(), Error> {
+    if force {
+        info!("force reloading configuration");
+    } else {
+        info!("reloading configuration");
+    }
 
     // Load config from disk.
     let old_config = config();
     let new_config = load(&old_config.config_path, &old_config.users_path)?;
     let databases = from_config(&new_config);
 
+    // Terminate after checking config for validity.
+    if force {
+        terminate_active_connections();
+    }
+
     // Replace databases.
     replace_databases(databases, true)?;
 

pgdog/src/backend/pool/cluster.rs

@@ -7,6 +7,7 @@ use pgdog_config::{
     users::PasswordKind,
 };
 use std::{sync::Arc, time::Duration};
+use tokio_util::sync::CancellationToken;
 
 use crate::backend::schema::SchemaCache;
 use crate::backend::server::ServerRequest;
@@ -87,6 +88,7 @@ pub(crate) struct Cluster {
     canonical_oids: Option<Arc<CanonicalOids>>,
     read_only: bool,
     failover_signal: ClusterFailoverSignalWatcher,
+    cancellation_token: CancellationToken,
 }
 
 /// Bare test clusters carry the same defaults the config would apply,
@@ -136,6 +138,7 @@ impl Default for Cluster {
             canonical_oids: Default::default(),
             read_only: Default::default(),
             failover_signal: ClusterFailoverSignalWatcher::default(),
+            cancellation_token: Default::default(),
         }
     }
 }
@@ -415,9 +418,25 @@ impl Cluster {
             canonical_oids,
             read_only,
             failover_signal,
+            cancellation_token: Default::default(),
         }
     }
 
+    pub(crate) fn get_cancellation_token(&self) -> CancellationToken {
+        self.cancellation_token.clone()
+    }
+
+    /// Terminates all active connections for the `Cluster`
+    /// and marks all `Pool`s as offline to refuse future connections.
+    pub(crate) fn terminate_active_connections(&self) {
+        for shard in self.shards() {
+            for pool in shard.pools() {
+                pool.set_offline();
+            }
+        }
+        self.cancellation_token.cancel();
+    }
+
     /// Change config to work with logical replication streaming.
     pub(crate) fn logical_stream(&self) -> Self {
         let mut cluster = self.clone();

pgdog/src/backend/pool/connection/mod.rs

@@ -1,8 +1,10 @@
 //! Server connection requested by a frontend.
 
+use futures::future::try_join_all;
 use mirror::MirrorHandler;
 use pgdog_config::users::PasswordKind;
 use tokio::select;
+use tokio_util::sync::CancellationToken;
 use tracing::debug;
 
 use crate::{
@@ -22,7 +24,7 @@ use crate::{
 };
 
 use super::{
-    super::{Error, pool::Guard},
+    super::{Error, Server, pool::Guard},
     Address, Cluster, Request,
 };
 
@@ -51,6 +53,9 @@ pub(crate) struct Connection {
     database: String,
     binding: Binding,
     cluster: Option<Cluster>,
+    /// Each client polls own child node instead of contending on the shared `Cluster` node.
+    /// Cancelled when an admin terminates the cluster (`FORCE_RELOAD`)
+    cancellation_token: CancellationToken,
     mirrors: Vec<MirrorHandler>,
     pub_sub: PubSubClient,
 }
@@ -65,6 +70,7 @@ impl Connection {
                 Binding::NotConnected
             },
             cluster: None,
+            cancellation_token: CancellationToken::new(),
             user: user.to_owned(),
             database: database.to_owned(),
             mirrors: vec![],
@@ -396,6 +402,7 @@ impl Connection {
         let databases = databases();
         let cluster = databases.cluster(user)?;
 
+        self.cancellation_token = cluster.get_cancellation_token().child_token();
         self.cluster = Some(cluster.clone());
         let source_db = cluster.name();
         self.mirrors = databases
@@ -466,6 +473,29 @@ impl Connection {
         })
     }
 
+    /// Cancel the query the server(s) are running for this client
+    pub(crate) async fn cancel_query(&self) -> Result<(), Error> {
+        let servers: Vec<&Guard> = match self.binding {
+            Binding::Direct(ref server, ..) => vec![server],
+            Binding::MultiShard(ref servers, _) => servers.iter().collect(),
+            _ => return Ok(()),
+        };
+
+        try_join_all(
+            servers
+                .iter()
+                .map(|server| Server::cancel(server.addr(), server.key().clone())),
+        )
+        .await?;
+
+        Ok(())
+    }
+
+    /// Token cancelled when an admin terminates this connection's `Cluster`.
+    pub(crate) fn cancellation_token(&self) -> CancellationToken {
+        self.cancellation_token.clone()
+    }
+
     /// Get cluster if any.
     #[inline]
     pub(crate) fn cluster(&self) -> Result<&Cluster, Error> {