From 5a12780516e428f04a759ea7d0dc37a9010b5762 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Sun, 14 Sep 2025 21:22:46 +0000 Subject: [PATCH 1/3] Refactor coalescer to distribute items fairly using round-robin Co-authored-by: miles.frankel --- src/server/coalescer.rs | 27 +++++++++------------------ 1 file changed, 9 insertions(+), 18 deletions(-) diff --git a/src/server/coalescer.rs b/src/server/coalescer.rs index 80a633b..79db7d7 100644 --- a/src/server/coalescer.rs +++ b/src/server/coalescer.rs @@ -118,28 +118,19 @@ impl PollCoalescer { } }; - // distribute the items to the pending polls. - // TODO: simplify this - // TODO: reuse across loops + // distribute the items to the pending polls fairly (round-robin) let mut distributed: Vec = (0..batch.len()).map(|_| Vec::new()).collect(); let mut remaining: Vec = batch.iter().map(|p| p.num_items).collect(); - let mut idx = 0usize; + let mut active: VecDeque = + (0..batch.len()).filter(|&i| remaining[i] > 0).collect(); + for item in items.into_iter() { - let mut placed = false; - for _ in 0..batch.len() { - let target = idx % batch.len(); - if remaining[target] > 0 { - distributed[target].push(item); - remaining[target] -= 1; - idx = idx.wrapping_add(1); - placed = true; - break; - } - idx = idx.wrapping_add(1); - } - if !placed { - break; + let Some(idx) = active.pop_front() else { break }; + distributed[idx].push(item); + remaining[idx] -= 1; + if remaining[idx] > 0 { + active.push_back(idx); } } From ef7bf8781a625bf058fa179c119f2e2711c60277 Mon Sep 17 00:00:00 2001 From: Miles Frankel Date: Sun, 14 Sep 2025 23:59:30 +0200 Subject: [PATCH 2/3] upd --- src/server/coalescer.rs | 25 +++++++++++- src/server/mod.rs | 19 +++++++++ tests/server_poll.rs | 87 ++++++++++++++++++++++++++++++++++++++++- 3 files changed, 128 insertions(+), 3 deletions(-) diff --git a/src/server/coalescer.rs b/src/server/coalescer.rs index 79db7d7..59c0bc7 100644 --- a/src/server/coalescer.rs +++ b/src/server/coalescer.rs @@ -4,13 +4,14 @@ use tokio::time::Duration; use capnp::message::{Builder, HeapAllocator, TypedReader}; +use crate::errors; use crate::storage::{RetriedStorage, Storage}; use std::collections::VecDeque; pub type PolledItems = Vec, crate::protocol::polled_item::Owned>>; pub type CoalescedPollResult = - std::result::Result, std::sync::Arc>; + std::result::Result, std::sync::Arc>; #[derive(Clone, Copy)] struct PollCoalescingConfig { @@ -54,6 +55,24 @@ impl PollCoalescer { this } + /// Create a coalescer with a custom batch window in milliseconds (tests/tuning). + pub fn with_batch_window_ms( + storage: Arc>, + batch_window_ms: u64, + ) -> Self { + let this = Self { + storage, + cfg: PollCoalescingConfig { + batch_window_ms, + ..PollCoalescingConfig::default() + }, + pending: Arc::new(tokio::sync::Mutex::new(VecDeque::new())), + wake: Arc::new(Notify::new()), + }; + this.spawn_batcher(); + this + } + fn spawn_batcher(&self) { let storage = Arc::clone(&self.storage); let cfg = self.cfg; @@ -126,7 +145,9 @@ impl PollCoalescer { (0..batch.len()).filter(|&i| remaining[i] > 0).collect(); for item in items.into_iter() { - let Some(idx) = active.pop_front() else { break }; + let Some(idx) = active.pop_front() else { + panic!( "invariant violated: active is empty; this means we requested more items than we wanted"); + }; distributed[idx].push(item); remaining[idx] -= 1; if remaining[idx] > 0 { diff --git a/src/server/mod.rs b/src/server/mod.rs index 192542b..384ca3e 100644 --- a/src/server/mod.rs +++ b/src/server/mod.rs @@ -37,6 +37,25 @@ impl Server { coalescer, } } + + /// Build a server with a custom coalescing window (ms) for polling. + pub fn new_with_coalescing_window_ms( + storage: Arc>, + notify: Arc, + shutdown_tx: watch::Sender, + batch_window_ms: u64, + ) -> Self { + let coalescer = Arc::new(PollCoalescer::with_batch_window_ms( + Arc::clone(&storage), + batch_window_ms, + )); + Self { + storage, + notify, + shutdown_tx, + coalescer, + } + } } const BACKGROUND_LEASE_EXPIRY_INTERVAL_SECS: u64 = 1; diff --git a/tests/server_poll.rs b/tests/server_poll.rs index b86bf39..9fff3c4 100644 --- a/tests/server_poll.rs +++ b/tests/server_poll.rs @@ -58,7 +58,14 @@ fn start_test_server() -> TestServerHandle { Arc::clone(¬ify), shutdown_tx.clone(), ); - let server = Server::new(storage, notify, shutdown_tx.clone()); + // Use a longer coalescing window in tests to increase batching and reduce flakiness. + // 50ms is a reasonable tradeoff for deterministic coalescing without slowing tests too much. + let server = Server::new_with_coalescing_window_ms( + storage, + notify, + shutdown_tx.clone(), + 50, + ); let queue_client: QueueClient = capnp_rpc::new_client(server); // Indicate that the server is ready to accept connections @@ -437,3 +444,81 @@ async fn coalesced_concurrent_polls_distribute_items() { }) .await; } + +#[tokio::test(flavor = "current_thread")] +async fn coalesced_polls_round_robin_fairness_unequal_demands() { + let handle = start_test_server(); + let addr = handle.addr; + + with_client(addr, |queue_client| async move { + // Add 6 immediately visible items + let mut add = queue_client.add_request(); + { + let req = add.get().init_req(); + let mut items = req.init_items(6); + for i in 0..6 { + let mut item = items.reborrow().get(i as u32); + item.set_contents(format!("fair-{i}").as_bytes()); + item.set_visibility_timeout_secs(0); + } + } + let _ = add.send().promise.await.unwrap(); + + // Spawn 3 concurrent polls with unequal demands: 5, 3, 1 + // Round-robin fairness should distribute the 6 items as counts {3,2,1} (in some order). + let quotas = [5u32, 3u32, 1u32]; + let mut handles = Vec::new(); + for idx in 0..quotas.len() { + let q = quotas[idx]; + let client = queue_client.clone(); + handles.push(tokio::task::spawn_local(async move { + let mut poll = client.poll_request(); + { + let mut req = poll.get().init_req(); + req.set_lease_validity_secs(30); + req.set_num_items(q); + req.set_timeout_secs(1); + } + let reply = poll.send().promise.await.unwrap(); + let resp = reply.get().unwrap().get_resp().unwrap(); + let items = resp.get_items().unwrap(); + (idx, items.len() as usize) + })); + } + + // Gather per-poll counts + let mut counts: Vec<(usize, usize)> = Vec::new(); + for h in handles { + counts.push(h.await.unwrap()); + } + + // Ensure total items returned sums to 6 + let total_returned: usize = counts.iter().map(|(_, c)| *c).sum(); + assert_eq!( + total_returned, 6, + "all 6 items should be distributed across polls" + ); + + // Fairness assertion: with quotas [5,3,1] and 6 items total, round-robin distribution + // yields the multiset of delivered counts {3,2,1} regardless of poll arrival order. + let mut delivered: Vec = counts.iter().map(|(_, c)| *c).collect(); + delivered.sort_unstable_by(|a, b| b.cmp(a)); // descending + assert_eq!( + delivered, + vec![3, 2, 1], + "expected fair round-robin distribution {{3,2,1}}" + ); + + // Also assert no poll exceeded its requested num_items + for (i, c) in counts.iter() { + assert!( + *c as u32 <= quotas[*i], + "poll {} received {} which exceeds its quota {}", + i, + c, + quotas[*i] + ); + } + }) + .await; +} From fe5d5ebab7c6e74c100d37708aded9564a55ca42 Mon Sep 17 00:00:00 2001 From: Miles Frankel Date: Mon, 15 Sep 2025 00:05:37 +0200 Subject: [PATCH 3/3] clip --- tests/server_poll.rs | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/tests/server_poll.rs b/tests/server_poll.rs index 9fff3c4..b56b1da 100644 --- a/tests/server_poll.rs +++ b/tests/server_poll.rs @@ -468,8 +468,7 @@ async fn coalesced_polls_round_robin_fairness_unequal_demands() { // Round-robin fairness should distribute the 6 items as counts {3,2,1} (in some order). let quotas = [5u32, 3u32, 1u32]; let mut handles = Vec::new(); - for idx in 0..quotas.len() { - let q = quotas[idx]; + for (idx, &q) in quotas.iter().enumerate() { let client = queue_client.clone(); handles.push(tokio::task::spawn_local(async move { let mut poll = client.poll_request();