From 885115ed1377e846a320be06db2897ac32c469d5 Mon Sep 17 00:00:00 2001 From: jkaczman Date: Sat, 5 Sep 2026 10:09:47 -0400 Subject: [PATCH 1/4] feat: add admin FORCE_RELOAD command --- .../rust/tests/integration/admin_reload.rs | 47 +++++++++++++++++++ integration/rust/tests/integration/mod.rs | 1 + pgdog/src/admin/admin_reload.rs | 24 ++++++++++ pgdog/src/admin/mod.rs | 1 + pgdog/src/admin/parser.rs | 7 ++- pgdog/src/admin/pause.rs | 2 +- pgdog/src/backend/databases.rs | 15 ++++++ pgdog/src/backend/pool/cluster.rs | 36 ++++++++++++++ pgdog/src/backend/pool/inner.rs | 3 ++ pgdog/src/backend/pool/pool_impl.rs | 25 ++++++++-- pgdog/src/backend/pool/test/mod.rs | 6 +-- 11 files changed, 158 insertions(+), 9 deletions(-) create mode 100644 integration/rust/tests/integration/admin_reload.rs create mode 100644 pgdog/src/admin/admin_reload.rs diff --git a/integration/rust/tests/integration/admin_reload.rs b/integration/rust/tests/integration/admin_reload.rs new file mode 100644 index 000000000..c96c509a2 --- /dev/null +++ b/integration/rust/tests/integration/admin_reload.rs @@ -0,0 +1,47 @@ +use crate::setup::{admin_sqlx, connections_sqlx}; + +/// +/// 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; + 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(); +} diff --git a/integration/rust/tests/integration/mod.rs b/integration/rust/tests/integration/mod.rs index 03badb9eb..082df0d21 100644 --- a/integration/rust/tests/integration/mod.rs +++ b/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; diff --git a/pgdog/src/admin/admin_reload.rs b/pgdog/src/admin/admin_reload.rs new file mode 100644 index 000000000..68ac50136 --- /dev/null +++ b/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 { + Ok(ForceReload) + } + + async fn execute(&self) -> Result, Error> { + terminate_active_connections().await?; + reload()?; + + Ok(vec![]) + } +} diff --git a/pgdog/src/admin/mod.rs b/pgdog/src/admin/mod.rs index 6565801da..3c5e25255 100644 --- a/pgdog/src/admin/mod.rs +++ b/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; diff --git a/pgdog/src/admin/parser.rs b/pgdog/src/admin/parser.rs index 771afa22c..f7d786427 100644 --- a/pgdog/src/admin/parser.rs +++ b/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(_)); diff --git a/pgdog/src/admin/pause.rs b/pgdog/src/admin/pause.rs index 05d05feb7..9c3e42aae 100644 --- a/pgdog/src/admin/pause.rs +++ b/pgdog/src/admin/pause.rs @@ -59,7 +59,7 @@ impl Command for Pause { if self.resume { pool.resume(); } else { - pool.pause(); + pool.pause(false); } } } diff --git a/pgdog/src/backend/databases.rs b/pgdog/src/backend/databases.rs index 08f67a164..ee6839459 100644 --- a/pgdog/src/backend/databases.rs +++ b/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); diff --git a/pgdog/src/backend/pool/cluster.rs b/pgdog/src/backend/pool/cluster.rs index 53fd035f0..e584d4e5f 100644 --- a/pgdog/src/backend/pool/cluster.rs +++ b/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,41 @@ 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 { + // TODO: What happens if the pool is no longer working? (test this) + // Obtain a transaction -> poison pool somehow -> call FORCE_RELOAD + + 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. + pool.pause(true); + + // Connect outside of `Pool` idle connections to prevent waiting for an available connection. + // This also bypasses [`Pool.pause`] + 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?; + + // TODO: Should we be removing the in-memory PIDs from `Taken`? + } + } + } + } + + 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. diff --git a/pgdog/src/backend/pool/inner.rs b/pgdog/src/backend/pool/inner.rs index 2f55a0046..014c6d19a 100644 --- a/pgdog/src/backend/pool/inner.rs +++ b/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, diff --git a/pgdog/src/backend/pool/pool_impl.rs b/pgdog/src/backend/pool/pool_impl.rs index 1c351dd74..3e544b79c 100644 --- a/pgdog/src/backend/pool/pool_impl.rs +++ b/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::{ @@ -314,7 +314,13 @@ impl Pool { // 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. + if !from_guard.remove_pause_on_transfer { + to_guard.paused = true; + } else { + // TODO: Do we need to notify waiters? + } } from_guard.online = false; @@ -322,6 +328,7 @@ impl Pool { for server in idle { to_guard.put(server, now)?; } + to_guard.set_taken(taken); } @@ -336,9 +343,12 @@ 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; + } guard.paused = true; guard.dump_idle(); } @@ -352,12 +362,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 { + // Collect into a Vec to drop the pool lock + self.lock().cancel_keys().cloned().collect() + } + /// Resume the pool. pub(crate) fn resume(&self) { { diff --git a/pgdog/src/backend/pool/test/mod.rs b/pgdog/src/backend/pool/test/mod.rs index 4201a3c71..b12c2eb3b 100644 --- a/pgdog/src/backend/pool/test/mod.rs +++ b/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); From 900ca277f9b6af6212e3597f91bce16aacbf908a Mon Sep 17 00:00:00 2001 From: jkaczman Date: Sat, 5 Sep 2026 10:44:30 -0400 Subject: [PATCH 2/4] Don't dump idle connections in when remove_on_transfer flag is set --- pgdog/src/backend/pool/pool_impl.rs | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/pgdog/src/backend/pool/pool_impl.rs b/pgdog/src/backend/pool/pool_impl.rs index 3e544b79c..1939a47ca 100644 --- a/pgdog/src/backend/pool/pool_impl.rs +++ b/pgdog/src/backend/pool/pool_impl.rs @@ -348,9 +348,10 @@ impl Pool { 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. From 6298dcd6187aaafddeff3a4fcafcb7bc1c8308db Mon Sep 17 00:00:00 2001 From: jkaczman Date: Sun, 6 Sep 2026 15:50:31 -0400 Subject: [PATCH 3/4] self.shutdown() will take care of notifying waiters --- pgdog/src/backend/pool/pool_impl.rs | 14 ++++---------- 1 file changed, 4 insertions(+), 10 deletions(-) diff --git a/pgdog/src/backend/pool/pool_impl.rs b/pgdog/src/backend/pool/pool_impl.rs index 1939a47ca..a6bd2647b 100644 --- a/pgdog/src/backend/pool/pool_impl.rs +++ b/pgdog/src/backend/pool/pool_impl.rs @@ -313,17 +313,11 @@ impl Pool { let mut to_guard = destination.lock(); // Propagate pause state so a paused database stays paused after reload. - if from_guard.paused { - // 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. - if !from_guard.remove_pause_on_transfer { - to_guard.paused = true; - } else { - // TODO: Do we need to notify waiters? - } - } - + // 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)?; From 8838e5dbee069d6f3fbb2c441dd32b21c1f40e2a Mon Sep 17 00:00:00 2001 From: jkaczman Date: Sun, 6 Sep 2026 19:13:33 -0400 Subject: [PATCH 4/4] Refresh TODOs --- pgdog/src/backend/pool/cluster.rs | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/pgdog/src/backend/pool/cluster.rs b/pgdog/src/backend/pool/cluster.rs index e584d4e5f..315ca8a47 100644 --- a/pgdog/src/backend/pool/cluster.rs +++ b/pgdog/src/backend/pool/cluster.rs @@ -725,18 +725,21 @@ impl Cluster { for shard in self.shards() { let pools = shard.pools(); for pool in pools { - // TODO: What happens if the pool is no longer working? (test this) - // Obtain a transaction -> poison pool somehow -> call FORCE_RELOAD - 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 { @@ -745,8 +748,6 @@ impl Cluster { let request: ServerRequest = format!("SELECT pg_terminate_backend({});", key.pid).into(); server.execute(request).await?; - - // TODO: Should we be removing the in-memory PIDs from `Taken`? } } }