no explicit remove, only passive ttl

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<anyhow::Result<()>>,
         _guard: CancelChannelSizeGuard<'static>,
         expire: i64, // TTL for key
     },
@@ -42,11 +45,6 @@ pub enum CancelKeyOp {
         resp_tx: oneshot::Sender<anyhow::Result<String>>,
         _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<anyhow::Result<()>>,
+        _guard: CancelChannelSizeGuard<'static>,
+    },
     GetCancelData {
         resp_tx: oneshot::Sender<anyhow::Result<String>>,
         _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<Self>) -> Session {
+    pub(crate) fn get_key(self: Arc<Self>) -> 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<Infallible>,
+        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<Infallible, CancelError> {
         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;
+            }
+        }
     }
 }

proxy/src/console_redirect_proxy.rs

@@ -232,21 +232,22 @@ pub(crate) async fn handle_client<S: AsyncRead + AsyncWrite + Unpin + Send>(
     .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,
     }))

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<S> {
     pub(crate) aux: MetricsAuxInfo,
     pub(crate) session_id: uuid::Uuid,
     pub(crate) private_link_id: Option<SmolStr>,
-    pub(crate) cancel: cancellation::Session,
+
+    pub(crate) _cancel_on_shutdown: tokio::sync::oneshot::Sender<Infallible>,
 
     pub(crate) _req: NumConnectionRequestsGuard<'static>,
     pub(crate) _conn: NumClientConnectionsGuard<'static>,
@@ -96,8 +98,6 @@ impl<S: AsyncRead + AsyncWrite + Unpin> ProxyPassthrough<S> {
             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
     }
 }

proxy/src/proxy/mod.rs

@@ -372,13 +372,15 @@ pub(crate) async fn handle_client<S: AsyncRead + AsyncWrite + Unpin + Send>(
         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<S: AsyncRead + AsyncWrite + Unpin + Send>(
         private_link_id,
         compute: node,
         session_id: ctx.session_id(),
-        cancel: session,
+        _cancel_on_shutdown: cancel_on_shutdown,
         _req: request_gauge,
         _conn: conn_gauge,
     }))