From 4daac1d320dffd87d0bd651bbd008da765954104 Mon Sep 17 00:00:00 2001 From: MorganaFuture Date: Thu, 14 May 2026 00:09:45 +0300 Subject: [PATCH 1/2] Bound bi-stream handler concurrency at accept site A peer could open unbounded bi-streams and exhaust memory before the inner `agent.limits().sync` check inside `serve_sync` ever fired. Move the try_acquire up to `spawn_bipayload_handler`, before `tokio::spawn`, and write `MaxConcurrencyReached` on the stream when no permits remain. The duplicate check inside `serve_sync` is removed: keeping both would consume two permits per sync from the same 3-permit pool. --- crates/corro-agent/src/agent/bi.rs | 58 ++++++++++++++++++++++++-- crates/corro-agent/src/agent/tests.rs | 43 +++++++++++++++++++ crates/corro-agent/src/api/peer/mod.rs | 21 +--------- 3 files changed, 99 insertions(+), 23 deletions(-) diff --git a/crates/corro-agent/src/agent/bi.rs b/crates/corro-agent/src/agent/bi.rs index c6f0bb1d8..ba75e70d7 100644 --- a/crates/corro-agent/src/agent/bi.rs +++ b/crates/corro-agent/src/agent/bi.rs @@ -1,7 +1,9 @@ -use crate::api::peer::serve_sync; +use crate::api::peer::{encode_write_sync_msg, serve_sync}; +use bytes::BytesMut; use corro_types::{ agent::{Agent, Bookie}, broadcast::{BiPayload, BiPayloadV1}, + sync::{SyncMessage, SyncMessageV1, SyncRejectionV1}, }; use metrics::counter; use speedy::Readable; @@ -9,7 +11,7 @@ use std::time::Duration; use tokio::time::timeout; use tokio_stream::StreamExt; use tokio_util::codec::{FramedRead, LengthDelimitedCodec}; -use tracing::{debug, error, trace, warn}; +use tracing::{debug, error, info_span, trace, warn, Instrument}; use tripwire::Tripwire; /// Spawn a task that listens for incoming bi-directional sync streams @@ -51,11 +53,61 @@ pub fn spawn_bipayload_handler( conn.remote_address() ); - // TODO: implement concurrency limit for sync requests + // Bound the number of in-flight bi-handler tasks. All bi-streams + // are sync requests today, so we reuse `Limits.sync` rather than + // adding a parallel counter. Acquiring before `tokio::spawn` + // prevents an unbounded peer from exhausting tasks/memory just + // by opening streams. + let permit = match agent.limits().sync.clone().try_acquire_owned() { + Ok(p) => p, + Err(_) => { + counter!("corro.peer.stream.accept.rejected.total", "type" => "bi") + .increment(1); + warn!( + "rejecting bi-stream from {}: sync concurrency limit reached", + conn.remote_address() + ); + + // Best-effort: tell the peer why we rejected, then drop + // both stream halves. A slow peer must not wedge the + // accept loop, hence the timeout. + let mut tx = tx; + let _ = timeout(Duration::from_secs(5), async { + let mut codec = LengthDelimitedCodec::builder() + .max_frame_length(100 * 1_024 * 1_024) + .new_codec(); + let mut send_buf = BytesMut::new(); + let mut encode_buf = BytesMut::new(); + if let Err(e) = encode_write_sync_msg( + &mut codec, + &mut encode_buf, + &mut send_buf, + SyncMessage::V1(SyncMessageV1::Rejection( + SyncRejectionV1::MaxConcurrencyReached, + )), + &mut tx, + ) + .instrument(info_span!("write_bi_rejection")) + .await + { + debug!("could not write bi-stream rejection: {e}"); + } + let _ = tx.finish(); + }) + .await; + + // `rx` is dropped at scope end, closing the receive side. + continue; + } + }; + tokio::spawn({ let agent = agent.clone(); let bookie = bookie.clone(); async move { + // Hold the permit for the task's lifetime; drops on exit + // (including panic), releasing capacity for the next peer. + let _permit = permit; let mut framed = FramedRead::new( rx, LengthDelimitedCodec::builder() diff --git a/crates/corro-agent/src/agent/tests.rs b/crates/corro-agent/src/agent/tests.rs index f1b7659b9..7c26304f3 100644 --- a/crates/corro-agent/src/agent/tests.rs +++ b/crates/corro-agent/src/agent/tests.rs @@ -792,6 +792,49 @@ struct TestRecord { text: String, } +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn test_bi_stream_rejects_when_sync_at_capacity() -> eyre::Result<()> { + _ = tracing_subscriber::fmt::try_init(); + let (tripwire, tripwire_worker, tripwire_tx) = Tripwire::new_simple(); + + let ta1 = launch_test_agent(|conf| conf.build(), tripwire.clone()).await?; + let ta2 = launch_test_agent(|conf| conf.build(), tripwire.clone()).await?; + + // Pre-acquire all 3 sync permits on ta1 so the next inbound bi-stream + // must be rejected by `spawn_bipayload_handler`'s new acquire-or-reject + // path. Permits are released when the test ends (RAII). + let sem = ta1.agent.limits().sync.clone(); + let _p1 = sem.clone().try_acquire_owned().expect("permit 1"); + let _p2 = sem.clone().try_acquire_owned().expect("permit 2"); + let _p3 = sem.try_acquire_owned().expect("permit 3"); + + let (rtt_tx, _rtt_rx) = mpsc::channel(1024); + let ta2_transport = Transport::new(&ta2.agent.config().gossip, rtt_tx.clone()).await?; + + let res = parallel_sync( + &ta2.agent, + &ta2_transport, + vec![(ta1.agent.actor_id(), ta1.agent.gossip_addr())], + generate_sync(&ta2.bookie, ta2.agent.actor_id()).await, + ) + .await; + + // The rejection variant is decoded on the client side in + // `parallel_sync` (peer/mod.rs) and propagated via `SyncError::Rejection`. + let err = res.expect_err("sync should be rejected at capacity"); + let msg = format!("{err:#}"); + assert!( + msg.to_lowercase().contains("max concurrency"), + "unexpected error: {msg}" + ); + + drop((_p1, _p2, _p3)); + tripwire_tx.send(()).await.ok(); + tripwire_worker.await; + wait_for_all_pending_handles().await; + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 1)] async fn test_clear_empty_versions() -> eyre::Result<()> { _ = tracing_subscriber::fmt::try_init(); diff --git a/crates/corro-agent/src/api/peer/mod.rs b/crates/corro-agent/src/api/peer/mod.rs index dc4e92b8e..28c1de58f 100644 --- a/crates/corro-agent/src/api/peer/mod.rs +++ b/crates/corro-agent/src/api/peer/mod.rs @@ -959,7 +959,7 @@ fn encode_bipayload_msg( Ok(()) } -async fn encode_write_sync_msg( +pub(crate) async fn encode_write_sync_msg( codec: &mut LengthDelimitedCodec, encode_buf: &mut BytesMut, send_buf: &mut BytesMut, @@ -1470,25 +1470,6 @@ pub async fn serve_sync( trace!(actor_id = %their_actor_id, self_actor_id = %agent.actor_id(), "read clock"); - let _permit = match agent.limits().sync.try_acquire() { - Ok(permit) => permit, - Err(_) => { - // no permits! - encode_write_sync_msg( - &mut codec, - &mut encode_buf, - &mut send_buf, - SyncMessage::V1(SyncMessageV1::Rejection( - SyncRejectionV1::MaxConcurrencyReached, - )), - &mut write, - ) - .instrument(info_span!("write_sync_rejection")) - .await?; - return Ok(0); - } - }; - let sync_state = generate_sync(bookie, agent.actor_id()).await; // first, send the current sync state From 996bd6e6f714fce838a5420b7ed515ed1d61ee8e Mon Sep 17 00:00:00 2001 From: MorganaFuture Date: Mon, 18 May 2026 10:32:23 +0300 Subject: [PATCH 2/2] make configurable --- crates/corro-types/src/agent.rs | 3 ++- crates/corro-types/src/config.rs | 10 ++++++++++ 2 files changed, 12 insertions(+), 1 deletion(-) diff --git a/crates/corro-types/src/agent.rs b/crates/corro-types/src/agent.rs index 34c93b76e..75a63386b 100644 --- a/crates/corro-types/src/agent.rs +++ b/crates/corro-types/src/agent.rs @@ -123,6 +123,7 @@ pub struct Limits { impl Agent { pub fn new(config: AgentConfig) -> Self { + let max_concurrent_syncs = config.config.load().perf.max_concurrent_syncs; Self(Arc::new(AgentInner { actor_id: config.actor_id, pool: config.pool, @@ -144,7 +145,7 @@ impl Agent { schema: config.schema, cluster_id: ArcSwap::from_pointee(config.cluster_id), limits: Limits { - sync: Arc::new(Semaphore::new(3)), + sync: Arc::new(Semaphore::new(max_concurrent_syncs)), }, subs_manager: config.subs_manager, updates_manager: config.updates_manager, diff --git a/crates/corro-types/src/config.rs b/crates/corro-types/src/config.rs index 1bfeabc6a..5cfb04329 100755 --- a/crates/corro-types/src/config.rs +++ b/crates/corro-types/src/config.rs @@ -47,6 +47,10 @@ const fn default_processing_queue() -> usize { 20000 } +const fn default_max_concurrent_syncs() -> usize { + 3 +} + /// Used for the apply channel const fn default_huge_channel() -> usize { 2048 @@ -278,6 +282,11 @@ pub struct PerfConfig { // It's used to decide whether to wait for more changes for apply_queue_timeout ms or spawn a batch immediately #[serde(default = "default_batch_threshold_ratio")] pub apply_queue_batch_threshold_ratio: f64, + // Cap on concurrent inbound sync handlers. Enforced at the bi-stream + // accept site; further peers receive `MaxConcurrencyReached` and the + // stream is closed. + #[serde(default = "default_max_concurrent_syncs")] + pub max_concurrent_syncs: usize, } impl Default for PerfConfig { @@ -302,6 +311,7 @@ impl Default for PerfConfig { apply_queue_step_base: default_apply_batch_step(), apply_queue_max_batch_size: default_apply_batch_max(), apply_queue_batch_threshold_ratio: default_batch_threshold_ratio(), + max_concurrent_syncs: default_max_concurrent_syncs(), } } }