From 885115ed1377e846a320be06db2897ac32c469d5 Mon Sep 17 00:00:00 2001 From: jkaczman Date: Sat, 5 Sep 2026 10:09:47 -0400 Subject: [PATCH 1/8] 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/8] 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/8] 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/8] 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`? } } } From 1082bbc384134f648c94e1f7ae0631e7a278124b Mon Sep 17 00:00:00 2001 From: jkaczman Date: Tue, 8 Sep 2026 12:09:01 -0400 Subject: [PATCH 5/8] Use per-pool CancellationTokens instead of Postgres terminate_backend signaling func --- .../rust/tests/integration/admin_reload.rs | 53 ++++++++++++++++- .../{admin_reload.rs => force_reload.rs} | 6 +- pgdog/src/admin/mod.rs | 3 +- pgdog/src/admin/parser.rs | 3 +- pgdog/src/admin/pause.rs | 2 +- pgdog/src/admin/reload.rs | 3 +- pgdog/src/backend/databases.rs | 36 +++++++----- pgdog/src/backend/pool/cluster.rs | 37 ------------ pgdog/src/backend/pool/connection/binding.rs | 13 +++++ pgdog/src/backend/pool/inner.rs | 3 - pgdog/src/backend/pool/pool_impl.rs | 34 +++++------ pgdog/src/backend/pool/test/mod.rs | 6 +- pgdog/src/frontend/client/mod.rs | 58 +++++++++++++------ pgdog/src/frontend/client/query_engine/mod.rs | 6 ++ .../src/frontend/client/query_engine/query.rs | 29 ++++++++-- pgdog/src/frontend/client/test/mod.rs | 8 +-- pgdog/src/frontend/client/test/test_client.rs | 4 +- pgdog/src/frontend/error.rs | 3 + pgdog/src/frontend/listener.rs | 2 +- pgdog/src/net/messages/error_response.rs | 16 +++++ 20 files changed, 211 insertions(+), 114 deletions(-) rename pgdog/src/admin/{admin_reload.rs => force_reload.rs} (72%) diff --git a/integration/rust/tests/integration/admin_reload.rs b/integration/rust/tests/integration/admin_reload.rs index c96c509a2..305bdabb5 100644 --- a/integration/rust/tests/integration/admin_reload.rs +++ b/integration/rust/tests/integration/admin_reload.rs @@ -1,12 +1,17 @@ +use std::time::Duration; + +use tokio::time::{Instant, sleep}; + 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() { +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(); @@ -25,7 +30,8 @@ async fn admin_reload_test() { .err() .unwrap(); - // Standard Postgres error pertaining to pg_terminate_backend; + // Identical to Postgres error pertaining to normal backend termination. + // assert!( err.as_database_error() .unwrap() @@ -33,6 +39,9 @@ async fn admin_reload_test() { .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`) } @@ -44,4 +53,44 @@ async fn admin_reload_test() { // 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); + } } diff --git a/pgdog/src/admin/admin_reload.rs b/pgdog/src/admin/force_reload.rs similarity index 72% rename from pgdog/src/admin/admin_reload.rs rename to pgdog/src/admin/force_reload.rs index 68ac50136..cd7cf9047 100644 --- a/pgdog/src/admin/admin_reload.rs +++ b/pgdog/src/admin/force_reload.rs @@ -1,7 +1,7 @@ //! FORCE RELOAD command. use super::prelude::*; -use crate::backend::databases::{reload, terminate_active_connections}; +use crate::backend::databases::reload; pub(crate) struct ForceReload; @@ -16,8 +16,8 @@ impl Command for ForceReload { } async fn execute(&self) -> Result, Error> { - terminate_active_connections().await?; - reload()?; + // true = force reload + reload(true)?; Ok(vec![]) } diff --git a/pgdog/src/admin/mod.rs b/pgdog/src/admin/mod.rs index 3c5e25255..381647da8 100644 --- a/pgdog/src/admin/mod.rs +++ b/pgdog/src/admin/mod.rs @@ -4,11 +4,11 @@ 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; pub(crate) mod error; +pub(crate) mod force_reload; pub(crate) mod healthcheck; pub(crate) mod maintenance_mode; pub(crate) mod named_row; @@ -56,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::*; diff --git a/pgdog/src/admin/parser.rs b/pgdog/src/admin/parser.rs index f7d786427..545111420 100644 --- a/pgdog/src/admin/parser.rs +++ b/pgdog/src/admin/parser.rs @@ -1,8 +1,7 @@ //! Admin command parser. -use crate::admin::{admin_reload::ForceReload, show_guc::get_show_variable}; - use super::*; +use crate::admin::show_guc::get_show_variable; use tracing::debug; diff --git a/pgdog/src/admin/pause.rs b/pgdog/src/admin/pause.rs index 9c3e42aae..05d05feb7 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(false); + pool.pause(); } } } diff --git a/pgdog/src/admin/reload.rs b/pgdog/src/admin/reload.rs index e43c8a67a..3acee2a6a 100644 --- a/pgdog/src/admin/reload.rs +++ b/pgdog/src/admin/reload.rs @@ -16,7 +16,8 @@ impl Command for Reload { } async fn execute(&self) -> Result, Error> { - reload()?; + // false = NOT a force reload + reload(false)?; Ok(vec![]) } } diff --git a/pgdog/src/backend/databases.rs b/pgdog/src/backend/databases.rs index ee6839459..659d84cea 100644 --- a/pgdog/src/backend/databases.rs +++ b/pgdog/src/backend/databases.rs @@ -4,6 +4,7 @@ use std::collections::HashMap; use std::ops::Deref; use std::sync::Arc; +use crate::backend::Shard; use arc_swap::ArcSwap; use futures::future::try_join_all; use indexmap::IndexMap; @@ -20,6 +21,7 @@ use pgdog_config::{ use tracing::{debug, error, info, warn}; use crate::auth::AuthResult; +use crate::backend::Pool; use crate::backend::replication::ShardedSchemas; use crate::backend::schema::SchemaCache; use crate::config::PoolerMode; @@ -133,30 +135,34 @@ 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(()) +/// Terminates all active connections on all `Pool`s. +pub(crate) fn terminate_active_connections() { + databases() + .all() + .values() + .flat_map(Cluster::shards) + .flat_map(Shard::pools) + .for_each(Pool::cancel_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)?; diff --git a/pgdog/src/backend/pool/cluster.rs b/pgdog/src/backend/pool/cluster.rs index 315ca8a47..53fd035f0 100644 --- a/pgdog/src/backend/pool/cluster.rs +++ b/pgdog/src/backend/pool/cluster.rs @@ -8,7 +8,6 @@ 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; @@ -720,42 +719,6 @@ 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. diff --git a/pgdog/src/backend/pool/connection/binding.rs b/pgdog/src/backend/pool/connection/binding.rs index 4c32e8868..b6ed8a7e4 100644 --- a/pgdog/src/backend/pool/connection/binding.rs +++ b/pgdog/src/backend/pool/connection/binding.rs @@ -13,6 +13,7 @@ use crate::{ }; use futures::future::join_all; +use tokio_util::sync::CancellationToken; use super::*; use crate::util::safe_sleep; @@ -85,6 +86,18 @@ impl Binding { } } + /// Returns all `CancellationTokens` belonging to the `Binding`'s connected `Pool`s. + pub(crate) fn cancellation_tokens(&mut self) -> Vec { + match self { + Binding::Direct(guard, _) => vec![guard.pool.inner().cancellation_token.clone()], + Binding::MultiShard(guards, _) => guards + .iter() + .map(|guard| guard.pool.inner().cancellation_token.clone()) + .collect(), + _ => vec![], + } + } + pub(super) async fn read(&mut self) -> Result { match self { Binding::Direct(guard, _) => guard.read().await, diff --git a/pgdog/src/backend/pool/inner.rs b/pgdog/src/backend/pool/inner.rs index 014c6d19a..2f55a0046 100644 --- a/pgdog/src/backend/pool/inner.rs +++ b/pgdog/src/backend/pool/inner.rs @@ -31,8 +31,6 @@ 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 @@ -80,7 +78,6 @@ 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 a6bd2647b..e762647be 100644 --- a/pgdog/src/backend/pool/pool_impl.rs +++ b/pgdog/src/backend/pool/pool_impl.rs @@ -11,13 +11,14 @@ use parking_lot::{Mutex, RawMutex, lock_api::MutexGuard}; use pgdog_config::Role; use tokio::sync::Notify; use tokio::time::Instant; +use tokio_util::sync::CancellationToken; use tracing::{debug, error}; 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::{BackendKeyData, Liveness, Parameter, Parameters}; +use crate::net::{Liveness, Parameter, Parameters}; use super::inner::CheckInResult; use super::{ @@ -50,6 +51,7 @@ pub(crate) struct InnerSync { pub(super) lsn_stats: RwLock, pub(super) lsn_role_change: Notify, pub(super) oids: Arc, + pub(super) cancellation_token: CancellationToken, } impl std::fmt::Debug for Pool { @@ -81,6 +83,7 @@ impl Pool { lsn_stats: RwLock::new(LsnStats::default()), lsn_role_change: Notify::new(), oids, + cancellation_token: Default::default(), }), } } @@ -313,16 +316,13 @@ impl Pool { let mut to_guard = destination.lock(); // Propagate pause state so a paused database stays paused after reload. - // 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; + to_guard.paused = from_guard.paused; 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); } @@ -337,14 +337,9 @@ impl Pool { } /// Pause pool, closing all open connections. - /// 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) { + pub(crate) fn pause(&self) { let mut guard = self.lock(); - if !guard.paused && remove_on_transfer { - guard.remove_pause_on_transfer = true; - } else { - guard.dump_idle(); - } + guard.dump_idle(); guard.paused = true; } @@ -364,12 +359,6 @@ impl Pool { 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) { { @@ -407,6 +396,15 @@ impl Pool { self.comms().ready.notify_waiters(); } + /// Sets the `Pool` offline (to refuse more connections), and runs `cancel()` + /// on the Pool's `CancellationToken`, which causes active connections to terminate, + /// for and Clients to receive an `AdminTerminated` error. + pub(crate) fn cancel_active_connections(self) { + let mut guard = self.lock(); + guard.online = false; + self.inner.cancellation_token.cancel(); + } + /// Pool exclusive lock. #[inline] pub(super) fn lock(&self) -> MutexGuard<'_, RawMutex, Inner> { diff --git a/pgdog/src/backend/pool/test/mod.rs b/pgdog/src/backend/pool/test/mod.rs index b12c2eb3b..4201a3c71 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(false); + pool.pause(); // 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(false); + pool.pause(); 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(false); + source.pause(); assert!(source.lock().paused); assert!(!destination.lock().paused); diff --git a/pgdog/src/frontend/client/mod.rs b/pgdog/src/frontend/client/mod.rs index 16acef5c6..da59fbdca 100644 --- a/pgdog/src/frontend/client/mod.rs +++ b/pgdog/src/frontend/client/mod.rs @@ -7,9 +7,11 @@ use std::net::SocketAddr; use std::sync::Arc; use std::time::{Duration, Instant}; +use futures::future::select_all; use pgdog_config::users::PasswordKind; use timeouts::Timeouts; use tokio::{select, spawn}; +use tokio_util::sync::CancellationToken; use tracing::{Level as LogLevel, debug, enabled, error, info, trace, warn}; use super::{ClientRequest, Error, PreparedStatements}; @@ -523,6 +525,8 @@ impl Client { let client_state = query_engine.client_state(); + let cluster_cancellation: Vec = query_engine.cancellation_tokens(); + select! { _ = shutdown.cancelled(), if !offline => { continue; // Wake up task. @@ -534,7 +538,7 @@ impl Client { self.server_message(&mut query_engine, message).await?; } - buffer = self.buffer(client_state) => { + buffer = self.buffer(client_state, cluster_cancellation) => { let event = buffer?; // Only send requests to the backend if they are complete. @@ -633,7 +637,11 @@ impl Client { /// /// This ensures we don't check out a connection from the pool until the client /// sent a complete request. - async fn buffer(&mut self, state: State) -> Result { + async fn buffer( + &mut self, + state: State, + pool_cancellation_tokens: Vec, + ) -> Result { self.client_request.clear(); // Only start timer once we receive the first message. @@ -654,23 +662,39 @@ impl Client { .timeouts .client_idle_timeout(&state, &self.client_request); - let message = - match safe_timeout(idle_timeout, self.stream_buffer.read(&mut self.stream)).await { - Err(_) => { - self.stream - .fatal(ErrorResponse::client_idle_timeout(idle_timeout, &state)) - .await?; - return Ok(BufferEvent::DisconnectAbrupt); - } + let message = select! { + message = safe_timeout(idle_timeout, self.stream_buffer.read(&mut self.stream)) => { + message + } + // If any of the `CancellationTokens `trigger, exit early. Currently used for admin `FORCE_RELOAD`. + // If this returns an Error, it'll be propagated up to `Client`'s [`Box::pin(self.run())`] + // which will disconnect the `Client` (and `QueryEngine` transactions) + _ = async { select_all( + pool_cancellation_tokens.iter() + .map(|cancellation_token| Box::pin(cancellation_token.cancelled()))) + .await }, + if !pool_cancellation_tokens.is_empty() => { + return Err(Error::AdminTermination) + } + }; - Ok(Ok(message)) => message.stream(self.streaming).frontend(), - Ok(Err(err)) => { - if let Some(response) = err.as_fatal_error_response() { - self.stream.fatal(response).await?; - } - return Ok(BufferEvent::DisconnectAbrupt); + let message = match message { + Err(_) => { + self.stream + .fatal(ErrorResponse::client_idle_timeout(idle_timeout, &state)) + .await?; + return Ok(BufferEvent::DisconnectAbrupt); + } + + Ok(Ok(message)) => message.stream(self.streaming).frontend(), + + Ok(Err(err)) => { + if let Some(response) = err.as_fatal_error_response() { + self.stream.fatal(response).await?; } - }; + return Ok(BufferEvent::DisconnectAbrupt); + } + }; if timer.is_none() { timer = Some(Instant::now()); diff --git a/pgdog/src/frontend/client/query_engine/mod.rs b/pgdog/src/frontend/client/query_engine/mod.rs index 9e0a15696..c71617cb1 100644 --- a/pgdog/src/frontend/client/query_engine/mod.rs +++ b/pgdog/src/frontend/client/query_engine/mod.rs @@ -12,6 +12,7 @@ use crate::{ }; use fnv::FnvHashMap; use temp_table::TempTableState; +use tokio_util::sync::CancellationToken; use tracing::debug; pub(crate) mod advisory_lock; @@ -125,6 +126,11 @@ impl QueryEngine { self.stats.state } + /// Fetch all `CancellationToken`s for the backend. + pub(crate) fn cancellation_tokens(&mut self) -> Vec { + self.backend.cancellation_tokens() + } + /// Handle client request. pub(crate) async fn handle( &mut self, diff --git a/pgdog/src/frontend/client/query_engine/query.rs b/pgdog/src/frontend/client/query_engine/query.rs index 87bbbcc6d..3db7becd0 100644 --- a/pgdog/src/frontend/client/query_engine/query.rs +++ b/pgdog/src/frontend/client/query_engine/query.rs @@ -1,3 +1,5 @@ +use futures::future::select_all; +use tokio_util::sync::CancellationToken; use tracing::{info, trace}; use crate::{ @@ -65,12 +67,29 @@ impl QueryEngine { } } + let pool_cancellation_tokens: Vec = self.backend.cancellation_tokens(); let query_timeout = context.timeouts.query_timeout(&State::Active); - let result = safe_timeout( - query_timeout, - self.client_server_exchange(context, query_planner), - ) - .await; + + let result = tokio::select! { + result = safe_timeout( + query_timeout, + self.client_server_exchange(context, query_planner), + ) => { + result + } + // If any of the cancellation tokens trigger, exit early. Currently used for admin FORCE_RELOAD. + // If this returns an Error, it'll be propagated up to Client's Box::pin(self.run()) + // which will disconnect the client (and QueryEngine transactions) + _ = async { select_all( + pool_cancellation_tokens.iter() + .map(|cancellation_token| Box::pin(cancellation_token.cancelled()))) + .await }, + if !pool_cancellation_tokens.is_empty() => { + // I don't think we need to force-close here. After Databases::terminate_active_connections, + // we shutdown() the pools, and it should be handled there. + return Err(Error::AdminTermination); + } + }; match result { Ok(response) => response?, diff --git a/pgdog/src/frontend/client/test/mod.rs b/pgdog/src/frontend/client/test/mod.rs index 5b9906827..58bd83ed7 100644 --- a/pgdog/src/frontend/client/test/mod.rs +++ b/pgdog/src/frontend/client/test/mod.rs @@ -253,7 +253,7 @@ async fn test_abrupt_disconnect() { drop(conn); - let event = client.buffer(State::Idle).await.unwrap(); + let event = client.buffer(State::Idle, vec![]).await.unwrap(); assert_eq!(event, BufferEvent::DisconnectAbrupt); assert!(client.client_request.messages.is_empty()); @@ -272,7 +272,7 @@ async fn test_client_idle_timeout() { set(config).unwrap(); let start = Instant::now(); - let res = client.buffer(State::Idle).await.unwrap(); + let res = client.buffer(State::Idle, vec![]).await.unwrap(); assert_eq!(res, BufferEvent::DisconnectAbrupt); let err = read_one!(conn); @@ -283,7 +283,7 @@ async fn test_client_idle_timeout() { assert!( timeout( Duration::from_millis(50), - client.buffer(State::IdleInTransaction) + client.buffer(State::IdleInTransaction, vec![]) ) .await .is_err() @@ -635,7 +635,7 @@ async fn test_query_timeout() { let buf = buffer!({ Query::new("SELECT pg_sleep(0.2)") }); conn.write_all(&buf).await.unwrap(); - client.buffer(State::Idle).await.unwrap(); + client.buffer(State::Idle, vec![]).await.unwrap(); let result = client.client_messages(&mut engine).await; assert!(result.is_err()); diff --git a/pgdog/src/frontend/client/test/test_client.rs b/pgdog/src/frontend/client/test/test_client.rs index 11fb6214e..feac7a71a 100644 --- a/pgdog/src/frontend/client/test/test_client.rs +++ b/pgdog/src/frontend/client/test/test_client.rs @@ -237,7 +237,9 @@ impl TestClient { /// Process a request. pub(crate) async fn try_process(&mut self) -> Result<(), Box> { - self.client.buffer(self.engine.stats().state).await?; + self.client + .buffer(self.engine.stats().state, vec![]) + .await?; self.client.client_messages(&mut self.engine).await?; Ok(()) diff --git a/pgdog/src/frontend/error.rs b/pgdog/src/frontend/error.rs index 92e6950a4..6a8ced79c 100644 --- a/pgdog/src/frontend/error.rs +++ b/pgdog/src/frontend/error.rs @@ -57,6 +57,9 @@ pub(crate) enum Error { #[error("multi-tuple insert requires multi-shard binding")] MultiShardRequired, + #[error("terminating connection due to administrator command")] + AdminTermination, + // FIXME: layer errors better so we don't have // to reach so deep into a module. #[error("{0}")] diff --git a/pgdog/src/frontend/listener.rs b/pgdog/src/frontend/listener.rs index c2d6fb8f8..967ff128d 100644 --- a/pgdog/src/frontend/listener.rs +++ b/pgdog/src/frontend/listener.rs @@ -118,7 +118,7 @@ impl Listener { } _ = sighup.listen() => { - if let Err(err) = reload() { + if let Err(err) = reload(false) { error!("configuration reload error: {}", err); } } diff --git a/pgdog/src/net/messages/error_response.rs b/pgdog/src/net/messages/error_response.rs index 2b504d925..7aedacd50 100644 --- a/pgdog/src/net/messages/error_response.rs +++ b/pgdog/src/net/messages/error_response.rs @@ -207,6 +207,19 @@ impl ErrorResponse { } } + /// Terminating due to admin command (e.g. FORCE_RELOAD) + pub(crate) fn admin_termination() -> ErrorResponse { + ErrorResponse { + severity: "FATAL".into(), + code: "57P01".into(), + message: "terminating connection due to administrator command".into(), + detail: None, + context: None, + file: None, + routine: None, + } + } + pub(crate) fn syntax>(err: T) -> ErrorResponse { Self { severity: "ERROR".into(), @@ -260,6 +273,9 @@ impl ErrorResponse { use crate::backend::Error as BackendError; if let FrontendError::Backend(BackendError::ExecutionError(err)) = err { *(err.clone()) + } else if let FrontendError::AdminTermination = err { + // Allows us to set a custom code (to identically represent the same Postgres error) + ErrorResponse::admin_termination() } else { Self { severity: "FATAL".into(), From 0a792ef74ef7b273e3fa8519c6fcaf8700c188a5 Mon Sep 17 00:00:00 2001 From: jkaczman Date: Wed, 9 Sep 2026 13:59:43 -0400 Subject: [PATCH 6/8] Pool cancellation token -> Cluster --- pgdog/src/backend/databases.rs | 13 ++++--------- pgdog/src/backend/pool/cluster.rs | 19 +++++++++++++++++++ pgdog/src/backend/pool/connection/binding.rs | 13 ------------- pgdog/src/backend/pool/pool_impl.rs | 11 +++-------- pgdog/src/frontend/client/mod.rs | 15 ++++++--------- pgdog/src/frontend/client/query_engine/mod.rs | 19 +++++++++++++------ .../src/frontend/client/query_engine/query.rs | 18 ++++++++++-------- pgdog/src/frontend/client/test/mod.rs | 8 ++++---- pgdog/src/frontend/client/test/test_client.rs | 4 +--- 9 files changed, 60 insertions(+), 60 deletions(-) diff --git a/pgdog/src/backend/databases.rs b/pgdog/src/backend/databases.rs index 659d84cea..be0729749 100644 --- a/pgdog/src/backend/databases.rs +++ b/pgdog/src/backend/databases.rs @@ -1,10 +1,5 @@ //! Databases behind pgDog. -use std::collections::HashMap; -use std::ops::Deref; -use std::sync::Arc; - -use crate::backend::Shard; use arc_swap::ArcSwap; use futures::future::try_join_all; use indexmap::IndexMap; @@ -18,10 +13,12 @@ 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; -use crate::backend::Pool; use crate::backend::replication::ShardedSchemas; use crate::backend::schema::SchemaCache; use crate::config::PoolerMode; @@ -140,9 +137,7 @@ pub(crate) fn terminate_active_connections() { databases() .all() .values() - .flat_map(Cluster::shards) - .flat_map(Shard::pools) - .for_each(Pool::cancel_active_connections); + .for_each(Cluster::terminate_active_connections); } /// Re-create pools from config. diff --git a/pgdog/src/backend/pool/cluster.rs b/pgdog/src/backend/pool/cluster.rs index 53fd035f0..b603e4c32 100644 --- a/pgdog/src/backend/pool/cluster.rs +++ b/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>, 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(); diff --git a/pgdog/src/backend/pool/connection/binding.rs b/pgdog/src/backend/pool/connection/binding.rs index b6ed8a7e4..4c32e8868 100644 --- a/pgdog/src/backend/pool/connection/binding.rs +++ b/pgdog/src/backend/pool/connection/binding.rs @@ -13,7 +13,6 @@ use crate::{ }; use futures::future::join_all; -use tokio_util::sync::CancellationToken; use super::*; use crate::util::safe_sleep; @@ -86,18 +85,6 @@ impl Binding { } } - /// Returns all `CancellationTokens` belonging to the `Binding`'s connected `Pool`s. - pub(crate) fn cancellation_tokens(&mut self) -> Vec { - match self { - Binding::Direct(guard, _) => vec![guard.pool.inner().cancellation_token.clone()], - Binding::MultiShard(guards, _) => guards - .iter() - .map(|guard| guard.pool.inner().cancellation_token.clone()) - .collect(), - _ => vec![], - } - } - pub(super) async fn read(&mut self) -> Result { match self { Binding::Direct(guard, _) => guard.read().await, diff --git a/pgdog/src/backend/pool/pool_impl.rs b/pgdog/src/backend/pool/pool_impl.rs index e762647be..e78236057 100644 --- a/pgdog/src/backend/pool/pool_impl.rs +++ b/pgdog/src/backend/pool/pool_impl.rs @@ -11,7 +11,6 @@ use parking_lot::{Mutex, RawMutex, lock_api::MutexGuard}; use pgdog_config::Role; use tokio::sync::Notify; use tokio::time::Instant; -use tokio_util::sync::CancellationToken; use tracing::{debug, error}; use crate::backend::pool::LsnStats; @@ -51,7 +50,6 @@ pub(crate) struct InnerSync { pub(super) lsn_stats: RwLock, pub(super) lsn_role_change: Notify, pub(super) oids: Arc, - pub(super) cancellation_token: CancellationToken, } impl std::fmt::Debug for Pool { @@ -83,7 +81,6 @@ impl Pool { lsn_stats: RwLock::new(LsnStats::default()), lsn_role_change: Notify::new(), oids, - cancellation_token: Default::default(), }), } } @@ -396,13 +393,11 @@ impl Pool { self.comms().ready.notify_waiters(); } - /// Sets the `Pool` offline (to refuse more connections), and runs `cancel()` - /// on the Pool's `CancellationToken`, which causes active connections to terminate, - /// for and Clients to receive an `AdminTerminated` error. - pub(crate) fn cancel_active_connections(self) { + /// Sets the `Pool` offline (to refuse more connections) + /// Does not dump idle connections or shutdown. + pub(crate) fn set_offline(self) { let mut guard = self.lock(); guard.online = false; - self.inner.cancellation_token.cancel(); } /// Pool exclusive lock. diff --git a/pgdog/src/frontend/client/mod.rs b/pgdog/src/frontend/client/mod.rs index 5614c07b5..52cf07d18 100644 --- a/pgdog/src/frontend/client/mod.rs +++ b/pgdog/src/frontend/client/mod.rs @@ -7,7 +7,6 @@ use std::net::SocketAddr; use std::sync::Arc; use std::time::{Duration, Instant}; -use futures::future::select_all; use pgdog_config::users::PasswordKind; use timeouts::Timeouts; use tokio::{select, spawn}; @@ -529,7 +528,8 @@ impl Client { let client_state = query_engine.client_state(); - let cluster_cancellation: Vec = query_engine.cancellation_tokens(); + let cluster_cancellation: Option = + query_engine.get_cancellation_token(); select! { _ = shutdown.cancelled(), if !offline => { @@ -542,7 +542,7 @@ impl Client { self.server_message(&mut query_engine, message).await?; } - buffer = self.buffer(client_state, cluster_cancellation) => { + buffer = self.buffer(client_state, cluster_cancellation.as_ref()) => { let event = buffer?; // Only send requests to the backend if they are complete. @@ -644,7 +644,7 @@ impl Client { async fn buffer( &mut self, state: State, - pool_cancellation_tokens: Vec, + cluster_cancellation: Option<&CancellationToken>, ) -> Result { self.client_request.clear(); @@ -673,11 +673,8 @@ impl Client { // If any of the `CancellationTokens `trigger, exit early. Currently used for admin `FORCE_RELOAD`. // If this returns an Error, it'll be propagated up to `Client`'s [`Box::pin(self.run())`] // which will disconnect the `Client` (and `QueryEngine` transactions) - _ = async { select_all( - pool_cancellation_tokens.iter() - .map(|cancellation_token| Box::pin(cancellation_token.cancelled()))) - .await }, - if !pool_cancellation_tokens.is_empty() => { + _ = async { cluster_cancellation.unwrap().cancelled().await }, + if cluster_cancellation.is_some() => { return Err(Error::AdminTermination) } }; diff --git a/pgdog/src/frontend/client/query_engine/mod.rs b/pgdog/src/frontend/client/query_engine/mod.rs index 5b0e1e81f..ae6d321e6 100644 --- a/pgdog/src/frontend/client/query_engine/mod.rs +++ b/pgdog/src/frontend/client/query_engine/mod.rs @@ -1,5 +1,8 @@ use crate::{ - backend::pool::{Connection, Request}, + backend::{ + Cluster, + pool::{Connection, Request}, + }, config::config, frontend::{ BufferedQuery, Client, ClientComms, Command, DiscardTarget, Error, Router, RouterContext, @@ -10,6 +13,7 @@ use crate::{ net::{ErrorResponse, Message, Parameters}, state::State, }; +use tokio_util::sync::CancellationToken; use tracing::debug; pub(crate) mod advisory_lock; @@ -109,6 +113,14 @@ impl QueryEngine { Self::new(&client.params, &client.comms, client.admin) } + /// Fetch the `CancellationToken` for the backend (if any) + pub(crate) fn get_cancellation_token(&mut self) -> Option { + self.backend + .cluster() + .ok() + .map(Cluster::get_cancellation_token) + } + /// Wait for an async message from the backend. pub(crate) async fn read_backend(&mut self) -> Result { Ok(self.backend.read().await?) @@ -124,11 +136,6 @@ impl QueryEngine { self.stats.state } - /// Fetch all `CancellationToken`s for the backend. - pub(crate) fn cancellation_tokens(&mut self) -> Vec { - self.backend.cancellation_tokens() - } - /// Handle client request. pub(crate) async fn handle( &mut self, diff --git a/pgdog/src/frontend/client/query_engine/query.rs b/pgdog/src/frontend/client/query_engine/query.rs index 3d3c19b2d..8481aee7d 100644 --- a/pgdog/src/frontend/client/query_engine/query.rs +++ b/pgdog/src/frontend/client/query_engine/query.rs @@ -1,8 +1,8 @@ -use futures::future::select_all; use tokio_util::sync::CancellationToken; use tracing::{info, trace}; use crate::{ + backend::Cluster, frontend::{ client::TransactionType, router::parser::{explain_trace::ExplainTrace, rewrite::statement::plan::RewriteResult}, @@ -67,7 +67,12 @@ impl QueryEngine { } } - let pool_cancellation_tokens: Vec = self.backend.cancellation_tokens(); + let cluster_cancellation_token: Option = self + .backend + .cluster() + .ok() + .map(Cluster::get_cancellation_token); + let query_timeout = context.timeouts.query_timeout(&State::Active); let result = tokio::select! { @@ -77,14 +82,11 @@ impl QueryEngine { ) => { result } - // If any of the cancellation tokens trigger, exit early. Currently used for admin FORCE_RELOAD. + // If any of the cancellation tokens trigger, exit early. Currently us≤≤ed for admin FORCE_RELOAD. // If this returns an Error, it'll be propagated up to Client's Box::pin(self.run()) // which will disconnect the client (and QueryEngine transactions) - _ = async { select_all( - pool_cancellation_tokens.iter() - .map(|cancellation_token| Box::pin(cancellation_token.cancelled()))) - .await }, - if !pool_cancellation_tokens.is_empty() => { + _ = async {cluster_cancellation_token.unwrap().cancelled().await }, + if cluster_cancellation_token.is_some() => { // I don't think we need to force-close here. After Databases::terminate_active_connections, // we shutdown() the pools, and it should be handled there. return Err(Error::AdminTermination); diff --git a/pgdog/src/frontend/client/test/mod.rs b/pgdog/src/frontend/client/test/mod.rs index 58bd83ed7..89e94df51 100644 --- a/pgdog/src/frontend/client/test/mod.rs +++ b/pgdog/src/frontend/client/test/mod.rs @@ -253,7 +253,7 @@ async fn test_abrupt_disconnect() { drop(conn); - let event = client.buffer(State::Idle, vec![]).await.unwrap(); + let event = client.buffer(State::Idle, None).await.unwrap(); assert_eq!(event, BufferEvent::DisconnectAbrupt); assert!(client.client_request.messages.is_empty()); @@ -272,7 +272,7 @@ async fn test_client_idle_timeout() { set(config).unwrap(); let start = Instant::now(); - let res = client.buffer(State::Idle, vec![]).await.unwrap(); + let res = client.buffer(State::Idle, None).await.unwrap(); assert_eq!(res, BufferEvent::DisconnectAbrupt); let err = read_one!(conn); @@ -283,7 +283,7 @@ async fn test_client_idle_timeout() { assert!( timeout( Duration::from_millis(50), - client.buffer(State::IdleInTransaction, vec![]) + client.buffer(State::IdleInTransaction, None) ) .await .is_err() @@ -635,7 +635,7 @@ async fn test_query_timeout() { let buf = buffer!({ Query::new("SELECT pg_sleep(0.2)") }); conn.write_all(&buf).await.unwrap(); - client.buffer(State::Idle, vec![]).await.unwrap(); + client.buffer(State::Idle, None).await.unwrap(); let result = client.client_messages(&mut engine).await; assert!(result.is_err()); diff --git a/pgdog/src/frontend/client/test/test_client.rs b/pgdog/src/frontend/client/test/test_client.rs index feac7a71a..552641611 100644 --- a/pgdog/src/frontend/client/test/test_client.rs +++ b/pgdog/src/frontend/client/test/test_client.rs @@ -237,9 +237,7 @@ impl TestClient { /// Process a request. pub(crate) async fn try_process(&mut self) -> Result<(), Box> { - self.client - .buffer(self.engine.stats().state, vec![]) - .await?; + self.client.buffer(self.engine.stats().state, None).await?; self.client.client_messages(&mut self.engine).await?; Ok(()) From 39f0a1031a593f64fb40f2488ffd2374aad2481e Mon Sep 17 00:00:00 2001 From: jkaczman Date: Wed, 9 Sep 2026 14:31:05 -0400 Subject: [PATCH 7/8] Cancel query on Postgres side (+test) --- .../rust/tests/integration/admin_reload.rs | 15 +++++++++++++ pgdog/src/backend/pool/connection/mod.rs | 21 ++++++++++++++++++- .../src/frontend/client/query_engine/query.rs | 12 +++++++---- 3 files changed, 43 insertions(+), 5 deletions(-) diff --git a/integration/rust/tests/integration/admin_reload.rs b/integration/rust/tests/integration/admin_reload.rs index 305bdabb5..5e7cc7f51 100644 --- a/integration/rust/tests/integration/admin_reload.rs +++ b/integration/rust/tests/integration/admin_reload.rs @@ -1,3 +1,4 @@ +use sqlx::Row; use std::time::Duration; use tokio::time::{Instant, sleep}; @@ -92,5 +93,19 @@ async fn admin_force_reload_test() { // 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)")); + } } } diff --git a/pgdog/src/backend/pool/connection/mod.rs b/pgdog/src/backend/pool/connection/mod.rs index 0c41d97c6..0771c1b1f 100644 --- a/pgdog/src/backend/pool/connection/mod.rs +++ b/pgdog/src/backend/pool/connection/mod.rs @@ -1,5 +1,6 @@ //! Server connection requested by a frontend. +use futures::future::try_join_all; use mirror::MirrorHandler; use pgdog_config::users::PasswordKind; use tokio::select; @@ -22,7 +23,7 @@ use crate::{ }; use super::{ - super::{Error, pool::Guard}, + super::{Error, Server, pool::Guard}, Address, Cluster, Request, }; @@ -466,6 +467,24 @@ 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(()) + } + /// Get cluster if any. #[inline] pub(crate) fn cluster(&self) -> Result<&Cluster, Error> { diff --git a/pgdog/src/frontend/client/query_engine/query.rs b/pgdog/src/frontend/client/query_engine/query.rs index 8481aee7d..1c4a74f6a 100644 --- a/pgdog/src/frontend/client/query_engine/query.rs +++ b/pgdog/src/frontend/client/query_engine/query.rs @@ -82,13 +82,17 @@ impl QueryEngine { ) => { result } - // If any of the cancellation tokens trigger, exit early. Currently us≤≤ed for admin FORCE_RELOAD. + // If the cluster's cancellation token triggers, exit early. Currently used for admin FORCE_RELOAD. // If this returns an Error, it'll be propagated up to Client's Box::pin(self.run()) // which will disconnect the client (and QueryEngine transactions) - _ = async {cluster_cancellation_token.unwrap().cancelled().await }, + _ = async { cluster_cancellation_token.unwrap().cancelled().await }, if cluster_cancellation_token.is_some() => { - // I don't think we need to force-close here. After Databases::terminate_active_connections, - // we shutdown() the pools, and it should be handled there. + // Postgres is still running the query. Send a cancellation request before we stop on our end. + if let Err(err) = self.backend.cancel_query().await { + // Tell the administrator that we failed to cancel the query. + error!("failed to cancel query on admin termination: {err}"); + } + self.backend.force_close(); return Err(Error::AdminTermination); } }; From 009788260c8b9ea3e6a738e017d76bb4b0da06b8 Mon Sep 17 00:00:00 2001 From: jkaczman Date: Wed, 9 Sep 2026 17:32:34 -0400 Subject: [PATCH 8/8] Fix contention found in benchmarking; use child_token() instead of sharing a Cluster token --- pgdog/src/backend/pool/connection/mod.rs | 11 +++++++++++ pgdog/src/frontend/client/mod.rs | 10 ++++------ pgdog/src/frontend/client/query_engine/mod.rs | 14 ++++---------- .../src/frontend/client/query_engine/query.rs | 11 ++--------- pgdog/src/frontend/client/test/mod.rs | 18 ++++++++++++++---- pgdog/src/frontend/client/test/test_client.rs | 5 ++++- 6 files changed, 39 insertions(+), 30 deletions(-) diff --git a/pgdog/src/backend/pool/connection/mod.rs b/pgdog/src/backend/pool/connection/mod.rs index 0771c1b1f..dbc1a01a5 100644 --- a/pgdog/src/backend/pool/connection/mod.rs +++ b/pgdog/src/backend/pool/connection/mod.rs @@ -4,6 +4,7 @@ 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::{ @@ -52,6 +53,9 @@ pub(crate) struct Connection { database: String, binding: Binding, cluster: Option, + /// 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, pub_sub: PubSubClient, } @@ -66,6 +70,7 @@ impl Connection { Binding::NotConnected }, cluster: None, + cancellation_token: CancellationToken::new(), user: user.to_owned(), database: database.to_owned(), mirrors: vec![], @@ -397,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 @@ -485,6 +491,11 @@ impl Connection { 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> { diff --git a/pgdog/src/frontend/client/mod.rs b/pgdog/src/frontend/client/mod.rs index 52cf07d18..9d35b194c 100644 --- a/pgdog/src/frontend/client/mod.rs +++ b/pgdog/src/frontend/client/mod.rs @@ -528,8 +528,7 @@ impl Client { let client_state = query_engine.client_state(); - let cluster_cancellation: Option = - query_engine.get_cancellation_token(); + let cancellation_token = query_engine.cancellation_token(); select! { _ = shutdown.cancelled(), if !offline => { @@ -542,7 +541,7 @@ impl Client { self.server_message(&mut query_engine, message).await?; } - buffer = self.buffer(client_state, cluster_cancellation.as_ref()) => { + buffer = self.buffer(client_state, &cancellation_token) => { let event = buffer?; // Only send requests to the backend if they are complete. @@ -644,7 +643,7 @@ impl Client { async fn buffer( &mut self, state: State, - cluster_cancellation: Option<&CancellationToken>, + cancellation_token: &CancellationToken, ) -> Result { self.client_request.clear(); @@ -673,8 +672,7 @@ impl Client { // If any of the `CancellationTokens `trigger, exit early. Currently used for admin `FORCE_RELOAD`. // If this returns an Error, it'll be propagated up to `Client`'s [`Box::pin(self.run())`] // which will disconnect the `Client` (and `QueryEngine` transactions) - _ = async { cluster_cancellation.unwrap().cancelled().await }, - if cluster_cancellation.is_some() => { + _ = cancellation_token.cancelled() => { return Err(Error::AdminTermination) } }; diff --git a/pgdog/src/frontend/client/query_engine/mod.rs b/pgdog/src/frontend/client/query_engine/mod.rs index ae6d321e6..9196bcf8b 100644 --- a/pgdog/src/frontend/client/query_engine/mod.rs +++ b/pgdog/src/frontend/client/query_engine/mod.rs @@ -1,8 +1,5 @@ use crate::{ - backend::{ - Cluster, - pool::{Connection, Request}, - }, + backend::pool::{Connection, Request}, config::config, frontend::{ BufferedQuery, Client, ClientComms, Command, DiscardTarget, Error, Router, RouterContext, @@ -113,12 +110,9 @@ impl QueryEngine { Self::new(&client.params, &client.comms, client.admin) } - /// Fetch the `CancellationToken` for the backend (if any) - pub(crate) fn get_cancellation_token(&mut self) -> Option { - self.backend - .cluster() - .ok() - .map(Cluster::get_cancellation_token) + /// Token cancelled when an admin terminates this client's cluster (FORCE_RELOAD). + pub(crate) fn cancellation_token(&self) -> CancellationToken { + self.backend.cancellation_token() } /// Wait for an async message from the backend. diff --git a/pgdog/src/frontend/client/query_engine/query.rs b/pgdog/src/frontend/client/query_engine/query.rs index 1c4a74f6a..aa88e8f99 100644 --- a/pgdog/src/frontend/client/query_engine/query.rs +++ b/pgdog/src/frontend/client/query_engine/query.rs @@ -1,8 +1,6 @@ -use tokio_util::sync::CancellationToken; use tracing::{info, trace}; use crate::{ - backend::Cluster, frontend::{ client::TransactionType, router::parser::{explain_trace::ExplainTrace, rewrite::statement::plan::RewriteResult}, @@ -67,11 +65,7 @@ impl QueryEngine { } } - let cluster_cancellation_token: Option = self - .backend - .cluster() - .ok() - .map(Cluster::get_cancellation_token); + let cancellation_token = self.backend.cancellation_token(); let query_timeout = context.timeouts.query_timeout(&State::Active); @@ -85,8 +79,7 @@ impl QueryEngine { // If the cluster's cancellation token triggers, exit early. Currently used for admin FORCE_RELOAD. // If this returns an Error, it'll be propagated up to Client's Box::pin(self.run()) // which will disconnect the client (and QueryEngine transactions) - _ = async { cluster_cancellation_token.unwrap().cancelled().await }, - if cluster_cancellation_token.is_some() => { + _ = cancellation_token.cancelled() => { // Postgres is still running the query. Send a cancellation request before we stop on our end. if let Err(err) = self.backend.cancel_query().await { // Tell the administrator that we failed to cancel the query. diff --git a/pgdog/src/frontend/client/test/mod.rs b/pgdog/src/frontend/client/test/mod.rs index 89e94df51..febb77a5e 100644 --- a/pgdog/src/frontend/client/test/mod.rs +++ b/pgdog/src/frontend/client/test/mod.rs @@ -1,4 +1,5 @@ use std::time::{Duration, Instant}; +use tokio_util::sync::CancellationToken; use pgdog_config::{PoolerMode, QuerySizeLimitAction}; use tokio::{ @@ -253,7 +254,10 @@ async fn test_abrupt_disconnect() { drop(conn); - let event = client.buffer(State::Idle, None).await.unwrap(); + let event = client + .buffer(State::Idle, &CancellationToken::new()) + .await + .unwrap(); assert_eq!(event, BufferEvent::DisconnectAbrupt); assert!(client.client_request.messages.is_empty()); @@ -272,7 +276,10 @@ async fn test_client_idle_timeout() { set(config).unwrap(); let start = Instant::now(); - let res = client.buffer(State::Idle, None).await.unwrap(); + let res = client + .buffer(State::Idle, &CancellationToken::new()) + .await + .unwrap(); assert_eq!(res, BufferEvent::DisconnectAbrupt); let err = read_one!(conn); @@ -283,7 +290,7 @@ async fn test_client_idle_timeout() { assert!( timeout( Duration::from_millis(50), - client.buffer(State::IdleInTransaction, None) + client.buffer(State::IdleInTransaction, &CancellationToken::new()) ) .await .is_err() @@ -635,7 +642,10 @@ async fn test_query_timeout() { let buf = buffer!({ Query::new("SELECT pg_sleep(0.2)") }); conn.write_all(&buf).await.unwrap(); - client.buffer(State::Idle, None).await.unwrap(); + client + .buffer(State::Idle, &CancellationToken::new()) + .await + .unwrap(); let result = client.client_messages(&mut engine).await; assert!(result.is_err()); diff --git a/pgdog/src/frontend/client/test/test_client.rs b/pgdog/src/frontend/client/test/test_client.rs index 552641611..d4a6cbdbe 100644 --- a/pgdog/src/frontend/client/test/test_client.rs +++ b/pgdog/src/frontend/client/test/test_client.rs @@ -1,4 +1,5 @@ use std::{fmt::Debug, ops::Deref}; +use tokio_util::sync::CancellationToken; use bytes::{BufMut, Bytes, BytesMut}; use pgdog_config::RewriteMode; @@ -237,7 +238,9 @@ impl TestClient { /// Process a request. pub(crate) async fn try_process(&mut self) -> Result<(), Box> { - self.client.buffer(self.engine.stats().state, None).await?; + self.client + .buffer(self.engine.stats().state, &CancellationToken::new()) + .await?; self.client.client_messages(&mut self.engine).await?; Ok(())