From a78a52acb57bfbea653f51e22e365af13ba5c332 Mon Sep 17 00:00:00 2001
From: Conrad Ludgate <conradludgate@gmail.com>
Date: Tue, 10 Jun 2025 13:44:30 -0700
Subject: [PATCH] add in timeout for cancellation

---
 proxy/src/cancellation.rs | 17 +++++++++++++----
 1 file changed, 13 insertions(+), 4 deletions(-)

diff --git a/proxy/src/cancellation.rs b/proxy/src/cancellation.rs
index 8ce0f1365b..036f36c7f6 100644
--- a/proxy/src/cancellation.rs
+++ b/proxy/src/cancellation.rs
@@ -1,6 +1,7 @@
 use std::convert::Infallible;
 use std::net::{IpAddr, SocketAddr};
 use std::sync::{Arc, OnceLock};
+use std::time::Duration;
 
 use anyhow::anyhow;
 use futures::FutureExt;
@@ -11,6 +12,7 @@ use redis::{Cmd, FromRedisValue, Value};
 use serde::{Deserialize, Serialize};
 use thiserror::Error;
 use tokio::net::TcpStream;
+use tokio::time::timeout;
 use tracing::{debug, error, info};
 
 use crate::auth::AuthError;
@@ -237,10 +239,17 @@ impl CancellationHandler {
             return Err(CancelError::InternalError);
         };
 
-        let result = tx.call((guard, op)).await.map_err(|e| {
-            tracing::warn!("failed to receive GetCancelData response: {e}");
-            CancelError::InternalError
-        })?;
+        const TIMEOUT: Duration = Duration::from_secs(5);
+        let result = timeout(TIMEOUT, tx.call((guard, op)))
+            .await
+            .map_err(|_| {
+                tracing::warn!("timed out waiting to receive GetCancelData response");
+                CancelError::RateLimit
+            })?
+            .map_err(|e| {
+                tracing::warn!("failed to receive GetCancelData response: {e}");
+                CancelError::InternalError
+            })?;
 
         let cancel_state_str = String::from_owned_redis_value(result).map_err(|e| {
             tracing::warn!("failed to receive GetCancelData response: {e}");