Bound bi-stream handler concurrency at accept site

crates/corro-agent/src/agent/bi.rs

@@ -1,15 +1,17 @@
-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;
 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()

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();

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

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,

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(),
         }
     }
 }