feat: add admin FORCE_RELOAD command

integration/rust/tests/integration/admin_reload.rs

@@ -0,0 +1,47 @@
+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_reload_test() {
+    let admin = admin_sqlx().await;
+    let connections = connections_sqlx().await; // [pgdog, pgdog_sharded]
+
+    {
+        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();
+
+        // Standard Postgres error pertaining to pg_terminate_backend; <https://www.postgresql.org/docs/current/functions-admin.html>
+        assert!(
+            err.as_database_error()
+                .unwrap()
+                .message()
+                .contains("terminating connection due to administrator command")
+        );
+
+        // 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();
+}

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/admin_reload.rs

@@ -0,0 +1,24 @@
+//! FORCE RELOAD command.
+
+use super::prelude::*;
+use crate::backend::databases::{reload, terminate_active_connections};
+
+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> {
+        terminate_active_connections().await?;
+        reload()?;
+
+        Ok(vec![])
+    }
+}

pgdog/src/admin/mod.rs

@@ -4,6 +4,7 @@ use async_trait::async_trait;
 
 use crate::net::messages::Message;
 
+pub(crate) mod admin_reload;
 pub(crate) mod ban;
 pub(crate) mod copy_data;
 pub(crate) mod cutover;

pgdog/src/admin/parser.rs

@@ -1,6 +1,6 @@
 //! Admin command parser.
 
-use crate::admin::show_guc::get_show_variable;
+use crate::admin::{admin_reload::ForceReload, show_guc::get_show_variable};
 
 use super::*;
 
@@ -12,6 +12,7 @@ pub(crate) enum ParseResult {
     Reconnect(Reconnect),
     ShowClients(ShowClients),
     Reload(Reload),
+    ForceReload(ForceReload),
     ShowPools(ShowPools),
     ShowBans(ShowBans),
     ShowConfig(ShowConfig),
@@ -61,6 +62,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 +112,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 +209,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 +290,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/pause.rs

@@ -59,7 +59,7 @@ impl Command for Pause {
                     if self.resume {
                         pool.resume();
                     } else {
-                        pool.pause();
+                        pool.pause(false);
                     }
                 }
             }

pgdog/src/backend/databases.rs

@@ -133,12 +133,27 @@ pub(crate) async fn cancel_all(database: &str) -> Result<(), Error> {
     Ok(())
 }
 
+/// Terminates all active connections on all `Cluster`s.
+pub(crate) async fn terminate_active_connections() -> Result<(), Error> {
+    let clusters: Vec<_> = databases().all().values().cloned().collect();
+
+    try_join_all(
+        clusters
+            .iter()
+            .map(|cluster| cluster.terminate_active_connections()),
+    )
+    .await?;
+
+    Ok(())
+}
+
 /// Re-create pools from config.
 pub(crate) fn reload() -> Result<(), Error> {
     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);
 

pgdog/src/backend/pool/cluster.rs

@@ -8,6 +8,7 @@ use pgdog_config::{
 };
 use std::{sync::Arc, time::Duration};
 
+use crate::backend;
 use crate::backend::schema::SchemaCache;
 use crate::backend::server::ServerRequest;
 use crate::frontend::router::sharding::ShardedTable;
@@ -719,6 +720,42 @@ impl Cluster {
         Ok(())
     }
 
+    /// Terminates all active connections; more specifically, aimed towards terminating active, in-flight transactions.
+    pub(crate) async fn terminate_active_connections(&self) -> Result<(), backend::error::Error> {
+        for shard in self.shards() {
+            let pools = shard.pools();
+            for pool in pools {
+                let keys = pool.active_connections();
+                if !keys.is_empty() {
+                    // Prevent chance of more connections slipping through before we terminate backends.
+                    // When passing true, if the pool previously was NOT paused, it'll be resumed
+                    // again on the new `Pool` when the transfer happens later.
+                    //
+                    // TODO: If we error below on standalone or execute,
+                    // it'll leave the `Pool` in a pause state incorrectly.
+                    pool.pause(true);
+
+                    // Connect outside of `Pool` idle connections to prevent waiting for an available connection.
+                    // This also bypasses [`Pool.pause`]
+                    //
+                    // TODO: Say that this fails for some reason; it already internally re-tries multiple times.
+                    //       should we Error because not all transactions are terminated? Ignore it?
+                    let mut server = pool.standalone(backend::ConnectReason::Other).await?;
+
+                    for key in keys {
+                        // `pg_terminate_backend` will send a SIGTERM signal to the backend process corresponding with
+                        // the active connection belonging to the transaction.
+                        let request: ServerRequest =
+                            format!("SELECT pg_terminate_backend({});", key.pid).into();
+                        server.execute(request).await?;
+                    }
+                }
+            }
+        }
+
+        Ok(())
+    }
+
     /// Run a parameterized query on one shard, picked round-robin, and
     /// return all rows. The answer is only authoritative if every shard
     /// has the same data, e.g. an omnisharded table.

