diff --git a/CHANGELOG.md b/CHANGELOG.md index ef88a6b9..deb36654 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,3 +5,8 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/), and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). ## [Unreleased] + +### Fixed +- `Context::map` no longer opens one multiplexer channel per item. Items are distributed + round-robin over at most `concurrency_limit` lanes, so the number of channels a `map` opens + is bounded by the limit rather than by the size of the workload. diff --git a/crates/common/src/context.rs b/crates/common/src/context.rs index fad36f15..d6baee52 100644 --- a/crates/common/src/context.rs +++ b/crates/common/src/context.rs @@ -8,7 +8,6 @@ use std::sync::Arc; use futures::{ AsyncRead, AsyncWrite, future::{self, BoxFuture, Either}, - stream::{self, StreamExt, TryStreamExt}, }; #[cfg(any(test, feature = "test-utils"))] @@ -23,8 +22,9 @@ pub use test::{ use crate::{ContextId, io::Io, mux::Mux, thread_pool::ThreadPool}; /// Default maximum number of [`map`](Context::map) items processed -/// concurrently. Both parties must agree on this value, so it is a fixed -/// constant rather than data- or timing-dependent. +/// concurrently, and with it the number of channels a `map` opens. Both parties +/// must agree on this value, so it is a fixed constant rather than data- or +/// timing-dependent. pub const DEFAULT_CONCURRENCY_LIMIT: usize = 32; /// A task execution context. @@ -157,13 +157,41 @@ impl Context { /// Applies `f` to each item concurrently, returning the results in input /// order. + /// + /// # Channel usage + /// + /// Every child context allocates a channel from the multiplexer, and + /// multiplexers impose a hard limit on how many channels they will track. + /// The number of items handed to this method is a function of the workload + /// (e.g. one item per circuit call), so giving each item its own channel + /// would make channel usage unbounded and blow past that limit on larger + /// workloads (tlsn's mux, for example, caps streams at 512 and tears the + /// connection down beyond it). + /// + /// Bounding only how many items run *at once* is not enough: a mux frees a + /// channel when its stream is dropped, but that release is processed by the + /// connection task and lags behind the rate at which a sliding window opens + /// new ones, so a per-item channel layout still exhausts the mux's budget + /// on a large enough workload. + /// + /// Items are therefore distributed round-robin over at most + /// `concurrency_limit` *lanes*, each of which owns a single child context + /// and processes its items sequentially. The number of channels ever opened + /// is `min(items.len(), concurrency_limit)`, independent of the workload + /// size, and that is also the concurrency bound. + /// + /// The lane assignment (`index % lanes`) and the order of items within a + /// lane depend only on the item index, so both parties derive an identical + /// channel layout and an identical per-channel message order. Both must + /// configure the same limit — see + /// [`SessionBuilder::concurrency_limit`](crate::SessionBuilder::concurrency_limit). pub async fn map(&mut self, items: Vec, f: F) -> Result, ContextError> where F: for<'a> Fn(&'a mut Context, T) -> BoxFuture<'a, R> + Clone + Send + 'static, T: Send + 'static, R: Send + 'static, { - let (mux, pool, concurrency_limit) = match &self.mode { + let (pool, concurrency_limit) = match &self.mode { Mode::Single => { let mut results = Vec::with_capacity(items.len()); for item in items { @@ -172,41 +200,55 @@ impl Context { return Ok(results); } Mode::Multi { - mux, pool, concurrency_limit, - } => (mux.clone(), pool.clone(), *concurrency_limit), + .. + } => (pool.clone(), *concurrency_limit), }; + let len = items.len(); + if len == 0 { + // Still consume a fork index so that both parties stay in sync. + let _ = self.next_fork(); + return Ok(Vec::new()); + } + let parent_id = self.next_fork(); - // Each item lazily opens its own channel only once `buffered` polls it, - // so at most `limit` channels are open at any time. Channel IDs stay - // keyed by item index and results are yielded in input order, so the - // bound changes neither the wire protocol nor the output ordering. - stream::iter(items.into_iter().enumerate()) - .map(move |(i, item)| { - let i = u32::try_from(i).expect("more than u32::MAX items"); - let id = parent_id.child(i); - let (mux, pool, f) = (mux.clone(), pool.clone(), f.clone()); - async move { - let io = mux.open(id.as_ref()).map_err(ContextError::mux)?; - let mut ctx = Context { - id, - io, - mode: Mode::Multi { - mux, - pool: pool.clone(), - concurrency_limit, - }, - fork_counter: 0, - }; - Ok(run(pool.as_ref(), async move { f(&mut ctx, item).await }).await) + let lanes = len.min(concurrency_limit); + let mut queues: Vec> = (0..lanes) + .map(|_| Vec::with_capacity(len.div_ceil(lanes))) + .collect(); + for (i, item) in items.into_iter().enumerate() { + queues[i % lanes].push((i, item)); + } + + let mut tasks = Vec::with_capacity(lanes); + for (lane, queue) in queues.into_iter().enumerate() { + let lane = u32::try_from(lane).expect("lane count fits in u32"); + let mut ctx = self.child(parent_id.child(lane))?; + let f = f.clone(); + tasks.push(run(pool.as_ref(), async move { + let mut results = Vec::with_capacity(queue.len()); + for (i, item) in queue { + results.push((i, f(&mut ctx, item).await)); } - }) - .buffered(concurrency_limit) - .try_collect() - .await + results + })); + } + + // Restore input order. + let mut results: Vec> = (0..len).map(|_| None).collect(); + for lane_results in future::join_all(tasks).await { + for (i, result) in lane_results { + results[i] = Some(result); + } + } + + Ok(results + .into_iter() + .map(|result| result.expect("every item is assigned to exactly one lane")) + .collect()) } /// Runs `a` and `b` concurrently and returns both results. diff --git a/crates/common/src/context/test/tests.rs b/crates/common/src/context/test/tests.rs index b77afbf8..8edb562c 100644 --- a/crates/common/src/context/test/tests.rs +++ b/crates/common/src/context/test/tests.rs @@ -2,7 +2,7 @@ use serio::{SinkExt, stream::IoStreamExt}; -use crate::context::Context; +use crate::context::{Context, DEFAULT_CONCURRENCY_LIMIT}; use super::{ recording::{recording_mt_context, recording_st_context}, @@ -497,3 +497,48 @@ async fn test_map_respects_concurrency_limit() { // limit, so the window is fully utilized). assert_eq!(max_seen.load(Ordering::SeqCst), LIMIT); } + +/// Two parties mapping far more items than the concurrency limit must stay in +/// lockstep: both sides derive the same lanes, and results come back in input +/// order regardless of which lane handled which item. +#[tokio::test] +async fn test_recording_mt_map_beyond_concurrency_limit() { + let (session_0, session_1, recorded) = recording_mt_context(1024 * 1024); + + let mut ctx_0 = session_0.new_context().unwrap(); + let mut ctx_1 = session_1.new_context().unwrap(); + + // Far more items than the concurrency limit. + let items: Vec = (0..(DEFAULT_CONCURRENCY_LIMIT as u32 * 3 + 1)).collect(); + + let (recv_results, send_results) = futures::join!( + ctx_0.map(items.clone(), |ctx: &mut Context, _item: u32| Box::pin( + async move { + let msg: u32 = ctx.io_mut().expect_next().await.unwrap(); + msg + } + )), + ctx_1.map(items.clone(), |ctx: &mut Context, item: u32| Box::pin( + async move { + ctx.io_mut().send(item * 10).await.unwrap(); + } + )) + ); + + let recv_results = recv_results.unwrap(); + send_results.unwrap(); + + assert_eq!( + recv_results, + items.iter().map(|item| item * 10).collect::>() + ); + + // Channel usage is bounded by the lane count rather than growing with the + // number of items, which would exhaust the multiplexer's channel budget. + let recorded_data = recorded.lock().unwrap(); + assert_eq!( + recorded_data.channels.len(), + DEFAULT_CONCURRENCY_LIMIT, + "map must not open more than one channel per lane" + ); +}