From 6d124d6581731106ea2999db76edaf86db398738 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Sun, 14 Sep 2025 22:19:39 +0000 Subject: [PATCH 1/4] feat: Add poll coalescing configuration options Co-authored-by: miles.frankel --- src/bin/queueber/main.rs | 35 ++++++++++++++++++++-- src/server/coalescer.rs | 63 ++++++++++++++++++++++++++++++++++++---- src/server/mod.rs | 16 ++++++++++ 3 files changed, 107 insertions(+), 7 deletions(-) diff --git a/src/bin/queueber/main.rs b/src/bin/queueber/main.rs index 5fe624b..bab5811 100644 --- a/src/bin/queueber/main.rs +++ b/src/bin/queueber/main.rs @@ -8,7 +8,7 @@ use clap::Parser; use color_eyre::Result; use futures::AsyncReadExt; use queueber::{ - server::Server, + server::{PollCoalescingConfig, Server}, storage::{RetriedStorage, Storage}, }; use socket2::{Domain, Protocol, Socket, Type}; @@ -37,6 +37,22 @@ struct Args { /// Number of RPC worker threads. Defaults to available_parallelism. #[arg(long = "workers")] workers: Option, + + /// Enable poll request coalescing + #[arg(long = "coalesce", default_value_t = true)] + coalesce: bool, + + /// Max concurrent poll requests batched per DB call + #[arg(long = "coalesce-max-batch-size", default_value_t = 64)] + coalesce_max_batch_size: usize, + + /// Max total items returned per batched DB call + #[arg(long = "coalesce-max-batch-items", default_value_t = 512)] + coalesce_max_batch_items: usize, + + /// Max batching window in milliseconds + #[arg(long = "coalesce-batch-window-ms", default_value_t = 1)] + coalesce_batch_window_ms: u64, } // NOTE: to use the console you need "RUST_LOG=tokio=trace,runtime=trace" @@ -92,6 +108,12 @@ async fn main() -> Result<()> { let shutdown_tx_cloned = shutdown_tx.clone(); let mut shutdown_rx_worker = shutdown_rx.clone(); let addr_for_worker: SocketAddr = addr; + let coalesce_enabled = args.coalesce; + let coalesce_cfg = PollCoalescingConfig::new( + args.coalesce_max_batch_size, + args.coalesce_max_batch_items, + args.coalesce_batch_window_ms, + ); let thread_name = format!("rpc-worker-{}", i); let handle = std::thread::Builder::new() @@ -102,7 +124,16 @@ async fn main() -> Result<()> { .build() .expect("build worker runtime"); rt.block_on(async move { - let server = Server::new(storage_cloned, notify_cloned, shutdown_tx_cloned); + let server = if coalesce_enabled { + Server::new_with_coalescer_config( + storage_cloned, + notify_cloned, + shutdown_tx_cloned, + coalesce_cfg, + ) + } else { + Server::new(storage_cloned, notify_cloned, shutdown_tx_cloned) + }; let queue_client: queueber::protocol::queue::Client = capnp_rpc::new_client(server); // Each worker owns one Server; clone client per-connection. diff --git a/src/server/coalescer.rs b/src/server/coalescer.rs index 80a633b..d52de53 100644 --- a/src/server/coalescer.rs +++ b/src/server/coalescer.rs @@ -12,11 +12,11 @@ pub type PolledItems = pub type CoalescedPollResult = std::result::Result, std::sync::Arc>; -#[derive(Clone, Copy)] -struct PollCoalescingConfig { - max_batch_size: usize, - max_batch_items: usize, - batch_window_ms: u64, +#[derive(Clone, Copy, Debug)] +pub struct PollCoalescingConfig { + pub max_batch_size: usize, + pub max_batch_items: usize, + pub batch_window_ms: u64, } impl Default for PollCoalescingConfig { @@ -29,6 +29,16 @@ impl Default for PollCoalescingConfig { } } +impl PollCoalescingConfig { + pub fn new(max_batch_size: usize, max_batch_items: usize, batch_window_ms: u64) -> Self { + Self { + max_batch_size, + max_batch_items, + batch_window_ms, + } + } +} + struct PendingPoll { num_items: usize, lease_validity_secs: u64, @@ -54,6 +64,17 @@ impl PollCoalescer { this } + pub fn with_config(storage: Arc>, cfg: PollCoalescingConfig) -> Self { + let this = Self { + storage, + cfg, + 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; @@ -172,4 +193,36 @@ impl PollCoalescer { Err(_canceled) => Ok(None), } } + + #[cfg(test)] + pub(crate) fn config(&self) -> PollCoalescingConfig { + self.cfg + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test] + async fn with_config_sets_values() { + let dir = tempfile::tempdir().unwrap(); + let storage = Arc::new(RetriedStorage::new(Storage::new(dir.path()).unwrap())); + let cfg = PollCoalescingConfig::new(10, 100, 5); + let c = PollCoalescer::with_config(Arc::clone(&storage), cfg); + let got = c.config(); + assert_eq!(got.max_batch_size, 10); + assert_eq!(got.max_batch_items, 100); + assert_eq!(got.batch_window_ms, 5); + } + + #[tokio::test] + async fn default_config_is_reasonable() { + let dir = tempfile::tempdir().unwrap(); + let storage = Arc::new(RetriedStorage::new(Storage::new(dir.path()).unwrap())); + let c = PollCoalescer::new(storage); + let got = c.config(); + assert!(got.max_batch_size >= 1); + assert!(got.max_batch_items >= 1); + } } diff --git a/src/server/mod.rs b/src/server/mod.rs index 192542b..825d897 100644 --- a/src/server/mod.rs +++ b/src/server/mod.rs @@ -12,6 +12,7 @@ use crate::protocol::queue::{ use crate::storage::{RetriedStorage, Storage}; mod coalescer; use coalescer::PollCoalescer; +pub use coalescer::PollCoalescingConfig; // https://github.com/capnproto/capnproto-rust/tree/master/capnp-rpc // https://github.com/capnproto/capnproto-rust/blob/master/capnp-rpc/examples/hello-world/server.rs @@ -37,6 +38,21 @@ impl Server { coalescer, } } + + pub fn new_with_coalescer_config( + storage: Arc>, + notify: Arc, + shutdown_tx: watch::Sender, + cfg: PollCoalescingConfig, + ) -> Self { + let coalescer = Arc::new(PollCoalescer::with_config(Arc::clone(&storage), cfg)); + Self { + storage, + notify, + shutdown_tx, + coalescer, + } + } } const BACKGROUND_LEASE_EXPIRY_INTERVAL_SECS: u64 = 1; From 7ddb1a27dcd6e5aba2fef3bf8796831a524bf9c6 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Sun, 14 Sep 2025 22:45:02 +0000 Subject: [PATCH 2/4] cli: derive coalescer config from PollCoalescingConfig::default(); use Option flags to override --- src/bin/queueber/main.rs | 27 ++++++++++++++++----------- 1 file changed, 16 insertions(+), 11 deletions(-) diff --git a/src/bin/queueber/main.rs b/src/bin/queueber/main.rs index ebb86f1..caa35d1 100644 --- a/src/bin/queueber/main.rs +++ b/src/bin/queueber/main.rs @@ -39,16 +39,16 @@ struct Args { workers: Option, /// Max concurrent poll requests batched per DB call - #[arg(long = "coalesce-max-batch-size", default_value_t = 64)] - coalesce_max_batch_size: usize, + #[arg(long = "coalesce-max-batch-size")] + coalesce_max_batch_size: Option, /// Max total items returned per batched DB call - #[arg(long = "coalesce-max-batch-items", default_value_t = 512)] - coalesce_max_batch_items: usize, + #[arg(long = "coalesce-max-batch-items")] + coalesce_max_batch_items: Option, /// Max batching window in milliseconds - #[arg(long = "coalesce-batch-window-ms", default_value_t = 1)] - coalesce_batch_window_ms: u64, + #[arg(long = "coalesce-batch-window-ms")] + coalesce_batch_window_ms: Option, } // NOTE: to use the console you need "RUST_LOG=tokio=trace,runtime=trace" @@ -104,11 +104,16 @@ async fn main() -> Result<()> { let shutdown_tx_cloned = shutdown_tx.clone(); let mut shutdown_rx_worker = shutdown_rx.clone(); let addr_for_worker: SocketAddr = addr; - let coalesce_cfg = PollCoalescingConfig::new( - args.coalesce_max_batch_size, - args.coalesce_max_batch_items, - args.coalesce_batch_window_ms, - ); + let mut coalesce_cfg = PollCoalescingConfig::default(); + if let Some(v) = args.coalesce_max_batch_size { + coalesce_cfg.max_batch_size = v; + } + if let Some(v) = args.coalesce_max_batch_items { + coalesce_cfg.max_batch_items = v; + } + if let Some(v) = args.coalesce_batch_window_ms { + coalesce_cfg.batch_window_ms = v; + } let thread_name = format!("rpc-worker-{}", i); let handle = std::thread::Builder::new() From a320a22797952e1c1b59c90bfede991926454f67 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Sun, 14 Sep 2025 22:57:38 +0000 Subject: [PATCH 3/4] cli: document effective defaults in coalescer flag help --- src/bin/queueber/main.rs | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/bin/queueber/main.rs b/src/bin/queueber/main.rs index caa35d1..421e4d2 100644 --- a/src/bin/queueber/main.rs +++ b/src/bin/queueber/main.rs @@ -38,15 +38,15 @@ struct Args { #[arg(long = "workers")] workers: Option, - /// Max concurrent poll requests batched per DB call + /// Max concurrent poll requests batched per DB call (default: 64) #[arg(long = "coalesce-max-batch-size")] coalesce_max_batch_size: Option, - /// Max total items returned per batched DB call + /// Max total items returned per batched DB call (default: 512) #[arg(long = "coalesce-max-batch-items")] coalesce_max_batch_items: Option, - /// Max batching window in milliseconds + /// Max batching window in milliseconds (default: 1) #[arg(long = "coalesce-batch-window-ms")] coalesce_batch_window_ms: Option, } From 7602816b9d2a0fb52dde5a8feb9bbaa76a02212e Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Mon, 15 Sep 2025 14:25:14 +0000 Subject: [PATCH 4/4] cli: tie help defaults to PollCoalescingConfig defaults via associated constants --- src/bin/queueber/main.rs | 33 ++++++++++++++------------------- src/server/coalescer.rs | 24 ++++++++++++++---------- 2 files changed, 28 insertions(+), 29 deletions(-) diff --git a/src/bin/queueber/main.rs b/src/bin/queueber/main.rs index 421e4d2..14ab350 100644 --- a/src/bin/queueber/main.rs +++ b/src/bin/queueber/main.rs @@ -38,17 +38,17 @@ struct Args { #[arg(long = "workers")] workers: Option, - /// Max concurrent poll requests batched per DB call (default: 64) - #[arg(long = "coalesce-max-batch-size")] - coalesce_max_batch_size: Option, + /// Max concurrent poll requests batched per DB call + #[arg(long = "coalesce-max-batch-size", default_value_t = PollCoalescingConfig::DEFAULT_MAX_BATCH_SIZE)] + coalesce_max_batch_size: usize, - /// Max total items returned per batched DB call (default: 512) - #[arg(long = "coalesce-max-batch-items")] - coalesce_max_batch_items: Option, + /// Max total items returned per batched DB call + #[arg(long = "coalesce-max-batch-items", default_value_t = PollCoalescingConfig::DEFAULT_MAX_BATCH_ITEMS)] + coalesce_max_batch_items: usize, - /// Max batching window in milliseconds (default: 1) - #[arg(long = "coalesce-batch-window-ms")] - coalesce_batch_window_ms: Option, + /// Max batching window in milliseconds + #[arg(long = "coalesce-batch-window-ms", default_value_t = PollCoalescingConfig::DEFAULT_BATCH_WINDOW_MS)] + coalesce_batch_window_ms: u64, } // NOTE: to use the console you need "RUST_LOG=tokio=trace,runtime=trace" @@ -104,16 +104,11 @@ async fn main() -> Result<()> { let shutdown_tx_cloned = shutdown_tx.clone(); let mut shutdown_rx_worker = shutdown_rx.clone(); let addr_for_worker: SocketAddr = addr; - let mut coalesce_cfg = PollCoalescingConfig::default(); - if let Some(v) = args.coalesce_max_batch_size { - coalesce_cfg.max_batch_size = v; - } - if let Some(v) = args.coalesce_max_batch_items { - coalesce_cfg.max_batch_items = v; - } - if let Some(v) = args.coalesce_batch_window_ms { - coalesce_cfg.batch_window_ms = v; - } + let coalesce_cfg = PollCoalescingConfig::new( + args.coalesce_max_batch_size, + args.coalesce_max_batch_items, + args.coalesce_batch_window_ms, + ); let thread_name = format!("rpc-worker-{}", i); let handle = std::thread::Builder::new() diff --git a/src/server/coalescer.rs b/src/server/coalescer.rs index 7c1b9c5..f41aed6 100644 --- a/src/server/coalescer.rs +++ b/src/server/coalescer.rs @@ -20,17 +20,11 @@ pub struct PollCoalescingConfig { pub batch_window_ms: u64, } -impl Default for PollCoalescingConfig { - fn default() -> Self { - Self { - max_batch_size: 64, - max_batch_items: 512, - batch_window_ms: 1, - } - } -} - impl PollCoalescingConfig { + pub const DEFAULT_MAX_BATCH_SIZE: usize = 64; + pub const DEFAULT_MAX_BATCH_ITEMS: usize = 512; + pub const DEFAULT_BATCH_WINDOW_MS: u64 = 1; + pub fn new(max_batch_size: usize, max_batch_items: usize, batch_window_ms: u64) -> Self { Self { max_batch_size, @@ -40,6 +34,16 @@ impl PollCoalescingConfig { } } +impl Default for PollCoalescingConfig { + fn default() -> Self { + Self { + max_batch_size: Self::DEFAULT_MAX_BATCH_SIZE, + max_batch_items: Self::DEFAULT_MAX_BATCH_ITEMS, + batch_window_ms: Self::DEFAULT_BATCH_WINDOW_MS, + } + } +} + struct PendingPoll { num_items: usize, lease_validity_secs: u64,