Subrequest response finish can lose the terminal Done task if cancelled under channel backpressure

[YZL0v3ZZ]

Describe the bug

Pingora's legacy subrequest response writer can record a response body as complete before the terminal HttpTask::Done has actually been delivered to the subrequest output channel.

The vulnerable state transition is:

subrequest response uses a bounded output channel
-> response header/body tasks fill the channel
-> legacy finish() marks BodyWriter state as Complete(...)
-> finish() awaits tx.send(HttpTask::Done)
-> caller cancels/drops the finish future while the channel send is pending
-> Tokio drops the undelivered Done value
-> retrying finish() sees BodyMode::Complete(_) and becomes a no-op
-> the receiver can observe body bytes without the protocol-level terminal Done task

This is a cancellation-correctness issue because the durable in-memory state says the response stream is complete, while the external completion event that the subrequest pipe consumes was not published and is no longer recoverable by retrying finish().

Root cause

The root cause is visible in the legacy subrequest finish path. Source links below are pinned to commit e6e677fe9b58555140ab7bd14feff035392b3530.

1. Subrequest sessions use bounded channels for request and response tasks. A full output channel is a normal backpressure condition:

pingora-core/src/protocols/http/subrequest/server.rs#L155-L161

pub fn new_from_session(session: &GenericHttpSession) -> (Self, SubrequestHandle) {
    let v1_inner = SessionV1::new(Box::new(DummyIO::new(&session.to_h1_raw())));
    let digest = session.digest().cloned();
    // allow buffering a small number of tasks, otherwise exert backpressure
    const CHANNEL_BUFFER_SIZE: usize = 4;
    let (downstream_tx, downstream_rx) = mpsc::channel(CHANNEL_BUFFER_SIZE);
    let (upstream_tx, upstream_rx) = mpsc::channel(CHANNEL_BUFFER_SIZE);

2. The legacy subrequest HttpSession::finish() delegates to BodyWriter::finish() and awaits it as a normal caller-owned future:

pingora-core/src/protocols/http/subrequest/server.rs#L486-L494

/// Signal that there is no more body to write.
/// This call will try to flush the buffer if there is any un-flushed data.
/// For chunked encoding response, this call will also send the last chunk.
/// For upgraded sessions, this call will also close the reading of the client body.
pub async fn finish(&mut self) -> Result<Option<usize>> {
    let res = self
        .body_writer
        .finish(self.tx.as_mut().expect("tx valid before shutdown"))
        .await?;

3. BodyWriter::finish() treats BodyMode::Complete(_) as already finished and returns Ok(None). This is important because cancellation can leave this state behind before the terminal task is delivered:

pingora-core/src/protocols/http/subrequest/body.rs#L309-L315

pub async fn finish(&mut self, sender: &mut mpsc::Sender<HttpTask>) -> Result<Option<usize>> {
    match self.body_mode {
        BM::Complete(_) => Ok(None),
        BM::ContentLength(_, _) => self.do_finish_body(sender).await,
        BM::UntilClose(_) => self.do_finish_until_close_body(sender).await,
        BM::ToSelect => Ok(None),
    }
}

4. In the content-length finish path, the writer mutates self.body_mode before the await that sends HttpTask::Done:

pingora-core/src/protocols/http/subrequest/body.rs#L318-L332

async fn do_finish_body(&mut self, tx: &mut mpsc::Sender<HttpTask>) -> Result<Option<usize>> {
    match self.body_mode {
        BM::ContentLength(total, written) => {
            self.body_mode = BM::Complete(written);
            if written < total {
                return Error::e_explain(
                    PREMATURE_BODY_END,
                    format!("Content-length: {total} bytes written: {written} (subrequest)"),
                );
            }
            tx.send(HttpTask::Done).await.or_err(
                WriteError,
                "while sending done task to downstream (subrequest)",
            )?;
            Ok(Some(written))

5. The close-delimited finish path has the same ordering:

pingora-core/src/protocols/http/subrequest/body.rs#L338-L349

async fn do_finish_until_close_body(
    &mut self,
    tx: &mut mpsc::Sender<HttpTask>,
) -> Result<Option<usize>> {
    match self.body_mode {
        BM::UntilClose(written) => {
            self.body_mode = BM::Complete(written);
            tx.send(HttpTask::Done).await.or_err(
                WriteError,
                "while sending done task to downstream (subrequest)",
            )?;
            Ok(Some(written))

If tx.send(HttpTask::Done).await is pending and the future is dropped, Tokio does not deliver that owned Done value. The surviving state is still BodyMode::Complete(written), so a later finish() call cannot reconstruct the missing terminal task.

6. The terminal task is meaningful to the subrequest pipe. The pipe receives tasks from the subrequest output channel and forwards them to the parent session:

pingora-proxy/src/subrequest/pipe.rs#L244-L274

task = state.pipe_rx.as_mut().expect("pipe_rx always set after spawn").recv(), if !response_state.upstream_done() => {
    if let Some(task) = task {
        let mut tasks = vec![task];
        while let Some(maybe_task) = tokio::task::unconstrained(state.pipe_rx.as_mut().expect("pipe_rx always set after spawn").recv()).now_or_never() {
            match maybe_task {
                Some(task) => {
                    tasks.push(task);
                }
                None => {
                    debug!("upstream channel closed early");
                    response_state.maybe_set_upstream_done(true);
                    break;
                }
            }
        }
        if !tasks.is_empty() {
            let response_done = map_pipe_err(session.write_response_tasks(tasks).await, false, &mut state)?;

The bug therefore creates a protocol-level mismatch:

writer state: BodyMode::Complete(written)
receiver state: no HttpTask::Done received
retry behavior: finish() returns Ok(None), no Done sent

7. The same file contains a safer local pattern in the cancel-safe proxy task API: reserve channel capacity first, then synchronously update state and send Done through the acquired permit:

pingora-core/src/protocols/http/subrequest/server.rs#L1095-L1125

/// Dispatch the final `HttpTask::Done`, mirroring `body_writer::finish`.
async fn dispatch_finish(&mut self) -> Result<()> {
    // Reserve cancel-safely, then synchronously update body_mode and send.
    let tx_ref = self
        .tx
        .as_ref()
        .ok_or_else(|| Error::explain(InternalError, "subrequest tx already shut down"))?;
    let permit = match self.write_timeout {
        Some(t) => match timeout(t, tx_ref.reserve()).await {
            Ok(res) => res.or_err(WriteError, "subrequest channel closed")?,
            Err(_) => {
                return Error::e_explain(
                    WriteTimedout,
                    format!("reserving subrequest channel slot for finish, timeout: {t:?}"),
                );
            }
        },
        None => tx_ref
            .reserve()
            .await
            .or_err(WriteError, "subrequest channel closed")?,
    };

    match self.body_writer.body_mode {
        BodyMode::ContentLength(_total, written) => {
            self.body_writer.body_mode = BodyMode::Complete(written);
            permit.send(HttpTask::Done);
        }
        BodyMode::UntilClose(written) => {
            self.body_writer.body_mode = BodyMode::Complete(written);
            permit.send(HttpTask::Done);

That ordering avoids this class of bug: if reserve().await is cancelled, no completion state has been committed yet. In the legacy path, the completion state is committed before the cancellable send.

Pingora info

Please include the following information about your environment:

Pingora version: audited and reproduced against commit e6e677fe9b58555140ab7bd14feff035392b3530
Rust version: not captured in the shared reproduction log
Operating system version: reproduced as a Rust unit test on Linux x86 (server36)

Steps to reproduce

The following is a focused whitebox reproduction. It intentionally asserts the current bad behavior, so the test passing means the cancellation window is reachable in the current implementation. After a fix, the core assertion should be inverted into a regression property such as: after a cancelled finish, retrying finish() must still be able to deliver exactly one terminal HttpTask::Done, or the writer state must remain resumable until the terminal task is committed.

Steps:

1. Apply the test-only code below inside the existing #[cfg(test)] mod tests_stream in pingora-core/src/protocols/http/subrequest/server.rs.
2. Run the targeted test command shown below.
3. Observe that cancelling a blocked legacy finish() leaves BodyMode::Complete(3) while no HttpTask::Done is delivered, and a later finish() retry is a no-op.

Full whitebox test code for pingora-core/src/protocols/http/subrequest/server.rs

#[tokio::test]
async fn cancelled_legacy_finish_marks_complete_without_done() {
    let (mut http_stream, mut handle) = build_req().await;
    let response = test_header(StatusCode::OK);

    http_stream
        .write_response_header_ref(&response)
        .await
        .expect("test async operation should succeed");

    for body in ["a", "b", "c"] {
        http_stream
            .write_body(Bytes::from(body))
            .await
            .expect("test async operation should succeed")
            .expect("body task should be queued");
    }

    assert_eq!(http_stream.body_writer.body_mode, BodyMode::UntilClose(3));

    let mut finish = Box::pin(http_stream.finish());
    let waker = futures::task::noop_waker();
    let mut cx = std::task::Context::from_waker(&waker);
    assert!(
        matches!(
            std::future::Future::poll(finish.as_mut(), &mut cx),
            std::task::Poll::Pending
        ),
        "legacy finish should block on the full subrequest channel"
    );
    drop(finish);

    assert_eq!(http_stream.body_writer.body_mode, BodyMode::Complete(3));

    let first = recv_task(&mut handle.rx);
    assert!(matches!(first, HttpTask::Header(..)));

    let retried_finish = http_stream
        .finish()
        .await
        .expect("test async operation should succeed");
    assert_eq!(
        retried_finish, None,
        "retrying finish is a no-op after cancellation marked the body complete"
    );

    let mut delivered_after_retry = Vec::new();
    while let Ok(task) = handle.rx.try_recv() {
        delivered_after_retry.push(task);
    }
    assert_eq!(delivered_after_retry.len(), 3);
    assert!(
        delivered_after_retry
            .iter()
            .all(|task| matches!(task, HttpTask::Body(..))),
        "only the previously queued body tasks remain; the cancelled Done task is not repaired"
    );
    assert_eq!(
        delivered_after_retry
            .iter()
            .filter(|task| matches!(task, HttpTask::Done))
            .count(),
        0,
        "the terminal Done task was lost when the blocked legacy finish future was cancelled"
    );
}

Targeted command:

cargo test -p pingora-core --lib cancelled_legacy_finish_marks_complete_without_done -- --nocapture

Observed result:

Finished `test` profile [unoptimized + debuginfo] target(s) in 45.23s
     Running unittests src/lib.rs (target/debug/deps/pingora_core-4ef134c88541a33d)

running 1 test
test protocols::http::subrequest::server::tests_stream::cancelled_legacy_finish_marks_complete_without_done ... ok

test result: ok. 1 passed; 0 failed; 0 ignored; 0 measured; 484 filtered out; finished in 0.00s

Why this reproduces the issue:

• The test creates a real subrequest HttpSession and SubrequestHandle.
• A 200 OK response without Content-Length puts the legacy body writer into BodyMode::UntilClose(0).
• The test writes three body tasks. Together with the response header, this fills the four-slot subrequest output channel.
• It then manually polls http_stream.finish() once with a noop waker. This deterministically reaches the exact cancellation window: BodyWriter::finish() has set BodyMode::Complete(3), but tx.send(HttpTask::Done).await is pending because the channel is full.
• Dropping the pinned finish future models a real cancellation source: losing a select!, timeout, task abort, parent future drop, or runtime shutdown.
• The test proves the bad surviving state by asserting BodyMode::Complete(3).
• It drains only the already-sent header to create channel capacity, then retries finish().
• The retry returns None because the writer believes it is already complete.
• Finally, the receiver still contains only the three previously queued body tasks. There is no HttpTask::Done, proving that the cancelled terminal send was not repaired.

This is a bug-existence validation test. It asserts current behavior rather than the desired post-fix behavior.

Expected results

Cancelling a subrequest response finish() should not make the terminal response event unrecoverable.

More specifically, once a subrequest writer records its response body as complete, one of the following should be true:

• the terminal HttpTask::Done has already been delivered to the output channel;
• the terminal task is still owned by durable/resumable writer state;
• or retrying finish() can still deliver the missing terminal task.

The safe state-machine ordering should be:

await channel capacity
-> if reservation succeeds, synchronously mark body complete and send Done
-> if reservation is cancelled, keep body state unfinished so retry remains valid

This is the same ordering already used by the cancel-safe proxy task path in dispatch_finish().

Observed results

The current legacy finish path can produce this state:

BodyWriter.body_mode == BodyMode::Complete(3)
SubrequestHandle.rx received Header + Body + Body + Body
SubrequestHandle.rx did not receive HttpTask::Done
retrying HttpSession::finish() returns Ok(None)
no later finish call can repair the missing Done task

Operationally, this can leave a subrequest pipe waiting for a terminal response task that will never be sent while the sender remains alive. If the sender is later dropped, the pipe may only observe channel close, which is weaker than an explicit protocol-level HttpTask::Done and loses the distinction between "the response stream completed" and "the finishing future disappeared while the terminal task was pending."

Additional context

The cancellation-safety invariant I expected is:

A subrequest response writer must not persist "body complete" unless the
corresponding terminal response task has either been delivered or remains
resumable after cancellation.

A robust fix would likely reuse the reserve-before-commit pattern that already exists in dispatch_finish():

• In the legacy BodyWriter::finish() path, await tx.reserve() before mutating body_mode.
• After a permit is acquired, synchronously set BodyMode::Complete(written) and call permit.send(HttpTask::Done).
• If reserving capacity is cancelled or times out, leave body_mode unchanged so a later finish() can retry.

Other possible fixes would be to add an explicit FinishInProgress state that owns the pending terminal task, or route legacy subrequest finalization through the existing cancel-safe proxy task machinery. The important property is that the completion marker and terminal task delivery must commit together; a cancellable await should not sit between them.