diff --git a/libs/proxy/tokio-postgres2/src/cancel_token.rs b/libs/proxy/tokio-postgres2/src/cancel_token.rs index 61e1e52d66..9bb0b094f1 100644 --- a/libs/proxy/tokio-postgres2/src/cancel_token.rs +++ b/libs/proxy/tokio-postgres2/src/cancel_token.rs @@ -17,7 +17,7 @@ pub struct CancelToken { /// The capability to request cancellation of in-progress queries on a /// connection. -#[derive(Clone, Serialize, Deserialize)] +#[derive(Debug, Clone, Serialize, Deserialize)] pub struct RawCancelToken { pub ssl_mode: SslMode, pub process_id: i32, diff --git a/proxy/src/cancellation.rs b/proxy/src/cancellation.rs index 987b8545c0..9cdcdba94f 100644 --- a/proxy/src/cancellation.rs +++ b/proxy/src/cancellation.rs @@ -1,3 +1,4 @@ +use std::convert::Infallible; use std::net::{IpAddr, SocketAddr}; use std::sync::Arc; @@ -27,13 +28,15 @@ use crate::redis::kv_ops::RedisKVClient; type IpSubnetKey = IpNet; -const CANCEL_KEY_TTL: i64 = 1_209_600; // 2 weeks cancellation key expire time +const CANCEL_KEY_TTL: std::time::Duration = std::time::Duration::from_secs(600); +const CANCEL_KEY_REFRESH: std::time::Duration = std::time::Duration::from_secs(570); // Message types for sending through mpsc channel pub enum CancelKeyOp { StoreCancelKey { key: String, value: String, + resp_tx: oneshot::Sender>, _guard: CancelChannelSizeGuard<'static>, expire: i64, // TTL for key }, @@ -42,11 +45,6 @@ pub enum CancelKeyOp { resp_tx: oneshot::Sender>, _guard: CancelChannelSizeGuard<'static>, }, - RemoveCancelKey { - key: String, - field: String, - _guard: CancelChannelSizeGuard<'static>, - }, } pub struct Pipeline { @@ -111,10 +109,12 @@ impl CancelKeyOp { CancelKeyOp::StoreCancelKey { key, value, + resp_tx, _guard, expire, } => { - pipe.add_command_no_reply(Cmd::hset(&key, "data", value)); + let reply = CancelReplyOp::StoreCancelKey { resp_tx, _guard }; + pipe.add_command_with_reply(Cmd::hset(&key, "data", value), reply); pipe.add_command_no_reply(Cmd::expire(key, expire)); } CancelKeyOp::GetCancelData { @@ -125,15 +125,16 @@ impl CancelKeyOp { let reply = CancelReplyOp::GetCancelData { resp_tx, _guard }; pipe.add_command_with_reply(Cmd::hget(key, "data"), reply); } - CancelKeyOp::RemoveCancelKey { key, field, _guard } => { - pipe.add_command_no_reply(Cmd::hdel(key, field)); - } } } } // Message types for sending through mpsc channel pub enum CancelReplyOp { + StoreCancelKey { + resp_tx: oneshot::Sender>, + _guard: CancelChannelSizeGuard<'static>, + }, GetCancelData { resp_tx: oneshot::Sender>, _guard: CancelChannelSizeGuard<'static>, @@ -143,6 +144,12 @@ pub enum CancelReplyOp { impl CancelReplyOp { fn send_err(self, e: anyhow::Error) { match self { + CancelReplyOp::StoreCancelKey { resp_tx, _guard } => { + resp_tx + .send(Err(e)) + .inspect_err(|_| tracing::debug!("could not send reply")) + .ok(); + } CancelReplyOp::GetCancelData { resp_tx, _guard } => { resp_tx .send(Err(e)) @@ -154,6 +161,14 @@ impl CancelReplyOp { fn send_value(self, v: redis::Value) { match self { + CancelReplyOp::StoreCancelKey { resp_tx, _guard } => { + let send = + FromRedisValue::from_owned_redis_value(v).context("could not parse value"); + resp_tx + .send(send) + .inspect_err(|_| tracing::debug!("could not send reply")) + .ok(); + } CancelReplyOp::GetCancelData { resp_tx, _guard } => { let send = FromRedisValue::from_owned_redis_value(v).context("could not parse value"); @@ -255,7 +270,7 @@ impl CancellationHandler { } } - pub(crate) fn get_key(self: &Arc) -> Session { + pub(crate) fn get_key(self: Arc) -> Session { // we intentionally generate a random "backend pid" and "secret key" here. // we use the corresponding u64 as an identifier for the // actual endpoint+pid+secret for postgres/pgbouncer. @@ -272,7 +287,7 @@ impl CancellationHandler { Session { key, redis_key, - cancellation_handler: Arc::clone(self), + cancellation_handler: self, } } @@ -400,7 +415,7 @@ impl CancellationHandler { /// This should've been a [`std::future::Future`], but /// it's impossible to name a type of an unboxed future /// (we'd need something like `#![feature(type_alias_impl_trait)]`). -#[derive(Clone, Serialize, Deserialize)] +#[derive(Debug, Clone, Serialize, Deserialize)] pub struct CancelClosure { socket_addr: SocketAddr, cancel_token: RawCancelToken, @@ -454,11 +469,24 @@ impl Session { &self.key } - // Send the store key op to the cancellation handler and set TTL for the key - pub(crate) fn write_cancel_key( - &self, - cancel_closure: CancelClosure, + /// Ensure the cancel key is continously refreshed, + /// but stop when the channel is dropped. + pub(crate) async fn maintain_cancel_key( + self, + cancel: tokio::sync::oneshot::Receiver, + cancel_closure: &CancelClosure, ) -> Result<(), CancelError> { + tokio::select! { + res = self.maintain_redis_cancel_key(cancel_closure) => match res? {}, + _ = cancel => Ok(()), + } + } + + /// Ensure the cancel key is continously refreshed. + async fn maintain_redis_cancel_key( + &self, + cancel_closure: &CancelClosure, + ) -> Result { let Some(tx) = &self.cancellation_handler.tx else { tracing::warn!("cancellation handler is not available"); return Err(CancelError::InternalError); @@ -469,42 +497,34 @@ impl Session { CancelError::InternalError })?; - let op = CancelKeyOp::StoreCancelKey { - key: self.redis_key.clone(), - value: closure_json, - _guard: Metrics::get() - .proxy - .cancel_channel_size - .guard(RedisMsgKind::HSet), - expire: CANCEL_KEY_TTL, - }; + loop { + let (resp_tx, resp_rx) = tokio::sync::oneshot::channel(); + let op = CancelKeyOp::StoreCancelKey { + key: self.redis_key.clone(), + value: closure_json.clone(), + resp_tx, + _guard: Metrics::get() + .proxy + .cancel_channel_size + .guard(RedisMsgKind::HSet), + expire: CANCEL_KEY_TTL.as_secs() as i64, + }; - let _ = tx.try_send(op).map_err(|e| { - let key = self.key; - tracing::warn!("failed to send StoreCancelKey for {key}: {e}"); - }); - Ok(()) - } + tracing::debug!( + src=%self.key, + dest=?cancel_closure.cancel_token, + "registering cancellation key" + ); - pub(crate) fn remove_cancel_key(&self) -> Result<(), CancelError> { - let Some(tx) = &self.cancellation_handler.tx else { - tracing::warn!("cancellation handler is not available"); - return Err(CancelError::InternalError); - }; + tx.send(op).await.map_err(|e| { + let key = self.key; + tracing::warn!("failed to send StoreCancelKey for {key}: {e}"); + CancelError::InternalError + })?; - let op = CancelKeyOp::RemoveCancelKey { - key: self.redis_key.clone(), - field: "data".to_string(), - _guard: Metrics::get() - .proxy - .cancel_channel_size - .guard(RedisMsgKind::HDel), - }; - - let _ = tx.try_send(op).map_err(|e| { - let key = self.key; - tracing::warn!("failed to send RemoveCancelKey for {key}: {e}"); - }); - Ok(()) + if resp_rx.await.is_ok() { + tokio::time::sleep(CANCEL_KEY_REFRESH).await; + } + } } } diff --git a/proxy/src/console_redirect_proxy.rs b/proxy/src/console_redirect_proxy.rs index 5331ea41fd..49e0c673a9 100644 --- a/proxy/src/console_redirect_proxy.rs +++ b/proxy/src/console_redirect_proxy.rs @@ -232,21 +232,22 @@ pub(crate) async fn handle_client( .or_else(|e| async { Err(stream.throw_error(e, Some(ctx)).await) }) .await?; - let cancellation_handler_clone = Arc::clone(&cancellation_handler); - let session = cancellation_handler_clone.get_key(); - - session.write_cancel_key(node.cancel_closure.clone())?; + let session = cancellation_handler.get_key(); prepare_client_connection(&node, *session.key(), &mut stream); let stream = stream.flush_and_into_inner().await?; + let cancel_closure = node.cancel_closure.clone(); + let (cancel_on_shutdown, cancel) = tokio::sync::oneshot::channel(); + tokio::spawn(async move { session.maintain_cancel_key(cancel, &cancel_closure).await }); + Ok(Some(ProxyPassthrough { client: stream, aux: node.aux.clone(), private_link_id: None, compute: node, session_id: ctx.session_id(), - cancel: session, + _cancel_on_shutdown: cancel_on_shutdown, _req: request_gauge, _conn: conn_gauge, })) diff --git a/proxy/src/pglb/passthrough.rs b/proxy/src/pglb/passthrough.rs index 6f651d383d..fa3df288be 100644 --- a/proxy/src/pglb/passthrough.rs +++ b/proxy/src/pglb/passthrough.rs @@ -1,3 +1,5 @@ +use std::convert::Infallible; + use futures::FutureExt; use smol_str::SmolStr; use tokio::io::{AsyncRead, AsyncWrite}; @@ -5,7 +7,6 @@ use tracing::debug; use utils::measured_stream::MeasuredStream; use super::copy_bidirectional::ErrorSource; -use crate::cancellation; use crate::compute::PostgresConnection; use crate::config::ComputeConfig; use crate::control_plane::messages::MetricsAuxInfo; @@ -68,7 +69,8 @@ pub(crate) struct ProxyPassthrough { pub(crate) aux: MetricsAuxInfo, pub(crate) session_id: uuid::Uuid, pub(crate) private_link_id: Option, - pub(crate) cancel: cancellation::Session, + + pub(crate) _cancel_on_shutdown: tokio::sync::oneshot::Sender, pub(crate) _req: NumConnectionRequestsGuard<'static>, pub(crate) _conn: NumClientConnectionsGuard<'static>, @@ -96,8 +98,6 @@ impl ProxyPassthrough { tracing::warn!(session_id = ?self.session_id, ?err, "could not cancel the query in the database"); } - drop(self.cancel.remove_cancel_key()); // we don't need a result. If the queue is full, we just log the error - res } } diff --git a/proxy/src/proxy/mod.rs b/proxy/src/proxy/mod.rs index 4211406f6c..980b82df36 100644 --- a/proxy/src/proxy/mod.rs +++ b/proxy/src/proxy/mod.rs @@ -372,13 +372,15 @@ pub(crate) async fn handle_client( Err(e) => Err(stream.throw_error(e, Some(ctx)).await)?, }; - let cancellation_handler_clone = Arc::clone(&cancellation_handler); - let session = cancellation_handler_clone.get_key(); + let session = cancellation_handler.get_key(); - session.write_cancel_key(node.cancel_closure.clone())?; prepare_client_connection(&node, *session.key(), &mut stream); let stream = stream.flush_and_into_inner().await?; + let cancel_closure = node.cancel_closure.clone(); + let (cancel_on_shutdown, cancel) = tokio::sync::oneshot::channel(); + tokio::spawn(async move { session.maintain_cancel_key(cancel, &cancel_closure).await }); + let private_link_id = match ctx.extra() { Some(ConnectionInfoExtra::Aws { vpce_id }) => Some(vpce_id.clone()), Some(ConnectionInfoExtra::Azure { link_id }) => Some(link_id.to_smolstr()), @@ -391,7 +393,7 @@ pub(crate) async fn handle_client( private_link_id, compute: node, session_id: ctx.session_id(), - cancel: session, + _cancel_on_shutdown: cancel_on_shutdown, _req: request_gauge, _conn: conn_gauge, }))