diff --git a/Cargo.lock b/Cargo.lock index 1bb607c5f5..5fb917dae8 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -5626,6 +5626,7 @@ dependencies = [ "catalog", "client", "common-base", + "common-batcher", "common-catalog", "common-config", "common-datasource", diff --git a/src/common/batcher/src/flush_policy/timing.rs b/src/common/batcher/src/flush_policy/timing.rs index fd44b3a775..f9d431a6df 100644 --- a/src/common/batcher/src/flush_policy/timing.rs +++ b/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 { - 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(&self, batch: &PendingBatch) -> 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 = diff --git a/src/common/batcher/src/worker_registry.rs b/src/common/batcher/src/worker_registry.rs index 1a77a01c2f..f5dbf48e48 100644 --- a/src/common/batcher/src/worker_registry.rs +++ b/src/common/batcher/src/worker_registry.rs @@ -77,6 +77,10 @@ impl WorkerRegistry { where F: FnOnce() -> Sender, { + 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(); diff --git a/src/frontend/Cargo.toml b/src/frontend/Cargo.toml index 7ad24a466e..677dbbbe55 100644 --- a/src/frontend/Cargo.toml +++ b/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 diff --git a/src/frontend/src/server.rs b/src/frontend/src/server.rs index e112e5ae2a..1812dc71a5 100644 --- a/src/frontend/src/server.rs +++ b/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 }), diff --git a/src/frontend/src/service_config/prom_store.rs b/src/frontend/src/service_config/prom_store.rs index d2e386caae..f8429a8692 100644 --- a/src/frontend/src/service_config/prom_store.rs +++ b/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(); diff --git a/src/servers/src/pending_rows_batcher.rs b/src/servers/src/pending_rows_batcher.rs index 6f6991a542..442e8d6ef0 100644 --- a/src/servers/src/pending_rows_batcher.rs +++ b/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> { - // 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; }