Skip to content

fix: reuse the upstream connection after a response_filter error #952

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
stareezy-1 wants to merge 1 commit into
base: main
Choose a base branch
from
Open

fix: reuse the upstream connection after a response_filter error #952

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 10 additions & 0 deletions pingora-proxy/src/proxy_common.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -308,6 +308,11 @@ impl ResponseStateMachine {
pub(crate) enum PipeState {
Active = 0,
DownstreamComplete = 1,
/// The downstream half aborted while running a response filter. The
/// upstream half drains the remaining response body so the connection
/// can still be reused (e.g. a 3xx redirect followed by the outer retry
/// loop after a `response_filter` error).
DownstreamFilterAborted = 2,
}

impl PipeState {
Expand All @@ -317,6 +322,11 @@ impl PipeState {
pub(crate) fn is_downstream_complete(raw: u8) -> bool {
raw == PipeState::DownstreamComplete as u8
}

/// Whether `raw` is [`DownstreamFilterAborted`](Self::DownstreamFilterAborted).
pub(crate) fn is_downstream_filter_aborted(raw: u8) -> bool {
raw == PipeState::DownstreamFilterAborted as u8
}
}

#[cfg(test)]
Expand Down
155 changes: 135 additions & 20 deletions pingora-proxy/src/proxy_h1.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -121,18 +121,73 @@ where
let pipe_state = Arc::new(AtomicU8::new(PipeState::Active as u8));

// start bi-directional streaming
let ret = tokio::try_join!(
self.proxy_handle_downstream(
session,
tx_downstream,
rx_upstream,
ctx,
&mut downstream_custom_message_writer,
&mut downstream_custom_message_reader,
pipe_state.clone(),
),
self.proxy_handle_upstream(client_session, tx_upstream, rx_downstream, pipe_state),
);
// Unlike try_join!, an error from the downstream half does not cancel
// the upstream half: when the downstream half aborted in a response
// filter, the upstream half is allowed to drain the remaining
// response body so the connection can still be reused (issue #866).
let mut downstream = Box::pin(self.proxy_handle_downstream(
session,
tx_downstream,
rx_upstream,
ctx,
&mut downstream_custom_message_writer,
&mut downstream_custom_message_reader,
pipe_state.clone(),
));
let mut upstream = Box::pin(self.proxy_handle_upstream(
client_session,
tx_upstream,
rx_downstream,
pipe_state.clone(),
));

let (downstream_can_reuse, upstream_can_reuse, error) = tokio::select! {
r = &mut downstream => {
match r {
Ok(downstream_can_reuse) => {
match (&mut upstream).await {
Ok(upstream_can_reuse) => (downstream_can_reuse, upstream_can_reuse, None),
Err(e) => (false, false, Some(e)),
}
}
Err(e) => {
if PipeState::is_downstream_filter_aborted(
pipe_state.load(Ordering::Acquire),
) {
// The downstream half aborted in a response filter
// (e.g. a 3xx redirect followed by the outer retry
// loop). Let the upstream half drain the remaining
// response body; the connection is reusable only if
// the drain completed.
match (&mut upstream).await {
Ok(true) => (false, true, Some(e)),
Ok(false) => (false, false, Some(e)),
Err(up_e) => (false, false, Some(up_e)),
}
} else {
// Fail fast, dropping the upstream half; the caller
// closes the connection since client_reuse=false.
(false, false, Some(e))
}
}
}
}
r = &mut upstream => {
match r {
Ok(upstream_can_reuse) => {
match (&mut downstream).await {
Ok(downstream_can_reuse) => (downstream_can_reuse, upstream_can_reuse, None),
Err(e) => (false, false, Some(e)),
}
}
Err(e) => (false, false, Some(e)),
}
}
};
// Release the borrows held by the two halves before restoring the
// custom message handles below.
drop(downstream);
drop(upstream);

if let Some(custom_session) = session.downstream_session.as_custom_mut() {
if let Some(downstream_custom_message_writer) = downstream_custom_message_writer {
Expand All @@ -155,12 +210,7 @@ where
}
}

match ret {
Ok((downstream_can_reuse, upstream_can_reuse)) => {
(downstream_can_reuse, upstream_can_reuse, None)
}
Err(e) => (false, false, Some(e)),
}
(downstream_can_reuse, upstream_can_reuse, error)
}

pub(crate) async fn proxy_to_h1_upstream(
Expand Down Expand Up @@ -275,6 +325,15 @@ where
// unread bytes and is unsafe to reuse.
return Ok(false);
}
if PipeState::is_downstream_filter_aborted(
pipe_state.load(Ordering::Acquire),
) {
// The downstream half aborted in a response filter
// (e.g. a 3xx redirect followed by the outer retry
// loop). Drain the remaining upstream response so the
// connection can be returned to the pool (issue #866).
return self.drain_upstream_response(client_session).await;
}
// The pipe closed but the downstream half did not signal completion:
// this is an unexpected closure, so surface the original send error.
return Err(result.expect_err("send failure already checked via is_err() above"));
Expand Down Expand Up @@ -326,6 +385,48 @@ where
Ok(upstream_can_reuse)
}

