fix: align batcher worker lookup and configuration validation

Signed-off-by: WenyXu <wenymedia@gmail.com>

Cargo.lock

@@ -5626,6 +5626,7 @@ dependencies = [
  "catalog",
  "client",
  "common-base",
+ "common-batcher",
  "common-catalog",
  "common-config",
  "common-datasource",

src/common/batcher/src/flush_policy/timing.rs

@@ -33,16 +33,20 @@ pub struct TimingFlushPolicy {
 impl TimingFlushPolicy {
     /// Creates a policy, rejecting zero intervals and unrepresentable deadlines.
     pub fn try_new(flush_interval: Duration, max_batch_rows: NonZeroUsize) -> Option<Self> {
-        if flush_interval.is_zero() {
+        if !Self::validate(flush_interval) {
             return None;
         }
-        Instant::now().checked_add(flush_interval)?;
         Some(Self {
             flush_interval,
             max_batch_rows,
         })
     }
 
+    /// Checks that the interval is nonzero and its deadline is representable.
+    pub fn validate(flush_interval: Duration) -> bool {
+        !flush_interval.is_zero() && Instant::now().checked_add(flush_interval).is_some()
+    }
+
     /// Checks only the row trigger when a caller drives deadline events separately.
     pub fn reached_row_threshold<T>(&self, batch: &PendingBatch<T>) -> bool {
         !batch.is_empty() && batch.total_rows() >= self.max_batch_rows.get()
@@ -78,6 +82,16 @@ impl FlushPolicy for TimingFlushPolicy {
 mod tests {
     use super::*;
 
+    #[test]
+    fn test_validate_matches_construction() {
+        for interval in [Duration::ZERO, Duration::from_secs(1), Duration::MAX] {
+            assert_eq!(
+                TimingFlushPolicy::validate(interval),
+                TimingFlushPolicy::try_new(interval, NonZeroUsize::MIN).is_some()
+            );
+        }
+    }
+
     #[test]
     fn test_first_submission_deadline() {
         let policy =

src/common/batcher/src/worker_registry.rs

@@ -77,6 +77,10 @@ impl<K: Eq + Hash, T> WorkerRegistry<K, T> {
     where
         F: FnOnce() -> Sender<T>,
     {
+        if let Some(tx) = self.get(&key).await {
+            return tx;
+        }
+
         match self.workers.entry(key) {
             Entry::Occupied(mut entry) => {
                 if entry.get().is_closed() {
@@ -114,6 +118,27 @@ mod tests {
 
     use crate::worker_registry::WorkerRegistry;
 
+    #[tokio::test]
+    async fn test_live_lookup_uses_read_lock() {
+        let registry = Arc::new(WorkerRegistry::<_, ()>::new());
+        let (sender, _receiver) = registry.get_or_create(1, 1).await;
+        let guard = registry.workers.get(&1).unwrap();
+        let (result_tx, result_rx) = std::sync::mpsc::channel();
+        let worker_registry = registry.clone();
+        let runtime = tokio::runtime::Handle::current();
+        let thread = std::thread::spawn(move || {
+            let result = runtime.block_on(worker_registry.get_or_create(1, 1));
+            let _ = result_tx.send(result);
+        });
+        // Release the read guard before joining, even if lookup needs a write lock.
+        let result = result_rx.recv_timeout(std::time::Duration::from_secs(5));
+        drop(guard);
+        thread.join().unwrap();
+        let (reused, receiver) = result.expect("live lookup must not wait for an exclusive lock");
+        assert!(sender.same_channel(&reused));
+        assert!(receiver.is_none());
+    }
+
     #[tokio::test]
     async fn test_reuses_live_sender_and_returns_receiver_to_caller() {
         let registry = WorkerRegistry::new();

src/frontend/Cargo.toml

@@ -29,6 +29,7 @@ cache.workspace = true
 catalog.workspace = true
 client.workspace = true
 common-base.workspace = true
+common-batcher.workspace = true
 common-catalog.workspace = true
 common-config.workspace = true
 common-datasource.workspace = true

src/frontend/src/server.rs

@@ -567,7 +567,23 @@ mod tests {
         // returns `None`; in these cases no request can wait for a pending-row
         // flush, so the timeout must not be raised.
         type KnobMutator = fn(&mut FrontendOptions);
-        let cases: [(&str, KnobMutator); 4] = [
+        let cases: [(&str, KnobMutator); 9] = [
+            ("oversized flush concurrency", |opts| {
+                opts.prom_store.max_concurrent_flushes = usize::MAX
+            }),
+            ("oversized worker channel", |opts| {
+                opts.prom_store.worker_channel_capacity = usize::MAX
+            }),
+            ("oversized inflight limit", |opts| {
+                opts.prom_store.max_inflight_requests = usize::MAX
+            }),
+            ("oversized flow queue", |opts| {
+                opts.prom_store.flow_notification_queue_capacity =
+                    std::num::NonZeroUsize::new(usize::MAX).unwrap()
+            }),
+            ("unrepresentable deadline", |opts| {
+                opts.prom_store.pending_rows_flush_interval = Duration::MAX
+            }),
             ("zero max_batch_rows", |opts| {
                 opts.prom_store.max_batch_rows = 0
             }),

src/frontend/src/service_config/prom_store.rs

@@ -15,8 +15,10 @@
 use std::num::NonZeroUsize;
 use std::time::Duration;
 
+use common_batcher::flush_policy::timing::TimingFlushPolicy;
 use serde::{Deserialize, Serialize};
 use servers::prom_remote_write::validation::PromValidationMode;
+use tokio::sync::Semaphore;
 
 #[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
 pub struct PromStoreOptions {
@@ -64,18 +66,20 @@ fn default_flow_notification_queue_capacity() -> NonZeroUsize {
 }
 
 impl PromStoreOptions {
-    /// Returns whether the pending rows batcher can be enabled with these
-    /// options. Mirrors the enablement conditions of
-    /// `PendingRowsBatcher::try_new` in the servers crate, which returns
-    /// `None` when any of these knobs is zero.
+    /// Returns whether batching is enabled and its controls pass construction validation.
     pub fn pending_rows_batching_enabled(&self) -> bool {
         self.enable
             && self.with_metric_engine
-            && !self.pending_rows_flush_interval.is_zero()
+            && TimingFlushPolicy::validate(self.pending_rows_flush_interval)
             && self.max_batch_rows > 0
-            && self.max_concurrent_flushes > 0
-            && self.worker_channel_capacity > 0
-            && self.max_inflight_requests > 0
+            && [
+                self.max_concurrent_flushes,
+                self.worker_channel_capacity,
+                self.max_inflight_requests,
+                self.flow_notification_queue_capacity.get(),
+            ]
+            .into_iter()
+            .all(|capacity| (1..=Semaphore::MAX_PERMITS).contains(&capacity))
     }
 }
 
@@ -107,6 +111,42 @@ mod tests {
         default_worker_channel_capacity,
     };
 
+    #[test]
+    fn test_batching_capacity_boundaries() {
+        let max = tokio::sync::Semaphore::MAX_PERMITS;
+        let options = PromStoreOptions {
+            pending_rows_flush_interval: Duration::from_secs(5),
+            max_batch_rows: usize::MAX,
+            max_concurrent_flushes: max,
+            worker_channel_capacity: max,
+            max_inflight_requests: max,
+            flow_notification_queue_capacity: std::num::NonZeroUsize::new(max).unwrap(),
+            ..Default::default()
+        };
+        assert!(options.pending_rows_batching_enabled());
+        let cases = [
+            PromStoreOptions {
+                max_concurrent_flushes: max + 1,
+                ..options.clone()
+            },
+            PromStoreOptions {
+                worker_channel_capacity: max + 1,
+                ..options.clone()
+            },
+            PromStoreOptions {
+                max_inflight_requests: max + 1,
+                ..options.clone()
+            },
+            PromStoreOptions {
+                flow_notification_queue_capacity: std::num::NonZeroUsize::new(max + 1).unwrap(),
+                ..options
+            },
+        ];
+        for options in cases {
+            assert!(!options.pending_rows_batching_enabled());
+        }
+    }
+
     #[test]
     fn test_prom_store_options() {
         let default = PromStoreOptions::default();

src/servers/src/pending_rows_batcher.rs

@@ -52,7 +52,7 @@ use smallvec::SmallVec;
 use snafu::{OptionExt, ResultExt, ensure};
 use store_api::storage::{RegionId, TableId};
 use table::metadata::{TableInfo, TableInfoRef};
-use tokio::sync::{OwnedSemaphorePermit, broadcast, mpsc, oneshot};
+use tokio::sync::{OwnedSemaphorePermit, Semaphore, broadcast, mpsc, oneshot};
 
 use crate::error;
 use crate::error::{Error, Result};
@@ -378,15 +378,7 @@ impl PendingRowsBatcher {
         max_inflight_requests: usize,
         flow_notification_queue_capacity: NonZeroUsize,
     ) -> Option<Arc<Self>> {
-        // Disable the batcher if flush is disabled or configuration is invalid.
-        // Zero values for these knobs either cause panics (e.g., zero-capacity channels)
-        // or deadlocks (e.g., semaphores with no permits).
-        if flush_interval.is_zero()
-            || max_batch_rows == 0
-            || max_concurrent_flushes == 0
-            || worker_channel_capacity == 0
-            || max_inflight_requests == 0
-        {
+        if worker_channel_capacity == 0 || worker_channel_capacity > Semaphore::MAX_PERMITS {
             return None;
         }