fix(rustproxy-http,rustproxy-passthrough): improve upstream connection cleanup and graceful tunnel shutdown
This commit is contained in:
@@ -966,11 +966,18 @@ impl HttpProxyService {
|
||||
}
|
||||
};
|
||||
|
||||
tokio::spawn(async move {
|
||||
if let Err(e) = conn.await {
|
||||
debug!("HTTP/2 upstream connection error: {}", e);
|
||||
}
|
||||
});
|
||||
// Spawn the H2 connection driver; proactively evict from pool on exit
|
||||
// so the next request gets a fresh connection instead of a dead sender.
|
||||
{
|
||||
let pool = Arc::clone(&self.connection_pool);
|
||||
let key = pool_key.clone();
|
||||
tokio::spawn(async move {
|
||||
if let Err(e) = conn.await {
|
||||
debug!("HTTP/2 upstream connection error: {}", e);
|
||||
}
|
||||
pool.remove_h2(&key);
|
||||
});
|
||||
}
|
||||
|
||||
// Clone sender for potential pool registration; register only after first request succeeds
|
||||
let sender_for_pool = sender.clone();
|
||||
@@ -1111,11 +1118,17 @@ impl HttpProxyService {
|
||||
}
|
||||
};
|
||||
|
||||
tokio::spawn(async move {
|
||||
if let Err(e) = conn.await {
|
||||
debug!("H2 retry: upstream connection error: {}", e);
|
||||
}
|
||||
});
|
||||
// Spawn the H2 connection driver; proactively evict from pool on exit.
|
||||
{
|
||||
let pool = Arc::clone(&self.connection_pool);
|
||||
let key = pool_key.clone();
|
||||
tokio::spawn(async move {
|
||||
if let Err(e) = conn.await {
|
||||
debug!("H2 retry: upstream connection error: {}", e);
|
||||
}
|
||||
pool.remove_h2(&key);
|
||||
});
|
||||
}
|
||||
|
||||
// Build request with empty body using absolute URI for H2 pseudo-headers
|
||||
let scheme = if pool_key.use_tls { "https" } else { "http" };
|
||||
@@ -1234,11 +1247,17 @@ impl HttpProxyService {
|
||||
}
|
||||
}
|
||||
Ok(Ok((mut sender, conn))) => {
|
||||
tokio::spawn(async move {
|
||||
if let Err(e) = conn.await {
|
||||
debug!("HTTP/2 upstream connection error: {}", e);
|
||||
}
|
||||
});
|
||||
// Spawn the H2 connection driver; proactively evict from pool on exit.
|
||||
{
|
||||
let pool = Arc::clone(&self.connection_pool);
|
||||
let key = pool_key.clone();
|
||||
tokio::spawn(async move {
|
||||
if let Err(e) = conn.await {
|
||||
debug!("HTTP/2 upstream connection error: {}", e);
|
||||
}
|
||||
pool.remove_h2(&key);
|
||||
});
|
||||
}
|
||||
|
||||
// Save retry state before consuming parts/body (for bodyless requests like GET)
|
||||
// Clone BEFORE removing Host — H1 fallback needs Host header
|
||||
@@ -1578,6 +1597,19 @@ impl HttpProxyService {
|
||||
|
||||
if let Some(headers) = response.headers_mut() {
|
||||
*headers = resp_parts.headers;
|
||||
|
||||
// Strip hop-by-hop headers from the backend response.
|
||||
// RFC 9113 §8.2.2 forbids connection-specific headers in HTTP/2 responses;
|
||||
// forwarding them from an H1 backend can cause H2 stream resets.
|
||||
// Mirrors the request-path stripping at the forward methods above.
|
||||
headers.remove("connection");
|
||||
headers.remove("keep-alive");
|
||||
headers.remove("proxy-connection");
|
||||
headers.remove("transfer-encoding");
|
||||
headers.remove("te");
|
||||
headers.remove("trailer");
|
||||
// Note: "upgrade" is intentionally kept — needed for WebSocket 101 responses.
|
||||
|
||||
ResponseFilter::apply_headers(route, headers, None);
|
||||
}
|
||||
|
||||
@@ -1913,6 +1945,10 @@ impl HttpProxyService {
|
||||
let last_activity = Arc::new(AtomicU64::new(0));
|
||||
let start = std::time::Instant::now();
|
||||
|
||||
// Per-connection cancellation token: the watchdog cancels this instead of
|
||||
// aborting tasks, so the copy loops can shut down gracefully (TLS close_notify).
|
||||
let ws_cancel = CancellationToken::new();
|
||||
|
||||
// For H2 WebSocket: also update the connection-level activity tracker
|
||||
// to prevent the idle watchdog from killing the H2 connection
|
||||
let conn_act_c2u = conn_activity.as_ref().map(|ca| (Arc::clone(&ca.last_activity), ca.start));
|
||||
@@ -1922,13 +1958,17 @@ impl HttpProxyService {
|
||||
let metrics_c2u = Arc::clone(&metrics);
|
||||
let route_c2u = route_id_owned.clone();
|
||||
let ip_c2u = source_ip_owned.clone();
|
||||
let wsc1 = ws_cancel.clone();
|
||||
let c2u = tokio::spawn(async move {
|
||||
let mut buf = vec![0u8; 65536];
|
||||
let mut total = 0u64;
|
||||
loop {
|
||||
let n = match cr.read(&mut buf).await {
|
||||
Ok(0) | Err(_) => break,
|
||||
Ok(n) => n,
|
||||
let n = tokio::select! {
|
||||
result = cr.read(&mut buf) => match result {
|
||||
Ok(0) | Err(_) => break,
|
||||
Ok(n) => n,
|
||||
},
|
||||
_ = wsc1.cancelled() => break,
|
||||
};
|
||||
if uw.write_all(&buf[..n]).await.is_err() {
|
||||
break;
|
||||
@@ -1940,7 +1980,11 @@ impl HttpProxyService {
|
||||
ca.store(ca_start.elapsed().as_millis() as u64, Ordering::Relaxed);
|
||||
}
|
||||
}
|
||||
let _ = uw.shutdown().await;
|
||||
// Graceful shutdown with timeout (sends TLS close_notify / TCP FIN)
|
||||
let _ = tokio::time::timeout(
|
||||
std::time::Duration::from_secs(2),
|
||||
uw.shutdown(),
|
||||
).await;
|
||||
total
|
||||
});
|
||||
|
||||
@@ -1948,13 +1992,17 @@ impl HttpProxyService {
|
||||
let metrics_u2c = Arc::clone(&metrics);
|
||||
let route_u2c = route_id_owned.clone();
|
||||
let ip_u2c = source_ip_owned.clone();
|
||||
let wsc2 = ws_cancel.clone();
|
||||
let u2c = tokio::spawn(async move {
|
||||
let mut buf = vec![0u8; 65536];
|
||||
let mut total = 0u64;
|
||||
loop {
|
||||
let n = match ur.read(&mut buf).await {
|
||||
Ok(0) | Err(_) => break,
|
||||
Ok(n) => n,
|
||||
let n = tokio::select! {
|
||||
result = ur.read(&mut buf) => match result {
|
||||
Ok(0) | Err(_) => break,
|
||||
Ok(n) => n,
|
||||
},
|
||||
_ = wsc2.cancelled() => break,
|
||||
};
|
||||
if cw.write_all(&buf[..n]).await.is_err() {
|
||||
break;
|
||||
@@ -1966,14 +2014,20 @@ impl HttpProxyService {
|
||||
ca.store(ca_start.elapsed().as_millis() as u64, Ordering::Relaxed);
|
||||
}
|
||||
}
|
||||
let _ = cw.shutdown().await;
|
||||
// Graceful shutdown with timeout (sends TLS close_notify / TCP FIN)
|
||||
let _ = tokio::time::timeout(
|
||||
std::time::Duration::from_secs(2),
|
||||
cw.shutdown(),
|
||||
).await;
|
||||
total
|
||||
});
|
||||
|
||||
// Watchdog: monitors inactivity, max lifetime, and cancellation
|
||||
// Watchdog: monitors inactivity, max lifetime, and cancellation.
|
||||
// First cancels the per-connection token for graceful shutdown (close_notify/FIN),
|
||||
// then falls back to abort if the tasks are stuck (e.g. on a blocked write_all).
|
||||
let la_watch = Arc::clone(&last_activity);
|
||||
let c2u_handle = c2u.abort_handle();
|
||||
let u2c_handle = u2c.abort_handle();
|
||||
let c2u_abort = c2u.abort_handle();
|
||||
let u2c_abort = u2c.abort_handle();
|
||||
let inactivity_timeout = ws_inactivity_timeout;
|
||||
let max_lifetime = ws_max_lifetime;
|
||||
|
||||
@@ -1985,8 +2039,6 @@ impl HttpProxyService {
|
||||
_ = tokio::time::sleep(check_interval) => {}
|
||||
_ = cancel.cancelled() => {
|
||||
debug!("WebSocket tunnel cancelled by shutdown");
|
||||
c2u_handle.abort();
|
||||
u2c_handle.abort();
|
||||
break;
|
||||
}
|
||||
}
|
||||
@@ -1994,8 +2046,6 @@ impl HttpProxyService {
|
||||
// Check max lifetime
|
||||
if start.elapsed() >= max_lifetime {
|
||||
debug!("WebSocket tunnel exceeded max lifetime, closing");
|
||||
c2u_handle.abort();
|
||||
u2c_handle.abort();
|
||||
break;
|
||||
}
|
||||
|
||||
@@ -2005,13 +2055,18 @@ impl HttpProxyService {
|
||||
let elapsed_since_activity = start.elapsed().as_millis() as u64 - current;
|
||||
if elapsed_since_activity >= inactivity_timeout.as_millis() as u64 {
|
||||
debug!("WebSocket tunnel inactive for {}ms, closing", elapsed_since_activity);
|
||||
c2u_handle.abort();
|
||||
u2c_handle.abort();
|
||||
break;
|
||||
}
|
||||
}
|
||||
last_seen = current;
|
||||
}
|
||||
// Phase 1: Signal copy loops to exit gracefully (allows close_notify/FIN)
|
||||
ws_cancel.cancel();
|
||||
// Phase 2: Wait for graceful shutdown (2s shutdown timeout + 2s margin)
|
||||
tokio::time::sleep(std::time::Duration::from_secs(4)).await;
|
||||
// Phase 3: Force-abort if still stuck (e.g. blocked on write_all)
|
||||
c2u_abort.abort();
|
||||
u2c_abort.abort();
|
||||
});
|
||||
|
||||
let bytes_in = c2u.await.unwrap_or(0);
|
||||
|
||||
Reference in New Issue
Block a user