diff --git a/src/bin/queueber/main.rs b/src/bin/queueber/main.rs index 5fe624b..14ab350 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,18 @@ struct Args { /// Number of RPC worker threads. Defaults to available_parallelism. #[arg(long = "workers")] workers: 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 + #[arg(long = "coalesce-max-batch-items", default_value_t = PollCoalescingConfig::DEFAULT_MAX_BATCH_ITEMS)] + coalesce_max_batch_items: usize, + + /// 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" @@ -92,6 +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 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 +119,12 @@ 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 = Server::new_with_coalescer_config( + storage_cloned, + notify_cloned, + shutdown_tx_cloned, + coalesce_cfg, + ); 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 59c0bc7..f41aed6 100644 --- a/src/server/coalescer.rs +++ b/src/server/coalescer.rs @@ -13,19 +13,33 @@ 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 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, + max_batch_items, + batch_window_ms, + } + } } impl Default for PollCoalescingConfig { fn default() -> Self { Self { - max_batch_size: 64, - max_batch_items: 512, - batch_window_ms: 1, + max_batch_size: Self::DEFAULT_MAX_BATCH_SIZE, + max_batch_items: Self::DEFAULT_MAX_BATCH_ITEMS, + batch_window_ms: Self::DEFAULT_BATCH_WINDOW_MS, } } } @@ -55,6 +69,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 + } + /// Create a coalescer with a custom batch window in milliseconds (tests/tuning). pub fn with_batch_window_ms( storage: Arc>, @@ -184,4 +209,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 384ca3e..04d71ec 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 @@ -38,6 +39,21 @@ impl Server { } } + 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, + } + } + /// Build a server with a custom coalescing window (ms) for polling. pub fn new_with_coalescing_window_ms( storage: Arc>,