/// Maximum number of bytes drained from an upstream response after the
/// downstream half aborted in a response filter. Larger responses abandon
/// the connection instead of being reused, bounding the drain cost.
const MAX_RESPONSE_DRAIN_BYTES: usize = 1024 * 1024;

/// Drains the remaining upstream response body after the downstream half
/// aborted in a response filter, so the connection can be reused. Returns
/// true when the response was fully consumed within the drain bound.
async fn drain_upstream_response(&self, client_session: &mut HttpSessionV1) -> Result<bool> {
if client_session.was_upgraded() {
// Upgraded bodies have no framing to drain; never reuse.
return Ok(false);
}
let mut drained = 0usize;
loop {
match client_session.read_response_task().await {
Ok(task) => {
if task.is_end() {
return Ok(true);
}
let data_len = match &task {
HttpTask::Body(data, _) | HttpTask::UpgradedBody(data, _) => {
data.as_ref().map_or(0, |d| d.len())
}
_ => 0,
};
drained += data_len;
if drained > Self::MAX_RESPONSE_DRAIN_BYTES {
warn!(
"upstream response exceeds drain bound after filter abort; not reusing connection"
);
return Ok(false);
}
}
Err(e) => {
warn!("failed to drain upstream response after filter abort: {e}");
return Ok(false);
}
}
}
}

#[allow(clippy::too_many_arguments)]
async fn process_upstream_tasks(
&self,
Expand All @@ -336,6 +437,7 @@ where
serve_from_cache: &mut ServeFromCache,
range_body_filter: &mut proxy_cache::range_filter::RangeBodyFilter,
response_state: &mut ResponseStateMachine,
pipe_state: &AtomicU8,
) -> Result<Option<bool>>
where
SV: ProxyHttp + Send + Sync,
Expand Down Expand Up @@ -379,9 +481,20 @@ where
#[cfg(feature = "upstream_modules")]
session.upstream_modules_filter_task(&mut t).await?;
session.upstream_compression.response_filter(&mut t);
let task = self
let task = match self
.h1_response_filter(session, t, ctx, serve_from_cache, range_body_filter, false)
.await?;
.await
{
Ok(t) => t,
Err(e) => {
// A response_filter error (e.g. a 3xx redirect followed by
// the outer retry loop) aborts the downstream half. Signal
// the upstream half to drain the remaining response body so
// the connection can still be reused (issue #866).
pipe_state.store(PipeState::DownstreamFilterAborted as u8, Ordering::Release);
return Err(e);
}
};
if serve_from_cache.is_miss_header() {
response_state.enable_cached_response();
}
Expand Down Expand Up @@ -595,6 +708,7 @@ where
&mut serve_from_cache,
&mut range_body_filter,
&mut response_state,
&pipe_state,
).await? else {
// nothing sent downstream e.g. serve_from_cache
continue;
Expand All @@ -620,6 +734,7 @@ where
&mut serve_from_cache,
&mut range_body_filter,
&mut response_state,
&pipe_state,
).await? else {
// nothing sent downstream e.g. serve_from_cache
continue;
Expand Down
Loading
Loading