Skip to content

Add support for multi-subrequest forking #7

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
addykan wants to merge 1 commit into
base: richard/lossy-subrequest
Choose a base branch
from
Open

Add support for multi-subrequest forking #7

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
171 changes: 138 additions & 33 deletions pingora-core/src/protocols/http/body_fork.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -40,6 +40,13 @@ pub enum BodyForkPushError {
Rejected,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub enum BodyMultiForkPushError {
Rejected,

AllErrored(Vec<BodyForkPushError>),
}

/// Error from [`BodyForkReceiver::recv`].
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct BodyForkAborted;
Expand All @@ -61,23 +68,24 @@ struct BodyForkShared {
/// Dropping without calling `finish` aborts the fork (clears buffered data).
pub struct BodyForkSender {
inner: Arc<BodyForkShared>,
mapper: Box<dyn Fn(Bytes) -> Option<Bytes> + Send + Sync>,
}

/// Receive side: [`recv`](BodyForkReceiver::recv) drains all queued chunks.
pub struct BodyForkReceiver {
inner: Arc<BodyForkShared>,
}

pub struct BodyMultiForkSender {
mapper: Box<dyn Fn(Bytes) -> Option<Bytes> + Send + Sync>,
senders: Mutex<Vec<BodyForkSender>>,
}

/// Create a bounded body fork pair with an owned-chunk mapper.
///
/// The mapper runs before queue admission. Returning [`None`] rejects the chunk with
/// [`BodyForkPushError::Rejected`]. A successfully mapped chunk that cannot be queued is dropped
/// before [`BodyForkPushError::Full`] is returned.
pub fn body_fork_pair_with<F>(max_chunks: usize, mapper: F) -> (BodyForkSender, BodyForkReceiver)
where
F: Fn(Bytes) -> Option<Bytes> + Send + Sync + 'static,
{
pub fn body_fork_pair_with(max_chunks: usize) -> (BodyForkSender, BodyForkReceiver) {
let inner = Arc::new(BodyForkShared {
max_chunks,
state: Mutex::new(BodyForkState::Open(VecDeque::new())),
Expand All @@ -86,20 +94,74 @@ where
(
BodyForkSender {
inner: inner.clone(),
mapper: Box::new(mapper),
},
BodyForkReceiver { inner },
)
}

impl BodyForkSender {
/// Try to map and queue a body chunk.
pub fn try_push(&self, chunk: Bytes) -> Result<(), BodyForkPushError> {
let chunk = (self.mapper)(chunk).ok_or(BodyForkPushError::Rejected)?;
pub fn body_multi_fork_pair_with<F>(
max_chunks: usize,
forks: usize,
mapper: F,
) -> (BodyMultiForkSender, Vec<BodyForkReceiver>)
where
F: Fn(Bytes) -> Option<Bytes> + Send + Sync + 'static,
{
assert!(forks > 0, "at least one fork is required");
let (senders, receivers) = (0..forks).map(|_| body_fork_pair_with(max_chunks)).unzip();
(
BodyMultiForkSender {
mapper: Box::new(mapper),
senders: Mutex::new(senders),
},
receivers,
)
}

impl BodyMultiForkSender {
/// Tries to push the same chunk of bytes to all forks, calling the mapping function once before pushing.
pub fn try_push(&self, chunk: Bytes) -> Result<Vec<BodyForkPushError>, BodyMultiForkPushError> {
let mut senders = self.senders.lock();
let chunk = (self.mapper)(chunk).ok_or(BodyMultiForkPushError::Rejected)?;
if chunk.is_empty() {
return Ok(());
return Ok(Vec::new());
}

let mut errs = Vec::new();
let remaining_senders = senders
.drain(..)
.filter_map(|sender| {
let res = sender.try_push(chunk.clone());
match res {
Ok(_) => Some(sender),
Err(e) => {
errs.push(e);
None
}
}
})
.collect::<Vec<_>>();
*senders = remaining_senders;

if senders.is_empty() {
Err(BodyMultiForkPushError::AllErrored(errs))
} else {
Ok(errs)
}
}

/// Finishes all forks.
pub fn finish(self) {
let mut senders = self.senders.lock();
for sender in senders.drain(..) {
sender.finish();
}
}
}

impl BodyForkSender {
/// Try to map and queue a body chunk.
pub fn try_push(&self, chunk: Bytes) -> Result<(), BodyForkPushError> {
let mut g = self.inner.state.lock();
let pending = match &mut *g {
BodyForkState::Open(pending) => pending,
Expand Down Expand Up @@ -210,7 +272,7 @@ mod tests {

#[tokio::test]
async fn push_finish_recv() {
let (tx, mut rx) = body_fork_pair_with(32, Some);
let (tx, mut rx) = body_fork_pair_with(32);
tx.try_push(Bytes::from_static(b"a")).unwrap();
tx.try_push(Bytes::from_static(b"bc")).unwrap();
tx.finish();
Expand All @@ -220,7 +282,7 @@ mod tests {

#[tokio::test]
async fn max_chunks_rejects_push() {
let (tx, mut rx) = body_fork_pair_with(2, Some);
let (tx, mut rx) = body_fork_pair_with(2);
tx.try_push(Bytes::from_static(b"ab")).unwrap();
tx.try_push(Bytes::from_static(b"cd")).unwrap();
assert_eq!(
Expand All @@ -234,7 +296,7 @@ mod tests {

#[tokio::test]
async fn recv_frees_capacity() {
let (tx, mut rx) = body_fork_pair_with(2, Some);
let (tx, mut rx) = body_fork_pair_with(2);
tx.try_push(Bytes::from_static(b"a")).unwrap();
tx.try_push(Bytes::from_static(b"b")).unwrap();
assert_eq!(
Expand All @@ -253,22 +315,22 @@ mod tests {

#[tokio::test]
async fn drop_without_finish_clears_bytes() {
let (tx, mut rx) = body_fork_pair_with(32, Some);
let (tx, mut rx) = body_fork_pair_with(32);
tx.try_push(Bytes::from_static(b"x")).unwrap();
drop(tx);
assert_eq!(rx.recv().await, Err(BodyForkAborted));
}

#[tokio::test]
async fn finish_empty_body() {
let (tx, mut rx) = body_fork_pair_with(32, Some);
let (tx, mut rx) = body_fork_pair_with(32);
tx.finish();
assert_eq!(rx.recv().await, Ok(None));
}

#[tokio::test]
async fn recv_returns_available_without_waiting_for_end() {
let (tx, mut rx) = body_fork_pair_with(32, Some);
let (tx, mut rx) = body_fork_pair_with(32);
tx.try_push(Bytes::from_static(b"a")).unwrap();
tx.try_push(Bytes::from_static(b"b")).unwrap();
let chunks = rx.recv().await.unwrap().unwrap();
Expand All @@ -280,6 +342,49 @@ mod tests {
assert!(rx.recv().await.unwrap().is_none());
}

#[tokio::test]
async fn multi_fork_partial_failure_keeps_healthy_fork() {
let (tx, mut receivers) = body_multi_fork_pair_with(1, 2, Some);
let mut healthy = receivers.remove(0);
let mut stalled = receivers.remove(0);

assert!(tx.try_push(Bytes::from_static(b"a")).unwrap().is_empty());
assert_eq!(flatten(healthy.recv().await.unwrap().unwrap()), b"a");

assert_eq!(
tx.try_push(Bytes::from_static(b"b")),
Ok(vec![BodyForkPushError::Full])
);
assert_eq!(stalled.recv().await, Err(BodyForkAborted));
assert_eq!(flatten(healthy.recv().await.unwrap().unwrap()), b"b");

tx.finish();
assert_eq!(healthy.recv().await, Ok(None));
}

#[tokio::test]
async fn multi_fork_errors_when_last_healthy_fork_fails() {
let (tx, mut receivers) = body_multi_fork_pair_with(1, 2, Some);
let mut healthy = receivers.remove(0);
let mut stalled = receivers.remove(0);

assert!(tx.try_push(Bytes::from_static(b"a")).unwrap().is_empty());
assert_eq!(flatten(healthy.recv().await.unwrap().unwrap()), b"a");
assert_eq!(
tx.try_push(Bytes::from_static(b"b")),
Ok(vec![BodyForkPushError::Full])
);

assert_eq!(
tx.try_push(Bytes::from_static(b"c")),
Err(BodyMultiForkPushError::AllErrored(vec![
BodyForkPushError::Full
]))
);
assert_eq!(healthy.recv().await, Err(BodyForkAborted));
assert_eq!(stalled.recv().await, Err(BodyForkAborted));
}

struct TrackedBytes {
bytes: Bytes,
live_bytes: Arc<AtomicUsize>,
Expand Down Expand Up @@ -312,48 +417,48 @@ mod tests {

#[tokio::test]
async fn mapper_rejection_aborts_when_sender_is_dropped() {
let (tx, mut rx) = body_fork_pair_with(32, |_| None);
let (tx, mut receivers) = body_multi_fork_pair_with(32, 1, |_| None);
let mut rx = receivers.pop().unwrap();
assert_eq!(
tx.try_push(Bytes::from_static(b"rejected")),
Err(BodyForkPushError::Rejected)
Err(BodyMultiForkPushError::Rejected)
);
drop(tx);
assert_eq!(rx.recv().await, Err(BodyForkAborted));
}

#[tokio::test]
async fn full_drops_mapped_chunk_immediately() {
async fn full_drops_failed_fork_bytes_immediately() {
let live_bytes = Arc::new(AtomicUsize::new(0));
let (tx, mut rx) = body_fork_pair_with(1, tracked_mapper(live_bytes.clone()));
let (tx, mut rx_vec) = body_multi_fork_pair_with(1, 1, tracked_mapper(live_bytes.clone()));
let mut rx = rx_vec.pop().expect("expected one fork");

tx.try_push(Bytes::from_static(b"a")).unwrap();
assert_eq!(live_bytes.load(Ordering::SeqCst), 1);
assert_eq!(
tx.try_push(Bytes::from_static(b"bc")),
Err(BodyForkPushError::Full)
Err(BodyMultiForkPushError::AllErrored(vec![
BodyForkPushError::Full
]))
);
assert_eq!(
live_bytes.load(Ordering::SeqCst),
1,
"the mapped chunk rejected by queue admission must be dropped"
0,
"aborting the failed fork must drop its queued and rejected chunks"
);

tx.finish();
let Some(chunks) = rx.recv().await.unwrap() else {
panic!("expected queued chunk");
};
drop(chunks);
assert_eq!(live_bytes.load(Ordering::SeqCst), 0);
assert_eq!(rx.recv().await, Ok(None));
assert_eq!(rx.recv().await, Err(BodyForkAborted));
}

#[tokio::test]
async fn mapped_bytes_live_across_drained_batch_and_refilled_queue() {
const CHUNKS_PER_BATCH: usize = 32;

let live_bytes = Arc::new(AtomicUsize::new(0));
let (tx, mut rx) =
body_fork_pair_with(CHUNKS_PER_BATCH, tracked_mapper(live_bytes.clone()));
let (tx, mut rx_vec) =
body_multi_fork_pair_with(CHUNKS_PER_BATCH, 1, tracked_mapper(live_bytes.clone()));
let mut rx = rx_vec.pop().expect("expected one fork");

for _ in 0..CHUNKS_PER_BATCH {
tx.try_push(Bytes::from_static(b"x")).unwrap();
Expand Down Expand Up @@ -383,7 +488,7 @@ mod tests {

#[tokio::test]
async fn abort_wait_is_persistent() {
let (tx, mut rx) = body_fork_pair_with(1, Some);
let (tx, mut rx) = body_fork_pair_with(1);
tx.try_push(Bytes::from_static(b"x")).unwrap();
let Some(batch) = rx.recv().await.unwrap() else {
panic!("expected queued chunk");
Expand Down
9 changes: 5 additions & 4 deletions pingora-core/src/protocols/http/server.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -593,17 +593,18 @@ impl Session {
///
/// Returns [`None`] for subrequest/custom sessions or if a fork is already attached on the
/// underlying session. See [`SessionV1::attach_request_body_fork_with`].
pub fn attach_request_body_fork_with<F>(
pub fn attach_request_body_multi_fork_with<F>(
&mut self,
max_chunks: usize,
forks: usize,
mapper: F,
) -> Option<BodyForkReceiver>
) -> Option<Vec<BodyForkReceiver>>
where
F: Fn(Bytes) -> Option<Bytes> + Send + Sync + 'static,
{
match self {
Self::H1(s) => s.attach_request_body_fork_with(max_chunks, mapper),
Self::H2(s) => s.attach_request_body_fork_with(max_chunks, mapper),
Self::H1(s) => s.attach_request_body_multi_fork_with(max_chunks, forks, mapper),
Self::H2(s) => s.attach_request_body_multi_fork_with(max_chunks, forks, mapper),
Self::Subrequest(_) | Self::Custom(_) => None,
}
}
Expand Down
19 changes: 11 additions & 8 deletions pingora-core/src/protocols/http/v1/server.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -32,10 +32,9 @@ use tokio::io::{AsyncReadExt, AsyncWriteExt};

use super::body::{BodyReader, BodyWriter};
use super::common::*;
use crate::protocols::http::body_fork::{body_multi_fork_pair_with, BodyMultiForkSender};
use crate::protocols::http::{
body_buffer::FixedBuffer,
body_fork::{body_fork_pair_with, BodyForkReceiver, BodyForkSender},
date, HttpTask,
body_buffer::FixedBuffer, body_fork::BodyForkReceiver, date, HttpTask,
};
use crate::protocols::{Digest, SocketAddr, Stream};
use crate::utils::{BufRef, KVRef};
Expand Down Expand Up @@ -91,7 +90,7 @@ pub struct HttpSession {
/// after this session ends
keepalive_reuses_remaining: Option<u32>,
/// Optional lossy tee of request body bytes (see [`Self::attach_request_body_fork`]).
body_fork: Option<BodyForkSender>,
body_fork: Option<BodyMultiForkSender>,
}

impl HttpSession {
Expand Down Expand Up @@ -475,18 +474,19 @@ impl HttpSession {
///
/// The mapper runs for every forked chunk before queue admission. Returning [`None`] aborts
/// only the fork; the primary request continues with its original chunk.
pub fn attach_request_body_fork_with<F>(
pub fn attach_request_body_multi_fork_with<F>(
&mut self,
max_chunks: usize,
forks: usize,
mapper: F,

@molocule molocule Aug 11, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit, but could mapper be named to reflect that it just reserves memory

@addykan addykan Aug 11, 2026

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

the mapper is abstract inside this module - the concept of memory reservations only exists in fprs. I guess it's fine to rename given it's only us using it, but imo that bleeds the abstraction across two separate crates?

) -> Option<BodyForkReceiver>
) -> Option<Vec<BodyForkReceiver>>
where
F: Fn(Bytes) -> Option<Bytes> + Send + Sync + 'static,
{
if self.body_fork.is_some() {
return None;
}
let (tx, rx) = body_fork_pair_with(max_chunks, mapper);
let (tx, rx) = body_multi_fork_pair_with(max_chunks, forks, mapper);
self.body_fork = Some(tx);
Some(rx)
}
Expand Down Expand Up @@ -1488,7 +1488,10 @@ mod tests_stream {
let mock_io = Builder::new().read(&input[..]).build();
let mut http_stream = HttpSession::new(Box::new(mock_io));
http_stream.read_request().await.unwrap();
let mut fork = http_stream.attach_request_body_fork_with(1, Some).unwrap();
let mut forks = http_stream
.attach_request_body_multi_fork_with(1, 1, Some)
.unwrap();
let mut fork = forks.pop().expect("expected one fork");

let body = http_stream.read_body_bytes().await.unwrap().unwrap();
assert_eq!(body, b"abc".as_slice());
Expand Down
Loading
Loading