pgdog/src/backend/pool/inner.rs

@@ -31,6 +31,8 @@ pub(super) struct Inner {
     pub(super) online: bool,
     /// Pool is paused.
     pub(super) paused: bool,
+    // Pool's `paused` will not be propagated on transfer in `move_conns_to`
+    pub(super) remove_pause_on_transfer: bool,
     /// Track out of sync terminations.
     pub(super) out_of_sync: usize,
     /// How many times servers had to be re-synced
@@ -78,6 +80,7 @@ impl Inner {
             waiting: VecDeque::new(),
             online: false,
             paused: false,
+            remove_pause_on_transfer: false,
             force_close: 0,
             out_of_sync: 0,
             re_synced: 0,

pgdog/src/backend/pool/pool_impl.rs

@@ -17,7 +17,7 @@ use crate::backend::pool::LsnStats;
 use crate::backend::{ConnectReason, DisconnectReason, Server, ServerOptions};
 use crate::config::PoolerMode;
 use crate::net::messages::{BackendPid, FrontendPid};
-use crate::net::{Liveness, Parameter, Parameters};
+use crate::net::{BackendKeyData, Liveness, Parameter, Parameters};
 
 use super::inner::CheckInResult;
 use super::{
@@ -313,15 +313,16 @@ impl Pool {
             let mut to_guard = destination.lock();
 
             // Propagate pause state so a paused database stays paused after reload.
-            if from_guard.paused {
-                to_guard.paused = true;
-            }
-
+            // Only set if `remove_on_transfer_if_not_paused` is not set, which happens
+            // during admin FORCE_RELOAD command, and means that the `Pool` previously wasn't paused.
+            to_guard.paused = from_guard.paused && !from_guard.remove_pause_on_transfer;
             from_guard.online = false;
+
             let (idle, taken) = from_guard.move_conns_to(destination);
             for server in idle {
                 to_guard.put(server, now)?;
             }
+
             to_guard.set_taken(taken);
         }
 
@@ -336,11 +337,15 @@ impl Pool {
     }
 
     /// Pause pool, closing all open connections.
-    pub(crate) fn pause(&self) {
+    /// If `remove_on_transfer` is true, then the pause will be removed on `move_conns_to`
+    pub(crate) fn pause(&self, remove_on_transfer: bool) {
         let mut guard = self.lock();
-
+        if !guard.paused && remove_on_transfer {
+            guard.remove_pause_on_transfer = true;
+        } else {
+            guard.dump_idle();
+        }
         guard.paused = true;
-        guard.dump_idle();
     }
 
     /// Send a cancellation request for all running queries.
@@ -352,12 +357,19 @@ impl Pool {
             .cancel_keys()
             .map(|key| Server::cancel(&addr, key.clone()))
             .collect();
+
         try_join_all(futures)
             .await
             .map_err(|_| Error::FastShutdown)?;
         Ok(())
     }
 
+    /// Fetch cancel keys for all active connections belonging to the `Pool`
+    pub(crate) fn active_connections(&self) -> Vec<BackendKeyData> {
+        // Collect into a Vec to drop the pool lock
+        self.lock().cancel_keys().cloned().collect()
+    }
+
     /// Resume the pool.
     pub(crate) fn resume(&self) {
         {

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

@@ -226,7 +226,7 @@ async fn test_pause() {
     // Make sure we're not blocked still.
     drop(pool.get(&Request::default()).await.unwrap());
 
-    pool.pause();
+    pool.pause(false);
 
     // We'll hit the timeout now because we're waiting forever.
     let pause = Duration::from_millis(2_000);
@@ -250,7 +250,7 @@ async fn test_pause() {
     // Shutdown the pool while clients wait.
     // Makes sure they get woken up and kicked out of
     // the pool.
-    pool.pause();
+    pool.pause(false);
     let tracker = TaskTracker::new();
     let didnt_work = Arc::new(AtomicBool::new(false));
     for _ in 0..1000 {
@@ -1224,7 +1224,7 @@ async fn test_move_conns_to_propagates_pause_state() {
     destination.launch();
 
     // Pause the source pool.
-    source.pause();
+    source.pause(false);
     assert!(source.lock().paused);
     assert!(!destination.lock().paused);