Add request coalescer config to cli flags

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<usize>,
+
+    /// 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.

src/server/coalescer.rs

@@ -13,19 +13,33 @@ pub type PolledItems =
 pub type CoalescedPollResult =
     std::result::Result<Option<([u8; 16], PolledItems)>, std::sync::Arc<errors::Error>>;
 
-#[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<RetriedStorage<Storage>>, 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<RetriedStorage<Storage>>,
@@ -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);
+    }
 }

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<RetriedStorage<Storage>>,
+        notify: Arc<Notify>,
+        shutdown_tx: watch::Sender<bool>,
+        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<RetriedStorage<Storage>>,