From e0915a8380a4dd5bb292f1dce0cae5daa8db7293 Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Wed, 16 Sep 2026 09:21:49 +0000 Subject: [PATCH] refactor: reuse common batching components in Prom ingestion (#9114) * refactor: compose Prom ingestion with common batcher components Signed-off-by: WenyXu * refactor: share bulk insert IPC encoding Signed-off-by: WenyXu * refactor: centralize batcher worker channel creation Signed-off-by: WenyXu * refactor: share timestamp extraction for batched writes Signed-off-by: WenyXu * refactor: preserve sharded batcher worker lookup Signed-off-by: WenyXu * fix: align batcher worker lookup and configuration validation Signed-off-by: WenyXu * refactor: validate batcher limits inside fallible constructors Signed-off-by: WenyXu --------- Signed-off-by: WenyXu --- Cargo.lock | 11 + Cargo.toml | 2 + src/common/batcher/Cargo.toml | 19 + src/common/batcher/src/flush_limiter.rs | 73 +++ src/common/batcher/src/flush_policy.rs | 43 ++ src/common/batcher/src/flush_policy/timing.rs | 157 ++++++ src/common/batcher/src/flush_timer.rs | 141 +++++ src/common/batcher/src/lib.rs | 27 + src/common/batcher/src/notifier.rs | 145 +++++ src/common/batcher/src/pending_batch.rs | 115 ++++ src/common/batcher/src/pending_worker.rs | 260 +++++++++ src/common/batcher/src/request_limiter.rs | 115 ++++ src/common/batcher/src/worker_registry.rs | 265 +++++++++ src/common/grpc/src/flight.rs | 75 ++- src/datatypes/src/timestamp.rs | 60 +- src/frontend/Cargo.toml | 1 + src/frontend/src/server.rs | 18 +- src/frontend/src/service_config/prom_store.rs | 56 +- src/operator/src/bulk_insert.rs | 23 +- src/operator/src/req_convert/insert.rs | 2 + .../src/req_convert/insert/timestamps.rs | 38 ++ src/servers/Cargo.toml | 1 + src/servers/src/pending_rows_batcher.rs | 518 ++++++++++-------- 23 files changed, 1907 insertions(+), 258 deletions(-) create mode 100644 src/common/batcher/Cargo.toml create mode 100644 src/common/batcher/src/flush_limiter.rs create mode 100644 src/common/batcher/src/flush_policy.rs create mode 100644 src/common/batcher/src/flush_policy/timing.rs create mode 100644 src/common/batcher/src/flush_timer.rs create mode 100644 src/common/batcher/src/lib.rs create mode 100644 src/common/batcher/src/notifier.rs create mode 100644 src/common/batcher/src/pending_batch.rs create mode 100644 src/common/batcher/src/pending_worker.rs create mode 100644 src/common/batcher/src/request_limiter.rs create mode 100644 src/common/batcher/src/worker_registry.rs create mode 100644 src/operator/src/req_convert/insert/timestamps.rs diff --git a/Cargo.lock b/Cargo.lock index fd9eeaefdd..61005e9184 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2317,6 +2317,15 @@ dependencies = [ "zeroize", ] +[[package]] +name = "common-batcher" +version = "1.3.0-alpha.1" +dependencies = [ + "dashmap", + "futures", + "tokio", +] + [[package]] name = "common-catalog" version = "1.3.0-alpha.1" @@ -5669,6 +5678,7 @@ dependencies = [ "catalog", "client", "common-base", + "common-batcher", "common-catalog", "common-config", "common-datasource", @@ -13596,6 +13606,7 @@ dependencies = [ "chrono", "client", "common-base", + "common-batcher", "common-catalog", "common-decimal", "common-error", diff --git a/Cargo.toml b/Cargo.toml index a54a008180..730901925a 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -8,6 +8,7 @@ members = [ "src/client", "src/cmd", "src/common/base", + "src/common/batcher", "src/common/catalog", "src/common/config", "src/common/datasource", @@ -287,6 +288,7 @@ cli = { path = "src/cli" } client = { path = "src/client" } cmd = { path = "src/cmd", default-features = false } common-base = { path = "src/common/base" } +common-batcher = { path = "src/common/batcher" } common-catalog = { path = "src/common/catalog" } common-config = { path = "src/common/config" } common-datasource = { path = "src/common/datasource" } diff --git a/src/common/batcher/Cargo.toml b/src/common/batcher/Cargo.toml new file mode 100644 index 0000000000..5d6ace4b2d --- /dev/null +++ b/src/common/batcher/Cargo.toml @@ -0,0 +1,19 @@ +[package] +name = "common-batcher" +version.workspace = true +edition.workspace = true +license.workspace = true + +[lib] +path = "src/lib.rs" + +[lints] +workspace = true + +[dependencies] +dashmap.workspace = true +futures.workspace = true +tokio.workspace = true + +[dev-dependencies] +tokio = { workspace = true, features = ["test-util"] } diff --git a/src/common/batcher/src/flush_limiter.rs b/src/common/batcher/src/flush_limiter.rs new file mode 100644 index 0000000000..98b6b2a2f0 --- /dev/null +++ b/src/common/batcher/src/flush_limiter.rs @@ -0,0 +1,73 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::sync::Arc; + +use tokio::sync::{AcquireError, OwnedSemaphorePermit, Semaphore}; + +/// Shared execution budget. One permit represents one complete flush, not one RPC. +#[derive(Clone)] +pub struct FlushLimiter { + semaphore: Arc, +} + +impl FlushLimiter { + /// Returns `None` if the limit is zero or exceeds Tokio's supported semaphore capacity. + pub fn try_new(max_concurrent_flushes: usize) -> Option { + let permits = max_concurrent_flushes; + ((1..=Semaphore::MAX_PERMITS).contains(&permits)).then(|| Self { + semaphore: Arc::new(Semaphore::new(permits)), + }) + } + + /// The caller must retain the returned permit until its flush finishes. + pub async fn acquire(&self) -> Result { + self.semaphore.clone().acquire_owned().await + } +} + +#[cfg(test)] +mod tests { + use std::future::Future; + + use super::*; + + #[test] + fn test_invalid_capacity() { + for capacity in [0, Semaphore::MAX_PERMITS + 1, usize::MAX] { + assert!(FlushLimiter::try_new(capacity).is_none()); + } + for capacity in [1, Semaphore::MAX_PERMITS] { + assert!(FlushLimiter::try_new(capacity).is_some()); + } + } + + #[tokio::test] + async fn test_clones_share_flush_budget() { + let limiter = FlushLimiter::try_new(1).unwrap(); + let other = limiter.clone(); + let permit = limiter.acquire().await.unwrap(); + let waiting = other.acquire(); + tokio::pin!(waiting); + // Poll admission without a sleep or an assumption about scheduler timing. + assert!( + std::future::poll_fn(|cx| { + std::task::Poll::Ready(waiting.as_mut().poll(cx).is_pending()) + }) + .await + ); + drop(permit); + let _next = waiting.await.unwrap(); + } +} diff --git a/src/common/batcher/src/flush_policy.rs b/src/common/batcher/src/flush_policy.rs new file mode 100644 index 0000000000..14804f815e --- /dev/null +++ b/src/common/batcher/src/flush_policy.rs @@ -0,0 +1,43 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +pub mod timing; + +use tokio::time::Instant; + +use crate::pending_batch::PendingBatch; + +/// The event that caused the worker to reconsider flushing. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum FlushTrigger { + /// A complete submission was appended to the batch. + Submission, + /// The flush timer reached its deadline. + Deadline, +} + +/// Decides when a pending batch is ready without owning its payloads or execution. +/// +/// A deadline wakeup must either allow flushing or produce a future deadline; +/// otherwise the caller would repeatedly poll an expired timer. +pub trait FlushPolicy { + /// Returns the next timer deadline, or `None` when no timer is needed. + fn deadline(&self, batch: &PendingBatch) -> Option; + + /// Returns whether the batch should be flushed for this event at `now`. + /// + /// Workers never flush empty batches, regardless of this decision. + fn should_flush(&self, batch: &PendingBatch, now: Instant, trigger: FlushTrigger) + -> bool; +} diff --git a/src/common/batcher/src/flush_policy/timing.rs b/src/common/batcher/src/flush_policy/timing.rs new file mode 100644 index 0000000000..56a9172eaa --- /dev/null +++ b/src/common/batcher/src/flush_policy/timing.rs @@ -0,0 +1,157 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::time::Duration; + +use tokio::time::Instant; + +use crate::flush_policy::{FlushPolicy, FlushTrigger}; +use crate::pending_batch::PendingBatch; + +/// Decides whether the first-submission deadline or row threshold requires a flush. +/// +/// The row threshold is checked after appending a complete submission. It is not a +/// hard batch-size limit and never splits a submission. +#[derive(Debug, Clone, Copy)] +pub struct TimingFlushPolicy { + flush_interval: Duration, + max_batch_rows: usize, +} + +impl TimingFlushPolicy { + /// Creates a policy, rejecting zero limits, zero intervals, and unrepresentable deadlines. + pub fn try_new(flush_interval: Duration, max_batch_rows: usize) -> Option { + if max_batch_rows == 0 || !Self::validate(flush_interval) { + return None; + } + 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 + } +} + +impl FlushPolicy for TimingFlushPolicy { + /// Returns the deadline anchored to the first submission, or `None` for an empty batch. + /// + /// An unrepresentable deadline falls back to the first submission time so that + /// extreme timestamps cannot leave a nonempty batch waiting indefinitely. + fn deadline(&self, batch: &PendingBatch) -> Option { + batch + .first_submitted_at() + .map(|first| first.checked_add(self.flush_interval).unwrap_or(first)) + } + + /// Returns whether a nonempty batch has reached either trigger. + fn should_flush( + &self, + batch: &PendingBatch, + now: Instant, + trigger: FlushTrigger, + ) -> bool { + !batch.is_empty() + && (self.reached_row_threshold(batch) + || (trigger == FlushTrigger::Deadline + && self.deadline(batch).is_some_and(|deadline| now >= deadline))) + } +} + +#[cfg(test)] +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, 1).is_some() + ); + } + } + + #[test] + fn test_first_submission_deadline() { + let policy = TimingFlushPolicy::try_new(Duration::from_millis(10), 100).unwrap(); + let mut batch = PendingBatch::new(); + let first = Instant::now(); + assert_eq!(policy.deadline(&batch), None); + assert!(!policy.should_flush(&batch, first, FlushTrigger::Deadline)); + batch.push(1, 1, first); + batch.push(2, 1, first + Duration::from_millis(9)); + assert_eq!( + policy.deadline(&batch), + Some(first + Duration::from_millis(10)) + ); + assert!(!policy.should_flush( + &batch, + first + Duration::from_millis(9), + FlushTrigger::Deadline + )); + assert!(policy.should_flush( + &batch, + first + Duration::from_millis(10), + FlushTrigger::Deadline + )); + assert!(!policy.reached_row_threshold(&batch)); + assert!(!policy.should_flush( + &batch, + first + Duration::from_millis(10), + FlushTrigger::Submission, + )); + assert_eq!(batch.take(), vec![1, 2]); + let next = first + Duration::from_millis(20); + batch.push(3, 0, next); + assert!(!policy.should_flush(&batch, next, FlushTrigger::Deadline)); + assert!(policy.should_flush( + &batch, + next + Duration::from_millis(10), + FlushTrigger::Deadline + )); + } + + #[test] + fn test_row_threshold_preserves_complete_submissions() { + let policy = TimingFlushPolicy::try_new(Duration::from_secs(1), 3).unwrap(); + let now = Instant::now(); + for rows in [2, 3, 4] { + let mut batch = PendingBatch::new(); + batch.push(vec![0; rows], rows, now); + assert_eq!( + policy.should_flush(&batch, now, FlushTrigger::Submission), + rows >= 3 + ); + assert_eq!(policy.reached_row_threshold(&batch), rows >= 3); + assert_eq!(batch.take(), vec![vec![0; rows]]); + } + } + + #[test] + fn test_invalid_interval() { + assert!(TimingFlushPolicy::try_new(Duration::from_secs(1), 0).is_none()); + let max_batch_rows = 1; + assert!(TimingFlushPolicy::try_new(Duration::ZERO, max_batch_rows).is_none()); + assert!(TimingFlushPolicy::try_new(Duration::MAX, max_batch_rows).is_none()); + } +} diff --git a/src/common/batcher/src/flush_timer.rs b/src/common/batcher/src/flush_timer.rs new file mode 100644 index 0000000000..7ae8173201 --- /dev/null +++ b/src/common/batcher/src/flush_timer.rs @@ -0,0 +1,141 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::pin::Pin; + +use tokio::time::{Instant, Sleep, sleep_until}; + +/// Maintains one reusable timer for the current flush deadline. +/// +/// Creating a timer does not require a Tokio runtime. Arming it requires a Tokio +/// runtime with time enabled. Disabling retains the allocation for the next arm. +#[derive(Default)] +pub struct FlushTimer { + deadline: Option, + sleep: Option>>, +} + +impl FlushTimer { + /// Creates a disabled timer without allocating a Tokio sleep. + pub fn new() -> Self { + Self::default() + } + + /// Arms or disables the timer, resetting its sleep only when the deadline changes. + /// + /// # Panics + /// + /// Arming for the first time panics outside a Tokio runtime with time enabled. + pub fn set_deadline(&mut self, deadline: Option) { + if self.deadline == deadline { + return; + } + self.deadline = deadline; + if let Some(deadline) = deadline { + match &mut self.sleep { + Some(sleep) => sleep.as_mut().reset(deadline), + None => self.sleep = Some(Box::pin(sleep_until(deadline))), + } + } + } + + /// Waits until the armed deadline, or indefinitely while disabled. + /// + /// Cancelling this wait does not change the deadline. An elapsed timer stays + /// ready until the caller changes or disables its deadline. + pub async fn wait(&mut self) { + if self.deadline.is_some() + && let Some(sleep) = &mut self.sleep + { + sleep.as_mut().await; + } else { + std::future::pending::<()>().await; + } + } +} + +#[cfg(test)] +mod tests { + use std::future::{Future, poll_fn}; + use std::task::Poll; + use std::time::Duration; + + use super::*; + + async fn is_ready(timer: &mut FlushTimer) -> bool { + let wait = timer.wait(); + tokio::pin!(wait); + poll_fn(|cx| Poll::Ready(wait.as_mut().poll(cx).is_ready())).await + } + + #[test] + fn test_create_without_runtime() { + let mut timer = FlushTimer::new(); + timer.set_deadline(None); + assert!(timer.sleep.is_none()); + } + + #[tokio::test(start_paused = true)] + async fn test_deadline_and_cancelled_wait() { + let mut timer = FlushTimer::new(); + assert!(!is_ready(&mut timer).await); + let deadline = Instant::now() + Duration::from_millis(10); + timer.set_deadline(Some(deadline)); + assert!(!is_ready(&mut timer).await); + tokio::time::advance(Duration::from_millis(9)).await; + // Each poll drops its wait future. Repeating the same deadline must not + // move the timer forward or replace its sleep allocation. + let sleep = timer.sleep.as_ref().unwrap().as_ref().get_ref() as *const Sleep; + timer.set_deadline(Some(deadline)); + assert_eq!( + sleep, + timer.sleep.as_ref().unwrap().as_ref().get_ref() as *const Sleep + ); + assert!(!is_ready(&mut timer).await); + tokio::time::advance(Duration::from_millis(1)).await; + assert!(is_ready(&mut timer).await); + } + + #[tokio::test(start_paused = true)] + async fn test_disable_and_rearm_reuses_sleep() { + let mut timer = FlushTimer::new(); + timer.set_deadline(Some(Instant::now())); + assert!(is_ready(&mut timer).await); + let sleep = timer.sleep.as_ref().unwrap().as_ref().get_ref() as *const Sleep; + timer.set_deadline(None); + assert!(!is_ready(&mut timer).await); + assert!(!is_ready(&mut timer).await); + timer.set_deadline(Some(Instant::now() + Duration::from_millis(5))); + assert_eq!( + sleep, + timer.sleep.as_ref().unwrap().as_ref().get_ref() as *const Sleep + ); + assert!(!is_ready(&mut timer).await); + tokio::time::advance(Duration::from_millis(5)).await; + assert!(is_ready(&mut timer).await); + } + + #[tokio::test(start_paused = true)] + async fn test_change_armed_deadline() { + let mut timer = FlushTimer::new(); + let now = Instant::now(); + timer.set_deadline(Some(now + Duration::from_millis(5))); + assert!(!is_ready(&mut timer).await); + timer.set_deadline(Some(now + Duration::from_millis(10))); + tokio::time::advance(Duration::from_millis(5)).await; + assert!(!is_ready(&mut timer).await); + timer.set_deadline(Some(Instant::now())); + assert!(is_ready(&mut timer).await); + } +} diff --git a/src/common/batcher/src/lib.rs b/src/common/batcher/src/lib.rs new file mode 100644 index 0000000000..9ad29097af --- /dev/null +++ b/src/common/batcher/src/lib.rs @@ -0,0 +1,27 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Protocol-independent building blocks for timing-based batching. +//! +//! Payload preparation, grouping keys, write execution and request completion +//! remain the caller's responsibility. + +pub mod flush_limiter; +pub mod flush_policy; +pub mod flush_timer; +pub mod notifier; +pub mod pending_batch; +pub mod pending_worker; +pub mod request_limiter; +pub mod worker_registry; diff --git a/src/common/batcher/src/notifier.rs b/src/common/batcher/src/notifier.rs new file mode 100644 index 0000000000..efc795b938 --- /dev/null +++ b/src/common/batcher/src/notifier.rs @@ -0,0 +1,145 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::future::Future; +use std::num::NonZeroUsize; + +use futures::{StreamExt, stream}; +use tokio::sync::Semaphore; +use tokio::sync::mpsc::error::TrySendError; +use tokio::sync::mpsc::{self, Receiver, Sender}; + +/// Best-effort notification admission, independent of the write result. +/// +/// The caller owns the notification payload, delivery handler and failure metrics. +pub struct Notifier { + sender: Sender, +} + +impl Clone for Notifier { + fn clone(&self) -> Self { + Self { + sender: self.sender.clone(), + } + } +} + +impl Notifier { + /// Returns `None` if the queue capacity is zero or exceeds Tokio's supported capacity. + pub fn try_new(capacity: usize) -> Option<(Self, Receiver)> { + if capacity == 0 || capacity > Semaphore::MAX_PERMITS { + return None; + } + let (sender, receiver) = mpsc::channel(capacity); + Some((Self { sender }, receiver)) + } + + /// Configured queue capacity, independent of its current occupancy. + pub fn max_capacity(&self) -> usize { + self.sender.max_capacity() + } + + /// Never waits for delivery. Returns the payload on a full or closed queue + /// so the caller can record domain-specific diagnostics. + pub fn try_notify(&self, notification: T) -> Result<(), TrySendError> { + self.sender.try_send(notification) + } +} + +/// Delivers queued notifications with bounded concurrency. +/// +/// The caller schedules this future on its runtime. Closing all senders drains +/// queued and in-flight notifications. Handler failures are handled by the caller; +/// this worker does not retry or propagate them into write completion. +pub async fn run_notifier( + receiver: Receiver, + max_concurrent_notifications: NonZeroUsize, + handler: F, +) where + F: FnMut(T) -> Fut, + Fut: Future, +{ + stream::unfold(receiver, |mut receiver| async move { + receiver.recv().await.map(|item| (item, receiver)) + }) + .for_each_concurrent(max_concurrent_notifications.get(), handler) + .await; +} + +#[cfg(test)] +mod tests { + use std::future::{Future, poll_fn}; + use std::num::NonZeroUsize; + use std::sync::Arc; + use std::sync::atomic::{AtomicUsize, Ordering}; + use std::task::Poll; + + use tokio::sync::Semaphore; + use tokio::sync::mpsc::error::TrySendError; + + use crate::notifier::{Notifier, run_notifier}; + + #[test] + fn test_admission() { + assert!(Notifier::::try_new(0).is_none()); + assert!(Notifier::::try_new(usize::MAX).is_none()); + let (notifier, receiver) = Notifier::try_new(1).unwrap(); + let other = notifier.clone(); + assert_eq!(other.max_capacity(), 1); + assert_eq!(notifier.try_notify(1), Ok(())); + assert_eq!(other.try_notify(2), Err(TrySendError::Full(2))); + drop(receiver); + assert_eq!(other.try_notify(3), Err(TrySendError::Closed(3))); + } + + #[tokio::test] + async fn test_concurrency_and_drain() { + // One-factor comparison: only delivery concurrency changes. + for _ in 0..2 { + for concurrency in [1, 2] { + let (notifier, receiver) = Notifier::try_new(3).unwrap(); + for item in 0..3 { + notifier.try_notify(item).unwrap(); + } + drop(notifier); + let started = Arc::new(AtomicUsize::new(0)); + let completed = Arc::new(AtomicUsize::new(0)); + let gate = Arc::new(Semaphore::new(0)); + let worker = + run_notifier(receiver, NonZeroUsize::new(concurrency).unwrap(), |_| { + let started = started.clone(); + let completed = completed.clone(); + let gate = gate.clone(); + async move { + started.fetch_add(1, Ordering::SeqCst); + gate.acquire().await.unwrap().forget(); + completed.fetch_add(1, Ordering::SeqCst); + } + }); + tokio::pin!(worker); + poll_fn(|cx| { + assert!(worker.as_mut().poll(cx).is_pending()); + Poll::Ready(()) + }) + .await; + assert_eq!(started.load(Ordering::SeqCst), concurrency); + assert_eq!(completed.load(Ordering::SeqCst), 0); + gate.add_permits(3); + worker.await; + assert_eq!(started.load(Ordering::SeqCst), 3); + assert_eq!(completed.load(Ordering::SeqCst), 3); + } + } + } +} diff --git a/src/common/batcher/src/pending_batch.rs b/src/common/batcher/src/pending_batch.rs new file mode 100644 index 0000000000..44dc6565fa --- /dev/null +++ b/src/common/batcher/src/pending_batch.rs @@ -0,0 +1,115 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use tokio::time::Instant; + +/// Accumulates complete submissions in arrival order, without interpreting their data. +#[derive(Debug)] +pub struct PendingBatch { + items: Vec, + total_rows: usize, + first_submitted_at: Option, +} + +impl Default for PendingBatch { + fn default() -> Self { + Self { + items: Vec::new(), + total_rows: 0, + first_submitted_at: None, + } + } +} + +impl PendingBatch { + /// Creates an empty batch. + pub fn new() -> Self { + Self::default() + } + + /// Appends a complete submission without moving the first submission's timestamp. + pub fn push(&mut self, item: T, total_rows: usize, now: Instant) { + self.first_submitted_at.get_or_insert(now); + self.items.push(item); + // Saturation preserves threshold decisions even for caller-supplied oversized weights. + self.total_rows = self.total_rows.saturating_add(total_rows); + } + + /// Returns the accumulated row count, saturated at `usize::MAX`. + pub fn total_rows(&self) -> usize { + self.total_rows + } + + /// Returns the number of complete submissions. + pub fn len(&self) -> usize { + self.items.len() + } + + /// Returns whether there are no submissions, including zero-row submissions. + pub fn is_empty(&self) -> bool { + self.items.is_empty() + } + + /// Returns the time the first submission was appended to the current batch. + pub fn first_submitted_at(&self) -> Option { + self.first_submitted_at + } + + /// Takes all submissions in arrival order and resets the batch's accounting. + pub fn take(&mut self) -> Vec { + self.total_rows = 0; + self.first_submitted_at = None; + std::mem::take(&mut self.items) + } +} + +#[cfg(test)] +mod tests { + use std::time::Duration; + + use super::*; + + #[test] + fn test_complete_submissions_and_reset() { + let mut batch = PendingBatch::new(); + let first = Instant::now(); + assert!(batch.is_empty()); + assert_eq!(batch.first_submitted_at(), None); + batch.push(vec![1, 2], 2, first); + batch.push(vec![3, 4, 5], 3, first + Duration::from_secs(1)); + assert_eq!(batch.len(), 2); + assert_eq!(batch.total_rows(), 5); + assert_eq!(batch.first_submitted_at(), Some(first)); + assert_eq!(batch.take(), vec![vec![1, 2], vec![3, 4, 5]]); + assert!(batch.is_empty()); + assert_eq!(batch.total_rows(), 0); + assert_eq!(batch.first_submitted_at(), None); + let next = first + Duration::from_secs(2); + batch.push(vec![6], 1, next); + assert_eq!(batch.first_submitted_at(), Some(next)); + } + + #[test] + fn test_zero_rows_and_saturating_count() { + let mut batch = PendingBatch::new(); + let now = Instant::now(); + batch.push(1, 0, now); + assert!(!batch.is_empty()); + assert_eq!(batch.total_rows(), 0); + batch.push(2, usize::MAX, now); + batch.push(3, 1, now); + assert_eq!(batch.total_rows(), usize::MAX); + assert_eq!(batch.take(), vec![1, 2, 3]); + } +} diff --git a/src/common/batcher/src/pending_worker.rs b/src/common/batcher/src/pending_worker.rs new file mode 100644 index 0000000000..2adda17a5e --- /dev/null +++ b/src/common/batcher/src/pending_worker.rs @@ -0,0 +1,260 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use tokio::time::Instant; + +use crate::flush_policy::{FlushPolicy, FlushTrigger}; +use crate::flush_timer::FlushTimer; +use crate::pending_batch::PendingBatch; + +/// Batch state for one grouping key, independent of the caller's event loop. +/// +/// The caller owns channels, idle/shutdown decisions, execution limits, task +/// spawning and business completion. Taking a batch transfers its items to the +/// caller; no background task is started by this worker. +pub struct PendingWorker { + batch: PendingBatch, + flush_policy: P, + flush_timer: FlushTimer, +} + +impl PendingWorker { + /// Creates an empty worker without starting a timer or runtime task. + pub fn new(flush_policy: P) -> Self { + Self { + batch: PendingBatch::new(), + flush_policy, + flush_timer: FlushTimer::new(), + } + } + + /// Returns whether the worker has no submissions, including zero-row items. + pub fn is_empty(&self) -> bool { + self.batch.is_empty() + } + + /// Returns the pending row count before ownership is transferred to the caller. + pub fn total_rows(&self) -> usize { + self.batch.total_rows() + } + + /// Waits for a timer wakeup without taking or executing the batch. + /// + /// Cancellation is safe: selecting another event does not lose items or + /// reset the deadline. After a wakeup, call [`Self::take_ready`] with + /// [`FlushTrigger::Deadline`]. An empty or unarmed worker stays pending. + pub async fn wait_flush(&mut self) { + self.flush_timer.wait().await; + } + + /// Takes all pending submissions regardless of policy, for example on shutdown. + /// + /// Returns `None` for an empty batch. The caller decides whether to execute + /// inline, acquire a flush permit or discard the returned items. + pub fn take_pending(&mut self) -> Option> { + self.flush_timer.set_deadline(None); + if self.batch.is_empty() { + None + } else { + Some(self.batch.take()) + } + } +} + +impl PendingWorker { + /// Appends one complete submission and arms the policy's deadline. + /// + /// Call [`Self::take_ready`] with [`FlushTrigger::Submission`] afterwards to + /// check whether this submission triggered a flush. Timer-backed policies + /// require a Tokio runtime with time enabled. + /// + /// # Panics + /// + /// Panics if arming a new timer outside a Tokio runtime with time enabled. + pub fn submit(&mut self, item: T, total_rows: usize) { + self.batch.push(item, total_rows, Instant::now()); + self.refresh_deadline(); + } + + /// Takes the batch only when the policy permits flushing for this event. + /// + /// Like [`Self::submit`], this may arm a timer and requires a time-enabled + /// Tokio runtime when the policy starts using deadlines. + pub fn take_ready(&mut self, trigger: FlushTrigger) -> Option> { + if !self.batch.is_empty() + && self + .flush_policy + .should_flush(&self.batch, Instant::now(), trigger) + { + self.take_pending() + } else { + self.refresh_deadline(); + None + } + } + + fn refresh_deadline(&mut self) { + let deadline = if self.batch.is_empty() { + None + } else { + self.flush_policy.deadline(&self.batch) + }; + self.flush_timer.set_deadline(deadline); + } +} + +#[cfg(test)] +mod tests { + use std::future::Future; + use std::task::Poll; + use std::time::Duration; + + use super::*; + use crate::flush_policy::timing::TimingFlushPolicy; + + fn worker(rows: usize) -> PendingWorker { + PendingWorker::new(TimingFlushPolicy::try_new(Duration::from_millis(10), rows).unwrap()) + } + + async fn assert_wait_pending(worker: &mut PendingWorker) { + let wait = worker.wait_flush(); + tokio::pin!(wait); + assert!(std::future::poll_fn(|cx| Poll::Ready(wait.as_mut().poll(cx).is_pending())).await); + } + + #[tokio::test(start_paused = true)] + async fn test_first_deadline_and_cancelled_wait() { + let mut worker = worker(100); + let start = Instant::now(); + worker.submit(1, 1); + assert_wait_pending(&mut worker).await; + tokio::time::advance(Duration::from_millis(5)).await; + worker.submit(2, 1); + assert!(worker.take_ready(FlushTrigger::Submission).is_none()); + assert_wait_pending(&mut worker).await; + worker.wait_flush().await; + assert_eq!(start + Duration::from_millis(10), Instant::now()); + assert_eq!(Some(vec![1, 2]), worker.take_ready(FlushTrigger::Deadline)); + assert!(worker.is_empty()); + assert_eq!(0, worker.total_rows()); + assert_wait_pending(&mut worker).await; + } + + #[tokio::test(start_paused = true)] + async fn test_submission_does_not_consume_expired_deadline() { + let mut worker = worker(100); + worker.submit(1, 1); + tokio::time::advance(Duration::from_millis(10)).await; + worker.submit(2, 1); + assert!(worker.take_ready(FlushTrigger::Submission).is_none()); + worker.wait_flush().await; + assert_eq!(Some(vec![1, 2]), worker.take_ready(FlushTrigger::Deadline)); + } + + #[tokio::test(start_paused = true)] + async fn test_size_flush_and_rearm() { + let mut worker = worker(2); + worker.submit(1, 3); + assert_eq!(3, worker.total_rows()); + assert_eq!(Some(vec![1]), worker.take_ready(FlushTrigger::Submission)); + tokio::time::advance(Duration::from_secs(1)).await; + assert_wait_pending(&mut worker).await; + let start = Instant::now(); + worker.submit(2, 1); + worker.wait_flush().await; + assert_eq!(start + Duration::from_millis(10), Instant::now()); + assert_eq!(Some(vec![2]), worker.take_ready(FlushTrigger::Deadline)); + } + + #[tokio::test(start_paused = true)] + async fn test_take_pending_without_waiting() { + let mut worker = worker(100); + assert_eq!(None, worker.take_pending()); + worker.submit(1, 0); + assert!(!worker.is_empty()); + let now = Instant::now(); + assert_eq!(Some(vec![1]), worker.take_pending()); + assert_eq!(now, Instant::now()); + tokio::time::advance(Duration::from_secs(1)).await; + assert_wait_pending(&mut worker).await; + } + + #[tokio::test(start_paused = true)] + async fn test_row_threshold_ablation() { + for (rows, early) in [(2, true), (100, false)] { + let mut worker = worker(rows); + let start = Instant::now(); + worker.submit(1, 1); + tokio::time::advance(Duration::from_millis(5)).await; + worker.submit(2, 1); + let batch = worker.take_ready(FlushTrigger::Submission); + if early { + assert_eq!(Some(vec![1, 2]), batch); + assert_eq!(start + Duration::from_millis(5), Instant::now()); + } else { + assert!(batch.is_none()); + worker.wait_flush().await; + assert_eq!(start + Duration::from_millis(10), Instant::now()); + assert_eq!(Some(vec![1, 2]), worker.take_ready(FlushTrigger::Deadline)); + } + } + } + + #[tokio::test] + async fn test_non_timing_policy() { + struct TwoSubmissionsPolicy; + impl FlushPolicy for TwoSubmissionsPolicy { + fn deadline(&self, _: &PendingBatch) -> Option { + None + } + fn should_flush( + &self, + batch: &PendingBatch, + _: Instant, + _: FlushTrigger, + ) -> bool { + batch.len() >= 2 + } + } + let mut worker = PendingWorker::new(TwoSubmissionsPolicy); + worker.submit(1, 0); + assert_wait_pending(&mut worker).await; + assert!(worker.take_ready(FlushTrigger::Submission).is_none()); + worker.submit(2, 0); + assert_eq!( + Some(vec![1, 2]), + worker.take_ready(FlushTrigger::Submission) + ); + } + + #[tokio::test(start_paused = true)] + async fn test_caller_owns_event_loop() { + let mut worker = worker(100); + let (tx, mut rx) = tokio::sync::mpsc::channel(1); + tx.send(1).await.unwrap(); + tokio::select! { + _ = worker.wait_flush() => panic!("empty worker must not wake"), + item = rx.recv() => worker.submit(item.unwrap(), 1), + } + assert!(worker.take_ready(FlushTrigger::Submission).is_none()); + let start = Instant::now(); + tokio::select! { + _ = worker.wait_flush() => { + assert_eq!(Some(vec![1]), worker.take_ready(FlushTrigger::Deadline)); + } + _ = rx.recv() => panic!("sender is still open without another submission"), + } + assert_eq!(start + Duration::from_millis(10), Instant::now()); + } +} diff --git a/src/common/batcher/src/request_limiter.rs b/src/common/batcher/src/request_limiter.rs new file mode 100644 index 0000000000..c488bec02d --- /dev/null +++ b/src/common/batcher/src/request_limiter.rs @@ -0,0 +1,115 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::sync::Arc; + +use tokio::sync::{AcquireError, OwnedSemaphorePermit, Semaphore}; + +/// Limits unfinished original requests independently of their number of submissions. +/// +/// Acquire once for each original request, then clone its permit into all of that +/// request's submissions. Capacity is released when the last owner drops it. +#[derive(Clone)] +pub struct RequestLimiter { + semaphore: Arc, +} + +impl RequestLimiter { + /// Returns `None` if the limit is zero or exceeds Tokio's supported semaphore capacity. + pub fn try_new(max_inflight_requests: usize) -> Option { + let permits = max_inflight_requests; + ((1..=Semaphore::MAX_PERMITS).contains(&permits)).then(|| Self { + semaphore: Arc::new(Semaphore::new(permits)), + }) + } + + /// Waits for one request slot and returns a permit shared by its submissions. + /// + /// Cancelling a pending acquisition does not consume a slot. Once acquired, + /// callers must retain the permit until every submission has finished, even + /// if the original caller stops waiting for its response. + pub async fn acquire(&self) -> Result, AcquireError> { + self.semaphore.clone().acquire_owned().await.map(Arc::new) + } +} + +#[cfg(test)] +mod tests { + use std::future::{Future, poll_fn}; + use std::pin::Pin; + use std::task::Poll; + + use tokio::sync::Semaphore; + + use crate::request_limiter::RequestLimiter; + + async fn is_pending(mut future: Pin<&mut F>) -> bool { + poll_fn(|cx| Poll::Ready(future.as_mut().poll(cx).is_pending())).await + } + + #[test] + fn test_capacity_boundaries() { + for capacity in [1, Semaphore::MAX_PERMITS] { + assert!(RequestLimiter::try_new(capacity).is_some()); + } + for capacity in [0, Semaphore::MAX_PERMITS + 1, usize::MAX] { + assert!(RequestLimiter::try_new(capacity).is_none()); + } + } + + #[tokio::test] + async fn test_last_submission_releases_request_slot() { + let limiter = RequestLimiter::try_new(1).unwrap(); + let other = limiter.clone(); + let request = limiter.acquire().await.unwrap(); + let first_submission = request.clone(); + let second_submission = request.clone(); + let mut waiting = Box::pin(other.acquire()); + assert!(is_pending(waiting.as_mut()).await); + // The caller and one submission finish, but another submission remains. + drop(request); + drop(first_submission); + assert!(is_pending(waiting.as_mut()).await); + drop(second_submission); + let next = waiting.await.unwrap(); + assert_eq!(next.num_permits(), 1); + } + + #[tokio::test] + async fn test_cancelled_acquisition_does_not_leak_capacity() { + let limiter = RequestLimiter::try_new(1).unwrap(); + let permit = limiter.acquire().await.unwrap(); + let mut cancelled = Box::pin(limiter.acquire()); + assert!(is_pending(cancelled.as_mut()).await); + drop(cancelled); + let mut next = Box::pin(limiter.acquire()); + assert!(is_pending(next.as_mut()).await); + drop(permit); + let next = next.await.unwrap(); + drop(next); + let _available = limiter.acquire().await.unwrap(); + } + + #[tokio::test] + async fn test_independent_requests_consume_separate_slots() { + let limiter = RequestLimiter::try_new(2).unwrap(); + let first = limiter.acquire().await.unwrap(); + let second = limiter.acquire().await.unwrap(); + let mut third = Box::pin(limiter.acquire()); + assert!(is_pending(third.as_mut()).await); + drop(first); + let _third = third.await.unwrap(); + drop(second); + } +} diff --git a/src/common/batcher/src/worker_registry.rs b/src/common/batcher/src/worker_registry.rs new file mode 100644 index 0000000000..f5dbf48e48 --- /dev/null +++ b/src/common/batcher/src/worker_registry.rs @@ -0,0 +1,265 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::hash::Hash; + +use dashmap::DashMap; +use dashmap::mapref::entry::Entry; +use tokio::sync::mpsc::{self, Receiver, Sender}; + +/// Registers worker senders by key, without owning workers or their execution. +/// +/// Closed senders are replaced on lookup/creation. Cleanup checks channel +/// identity under the same lock as replacement, so an old worker cannot remove +/// a replacement registered under its key. +pub struct WorkerRegistry { + workers: DashMap>, +} + +impl Default for WorkerRegistry { + fn default() -> Self { + Self { + workers: DashMap::new(), + } + } +} + +impl WorkerRegistry { + pub fn new() -> Self { + Self::default() + } + + /// Returns a live sender, if one is currently registered. + /// The receiver can close after this method returns; callers must handle a + /// failed send and retry worker lookup without discarding the unsent item. + pub async fn get(&self, key: &K) -> Option> { + self.workers + .get(key) + .filter(|tx| !tx.is_closed()) + .map(|tx| tx.value().clone()) + } + + /// Returns the registered sender and, only when created, its receiver. + /// + /// Start the new worker before the next await so cancellation cannot leave + /// a registered channel without a consumer. Existing channels retain their + /// original capacity. New channel capacity must satisfy `mpsc::channel`. + pub async fn get_or_create(&self, key: K, capacity: usize) -> (Sender, Option>) { + let mut receiver = None; + let sender = self + .get_or_insert_with(key, || { + let (sender, rx) = mpsc::channel(capacity); + receiver = Some(rx); + sender + }) + .await; + (sender, receiver) + } + + /// Reuses a live sender or atomically creates its replacement. + /// + /// `create` runs synchronously under the registry shard lock. It should only + /// prepare the sender and capture any initialization state (such as the + /// receiver) for the caller; start the worker after this method returns. + /// Do not block or reenter the registry from `create`. + pub async fn get_or_insert_with(&self, key: K, create: F) -> Sender + 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() { + entry.insert(create()); + } + entry.get().clone() + } + Entry::Vacant(entry) => entry.insert(create()).value().clone(), + } + } + + /// Removes the key only if it still points to this worker's channel. + pub async fn remove_if_same(&self, key: &K, tx: &Sender) -> bool { + self.workers + .remove_if(key, |_, current| current.same_channel(tx)) + .is_some() + } + + /// Number of registered entries, including senders whose receivers closed. + pub async fn len(&self) -> usize { + self.workers.len() + } + + pub async fn is_empty(&self) -> bool { + self.workers.is_empty() + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + use std::sync::atomic::{AtomicUsize, Ordering}; + + use tokio::sync::{Barrier, mpsc}; + + 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(); + assert!(registry.is_empty().await); + assert!(registry.get(&1).await.is_none()); + let mut receiver = None; + let first = registry + .get_or_insert_with(1, || { + let (tx, rx) = mpsc::channel(1); + receiver = Some(rx); + tx + }) + .await; + let second = registry + .get_or_insert_with(1, || panic!("live worker must be reused")) + .await; + assert!(first.same_channel(&second)); + assert!(registry.get(&1).await.unwrap().same_channel(&first)); + first.send(7).await.unwrap(); + assert_eq!(Some(7), receiver.as_mut().unwrap().recv().await); + assert_eq!(1, registry.len().await); + assert!(registry.remove_if_same(&1, &first).await); + assert!(!registry.remove_if_same(&1, &first).await); + assert!(registry.is_empty().await); + } + + #[tokio::test] + async fn test_old_cleanup_does_not_remove_replacement() { + let registry = WorkerRegistry::<_, ()>::new(); + let (first, mut receiver) = mpsc::channel(1); + registry.get_or_insert_with("table", || first.clone()).await; + receiver.close(); + assert!(registry.get(&"table").await.is_none()); + // Closed entries remain accounted until replacement or explicit removal. + assert_eq!(1, registry.len().await); + let (second, _receiver) = mpsc::channel(1); + registry + .get_or_insert_with("table", || second.clone()) + .await; + assert!(!registry.remove_if_same(&"table", &first).await); + assert!(registry.get(&"table").await.unwrap().same_channel(&second)); + assert_eq!(1, registry.len().await); + assert!(registry.remove_if_same(&"table", &second).await); + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn test_concurrent_lookup_initializes_once() { + let registry = Arc::new(WorkerRegistry::<_, ()>::new()); + let barrier = Arc::new(Barrier::new(3)); + let count = Arc::new(AtomicUsize::new(0)); + let mut tasks = Vec::new(); + for _ in 0..2 { + let registry = registry.clone(); + let barrier = barrier.clone(); + let count = count.clone(); + tasks.push(tokio::spawn(async move { + barrier.wait().await; + let mut receiver = None; + let tx = registry + .get_or_insert_with(1, || { + count.fetch_add(1, Ordering::Relaxed); + let (tx, rx) = mpsc::channel(1); + receiver = Some(rx); + tx + }) + .await; + // Keep the newly created receiver alive while both lookups run. + barrier.wait().await; + (tx, receiver) + })); + } + barrier.wait().await; + barrier.wait().await; + let first = tasks.remove(0).await.unwrap(); + let second = tasks.remove(0).await.unwrap(); + assert!(first.0.same_channel(&second.0)); + assert_eq!(1, count.load(Ordering::Relaxed)); + assert_eq!(1, registry.len().await); + } + #[tokio::test] + async fn test_channel_creation_replacement_and_cleanup() { + let registry = WorkerRegistry::<_, usize>::new(); + let (first, receiver) = registry.get_or_create(1, 2).await; + let mut receiver = receiver.unwrap(); + let (reused, absent) = registry.get_or_create(1, 3).await; + assert!(absent.is_none()); + assert!(first.same_channel(&reused)); + assert_eq!(2, reused.max_capacity()); + first.send(7).await.unwrap(); + assert_eq!(Some(7), receiver.recv().await); + receiver.close(); + let (replacement, receiver) = registry.get_or_create(1, 3).await; + assert!(receiver.is_some()); + assert_eq!(3, replacement.max_capacity()); + assert!(!first.same_channel(&replacement)); + assert!(!registry.remove_if_same(&1, &first).await); + assert!(registry.get(&1).await.unwrap().same_channel(&replacement)); + assert!(registry.remove_if_same(&1, &replacement).await); + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn test_concurrent_channel_creation_returns_one_receiver() { + let registry = Arc::new(WorkerRegistry::<_, ()>::new()); + let barrier = Arc::new(Barrier::new(3)); + let mut tasks = Vec::new(); + for _ in 0..2 { + let registry = registry.clone(); + let barrier = barrier.clone(); + tasks.push(tokio::spawn(async move { + barrier.wait().await; + let channel = registry.get_or_create(1, 2).await; + barrier.wait().await; + channel + })); + } + barrier.wait().await; + barrier.wait().await; + let first = tasks.remove(0).await.unwrap(); + let second = tasks.remove(0).await.unwrap(); + assert!(first.0.same_channel(&second.0)); + assert_ne!(first.1.is_some(), second.1.is_some()); + assert_eq!(1, registry.len().await); + } +} diff --git a/src/common/grpc/src/flight.rs b/src/common/grpc/src/flight.rs index 1245a86efb..4f5a47cb53 100644 --- a/src/common/grpc/src/flight.rs +++ b/src/common/grpc/src/flight.rs @@ -37,6 +37,35 @@ use vec1::{Vec1, vec1}; use crate::error; use crate::error::{DecodeFlightDataSnafu, InvalidFlightDataSnafu, Result}; +/// Encodes a batch into the schema, record-batch header and payload used by bulk inserts. +/// +/// Uses the default Flight compression and rejects dictionary batches, whose +/// additional Flight messages cannot be represented by this three-part format. +pub fn record_batch_to_ipc( + record_batch: DfRecordBatch, +) -> Result<(ProstBytes, ProstBytes, ProstBytes)> { + let mut encoder = FlightEncoder::default(); + let schema = encoder.encode_schema(record_batch.schema().as_ref()); + let mut iter = encoder + .encode(FlightMessage::RecordBatch(record_batch)) + .into_iter(); + let flight_data = iter.next().context(InvalidFlightDataSnafu { + reason: "Failed to encode empty flight data", + })?; + if iter.next().is_some() { + return error::NotSupportedSnafu { + feat: "bulk insert RecordBatch with dictionary arrays", + } + .fail(); + } + + Ok(( + schema.data_header, + flight_data.data_header, + flight_data.data_body, + )) +} + /// Flight metadata key used to carry flow query extensions as JSON pairs. pub const FLOW_EXTENSIONS_METADATA_KEY: &str = "x-greptime-flow-extensions"; /// Flight metadata key used to carry query snapshot read upper bounds as JSON. @@ -374,8 +403,52 @@ mod test { use datatypes::arrow::buffer::OffsetBuffer; use datatypes::arrow::datatypes::{DataType, Field, Schema}; - use super::*; use crate::Error; + use crate::flight::*; + + #[test] + fn test_record_batch_to_ipc_preserves_wire_bytes() { + let schema = Arc::new(Schema::new(vec![Field::new("n", DataType::Int32, true)])); + for values in [vec![], vec![Some(1), None, Some(3)]] { + let batch = + DfRecordBatch::try_new(schema.clone(), vec![Arc::new(Int32Array::from(values))]) + .unwrap(); + // Reproduce the previous Prom encoding sequence and compare all bytes. + let mut encoder = FlightEncoder::default(); + let expected_schema = encoder.encode_schema(batch.schema().as_ref()); + let messages = encoder.encode(FlightMessage::RecordBatch(batch.clone())); + assert_eq!(messages.len(), 1); + let expected_batch = messages.first(); + assert_eq!( + record_batch_to_ipc(batch).unwrap(), + ( + expected_schema.data_header, + expected_batch.data_header.clone(), + expected_batch.data_body.clone(), + ) + ); + } + } + + #[test] + fn test_record_batch_to_ipc_rejects_dictionary() { + let schema = Arc::new(Schema::new(vec![Field::new_dictionary( + "tag", + DataType::UInt32, + DataType::Utf8, + true, + )])); + let dictionary = DictionaryArray::new( + UInt32Array::from_value(0, 3), + Arc::new(StringArray::from_iter_values(["x"])), + ); + let batch = DfRecordBatch::try_new(schema, vec![Arc::new(dictionary)]).unwrap(); + assert!(matches!( + record_batch_to_ipc(batch), + Err(Error::NotSupported { feat }) + if feat == "bulk insert RecordBatch with dictionary arrays" + )); + } #[test] fn test_try_decode() -> Result<()> { diff --git a/src/datatypes/src/timestamp.rs b/src/datatypes/src/timestamp.rs index bf6fdaa162..ac6ca1444f 100644 --- a/src/datatypes/src/timestamp.rs +++ b/src/datatypes/src/timestamp.rs @@ -13,7 +13,7 @@ // limitations under the License. use arrow_array::{ - ArrayRef, PrimitiveArray, TimestampMicrosecondArray, TimestampMillisecondArray, + Array, ArrayRef, PrimitiveArray, TimestampMicrosecondArray, TimestampMillisecondArray, TimestampNanosecondArray, TimestampSecondArray, }; use arrow_schema::DataType; @@ -179,8 +179,24 @@ pub fn timestamp_array_to_primitive( Some((ts_primitive, *unit)) } +/// Appends non-null timestamps in the source array's native time unit. +/// +/// Returns `None` for a non-timestamp array without changing `timestamps`. +pub fn append_timestamps(ts_array: &ArrayRef, timestamps: &mut Vec) -> Option<()> { + let (values, _) = timestamp_array_to_primitive(ts_array)?; + if values.null_count() == 0 { + timestamps.extend_from_slice(values.values()); + } else { + timestamps.extend(values.iter().flatten()); + } + Some(()) +} + #[cfg(test)] mod tests { + use std::sync::Arc; + + use arrow_array::Int64Array; use common_time::timezone::set_default_timezone; use super::*; @@ -213,4 +229,46 @@ mod tests { assert_eq!(ts, ts.as_scalar_ref()); assert_eq!(ts, ts.to_owned_scalar()); } + + #[test] + fn test_append_timestamps() { + let cases = [ + vec![Some(i64::MIN), Some(-1), Some(0), Some(i64::MAX)], + vec![Some(-1), None, Some(2), None], + vec![None, None], + vec![], + ]; + for values in cases { + let arrays: [ArrayRef; 4] = [ + Arc::new(TimestampSecondArray::from(values.clone())), + Arc::new(TimestampMillisecondArray::from(values.clone())), + Arc::new(TimestampMicrosecondArray::from(values.clone())), + Arc::new(TimestampNanosecondArray::from(values.clone())), + ]; + let mut expected = vec![42]; + expected.extend(values.iter().flatten().copied()); + for array in arrays { + for _ in 0..2 { + let (primitive, _) = timestamp_array_to_primitive(&array).unwrap(); + let mut reference = vec![42]; + reference.extend(primitive.iter().flatten()); + + let mut timestamps = vec![42]; + assert_eq!(append_timestamps(&array, &mut timestamps), Some(())); + assert_eq!(timestamps, expected); + assert_eq!(timestamps, reference); + } + } + } + } + + #[test] + fn test_append_timestamps_invalid_array_preserves_prefix() { + for values in [vec![], vec![Some(1), None]] { + let array: ArrayRef = Arc::new(Int64Array::from(values)); + let mut timestamps = vec![42, -1]; + assert_eq!(append_timestamps(&array, &mut timestamps), None); + assert_eq!(timestamps, vec![42, -1]); + } + } } 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/operator/src/bulk_insert.rs b/src/operator/src/bulk_insert.rs index e35779ec97..c47236d7db 100644 --- a/src/operator/src/bulk_insert.rs +++ b/src/operator/src/bulk_insert.rs @@ -20,7 +20,6 @@ use api::v1::region::{ BulkInsertRequest, RegionRequest, RegionRequestHeader, bulk_insert_request, region_request, }; use api::v1::{ArrowIpc, PartitionExprVersion}; -use arrow::array::Array; use arrow::record_batch::RecordBatch; use bytes::Bytes; use common_base::AffectedRows; @@ -28,12 +27,13 @@ use common_grpc::FlightData; use common_grpc::flight::{FlightEncoder, FlightMessage}; use common_telemetry::error; use common_telemetry::tracing_context::TracingContext; -use snafu::{OptionExt, ResultExt, ensure}; +use snafu::{ResultExt, ensure}; use store_api::storage::RegionId; use table::TableRef; use table::metadata::TableInfoRef; use crate::insert::Inserter; +use crate::req_convert::insert::extract_timestamps; use crate::{error, metrics}; impl Inserter { @@ -319,22 +319,3 @@ impl Inserter { }); } } - -/// Calculate the timestamp range of record batch. Return `None` if record batch is empty. -fn extract_timestamps(rb: &RecordBatch, timestamp_index_name: &str) -> error::Result> { - let ts_col = rb - .column_by_name(timestamp_index_name) - .context(error::ColumnNotFoundSnafu { - msg: timestamp_index_name, - })?; - if rb.num_rows() == 0 { - return Ok(vec![]); - } - let (primitive, _) = - datatypes::timestamp::timestamp_array_to_primitive(ts_col).with_context(|| { - error::InvalidTimeIndexTypeSnafu { - ty: ts_col.data_type().clone(), - } - })?; - Ok(primitive.iter().flatten().collect()) -} diff --git a/src/operator/src/req_convert/insert.rs b/src/operator/src/req_convert/insert.rs index f9bf83d926..72d6693c6b 100644 --- a/src/operator/src/req_convert/insert.rs +++ b/src/operator/src/req_convert/insert.rs @@ -17,6 +17,7 @@ mod fill_impure_default; mod row_to_region; mod stmt_to_region; mod table_to_region; +mod timestamps; use api::v1::SemanticType; pub use column_to_row::ColumnToRow; @@ -26,6 +27,7 @@ use snafu::{OptionExt, ResultExt}; pub use stmt_to_region::StatementToRegion; use table::metadata::TableInfo; pub use table_to_region::TableToRegion; +pub use timestamps::extract_timestamps; use crate::error::{ColumnNotFoundSnafu, MissingTimeIndexColumnSnafu, Result}; diff --git a/src/operator/src/req_convert/insert/timestamps.rs b/src/operator/src/req_convert/insert/timestamps.rs new file mode 100644 index 0000000000..f0315135f5 --- /dev/null +++ b/src/operator/src/req_convert/insert/timestamps.rs @@ -0,0 +1,38 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use arrow::record_batch::RecordBatch; +use datatypes::timestamp::append_timestamps; +use snafu::OptionExt; + +use crate::error::{self, Result}; + +/// Extracts non-null timestamps in the source column's native time unit. +pub fn extract_timestamps(rb: &RecordBatch, timestamp_index_name: &str) -> Result> { + let ts_col = rb + .column_by_name(timestamp_index_name) + .context(error::ColumnNotFoundSnafu { + msg: timestamp_index_name, + })?; + if rb.num_rows() == 0 { + return Ok(vec![]); + } + let mut timestamps = Vec::with_capacity(rb.num_rows()); + append_timestamps(ts_col, &mut timestamps).with_context(|| { + error::InvalidTimeIndexTypeSnafu { + ty: ts_col.data_type().clone(), + } + })?; + Ok(timestamps) +} diff --git a/src/servers/Cargo.toml b/src/servers/Cargo.toml index 9d5bc257fd..7eb008369b 100644 --- a/src/servers/Cargo.toml +++ b/src/servers/Cargo.toml @@ -36,6 +36,7 @@ bytes.workspace = true catalog.workspace = true chrono.workspace = true common-base.workspace = true +common-batcher.workspace = true common-catalog.workspace = true common-decimal.workspace = true common-error.workspace = true diff --git a/src/servers/src/pending_rows_batcher.rs b/src/servers/src/pending_rows_batcher.rs index f8c4a869d5..79ae645b98 100644 --- a/src/servers/src/pending_rows_batcher.rs +++ b/src/servers/src/pending_rows_batcher.rs @@ -23,22 +23,27 @@ use api::v1::region::{ BulkInsertRequest, RegionRequest, RegionRequestHeader, bulk_insert_request, region_request, }; use api::v1::{ArrowIpc, ColumnSchema, RowInsertRequests, Rows}; -use arrow::array::Array; use arrow::compute::{concat_batches, filter_record_batch}; use arrow::datatypes::{DataType as ArrowDataType, Schema as ArrowSchema, TimeUnit}; use arrow::record_batch::RecordBatch; use async_trait::async_trait; -use bytes::Bytes; use catalog::CatalogManagerRef; -use common_grpc::flight::{FlightEncoder, FlightMessage}; +use common_batcher::flush_limiter::FlushLimiter; +use common_batcher::flush_policy::FlushTrigger; +use common_batcher::flush_policy::timing::TimingFlushPolicy; +use common_batcher::notifier::{Notifier, run_notifier}; +use common_batcher::pending_worker::PendingWorker as PendingCore; +use common_batcher::request_limiter::RequestLimiter; +use common_batcher::worker_registry::WorkerRegistry; +use common_grpc::error::Error as GrpcError; +use common_grpc::flight::record_batch_to_ipc; use common_meta::cache::TableFlownodeSetCacheRef; use common_meta::node_manager::NodeManagerRef; use common_query::prelude::{GREPTIME_PHYSICAL_TABLE, greptime_timestamp, greptime_value}; +use common_runtime::spawn_global; use common_telemetry::tracing_context::TracingContext; use common_telemetry::{debug, error, warn}; -use dashmap::DashMap; -use dashmap::mapref::entry::Entry; -use futures::StreamExt; +use datatypes::timestamp::append_timestamps; use metric_engine::batch_modifier::{TagColumnInfo, modify_batch_sparse}; use partition::manager::PartitionRuleManagerRef; use partition::partition::PartitionRuleRef; @@ -81,7 +86,7 @@ pub fn pending_rows_batch_sync_enabled() -> bool { } const WORKER_IDLE_TIMEOUT_MULTIPLIER: u32 = 3; const PHYSICAL_REGION_ESSENTIAL_COLUMN_COUNT: usize = 3; -const MAX_CONCURRENT_FLOW_NOTIFICATIONS: usize = 8; +const MAX_CONCURRENT_FLOW_NOTIFICATIONS: NonZeroUsize = NonZeroUsize::new(8).unwrap(); #[async_trait] pub trait PendingRowsSchemaAlterer: Send + Sync { /// Batch-create multiple logical tables that are missing. @@ -292,15 +297,13 @@ struct TableResolutionPlan { struct PendingBatch { tables: HashMap, - total_row_count: usize, db_string: String, ctx: QueryContextRef, - waiters: Vec, } struct FlushWaiter { response_tx: oneshot::Sender>>, - _permit: OwnedSemaphorePermit, + _permit: Arc, } struct FlushBatch { @@ -322,7 +325,7 @@ enum WorkerCommand { total_rows: usize, ctx: QueryContextRef, response_tx: oneshot::Sender>>, - _permit: OwnedSemaphorePermit, + _permit: Arc, }, #[cfg(test)] Ack { ack_tx: oneshot::Sender<()> }, @@ -344,15 +347,15 @@ fn batch_key_from_ctx(ctx: &QueryContextRef) -> BatchKey { /// Prometheus remote write pending rows batcher. pub struct PendingRowsBatcher { - workers: Arc>, + workers: Arc>, flush_interval: Duration, - max_batch_rows: usize, + flush_policy: TimingFlushPolicy, partition_manager: PartitionRuleManagerRef, node_manager: NodeManagerRef, catalog_manager: CatalogManagerRef, - flow_notification_tx: mpsc::Sender, - flush_semaphore: Arc, - inflight_semaphore: Arc, + flow_notification_tx: Notifier, + flush_limiter: FlushLimiter, + request_limiter: RequestLimiter, worker_channel_capacity: usize, prom_store_with_metric_engine: bool, schema_alterer: PendingRowsSchemaAltererRef, @@ -376,24 +379,21 @@ 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; } + let flush_policy = TimingFlushPolicy::try_new(flush_interval, max_batch_rows)?; + let flush_limiter = FlushLimiter::try_new(max_concurrent_flushes)?; + + let request_limiter = RequestLimiter::try_new(max_inflight_requests)?; + let (flow_notification_tx, flow_notification_rx) = + Notifier::try_new(flow_notification_queue_capacity.get())?; + let (shutdown, _) = broadcast::channel(1); let pending_rows_batch_sync = pending_rows_batch_sync_enabled(); - let workers = Arc::new(DashMap::new()); - PENDING_WORKERS.set(workers.len() as i64); - let (flow_notification_tx, flow_notification_rx) = - mpsc::channel(flow_notification_queue_capacity.get()); + let workers = Arc::new(WorkerRegistry::new()); + PENDING_WORKERS.set(0); start_flow_notification_worker( flow_notification_rx, table_flownode_set_cache, @@ -403,15 +403,15 @@ impl PendingRowsBatcher { Some(Arc::new(Self { workers, flush_interval, - max_batch_rows, + flush_policy, partition_manager, node_manager, catalog_manager, flow_notification_tx, prom_store_with_metric_engine, schema_alterer, - flush_semaphore: Arc::new(Semaphore::new(max_concurrent_flushes)), - inflight_semaphore: Arc::new(Semaphore::new(max_inflight_requests)), + flush_limiter, + request_limiter, worker_channel_capacity, pending_rows_batch_sync, shutdown, @@ -433,9 +433,8 @@ impl PendingRowsBatcher { let _timer = PENDING_ROWS_BATCH_INGEST_STAGE_ELAPSED .with_label_values(&["submit_acquire_inflight_permit"]) .start_timer(); - self.inflight_semaphore - .clone() - .acquire_owned() + self.request_limiter + .acquire() .await .map_err(|_| error::BatcherChannelClosedSnafu.build())? }; @@ -457,7 +456,7 @@ impl PendingRowsBatcher { .start_timer(); for _ in 0..2 { - let worker = self.get_or_spawn_worker(batch_key.clone()); + let worker = self.get_or_spawn_worker(batch_key.clone()).await; let Some(worker_cmd) = cmd.take() else { break; }; @@ -470,7 +469,8 @@ impl PendingRowsBatcher { self.workers.as_ref(), &batch_key, &worker.tx, - ); + ) + .await; } } } @@ -832,39 +832,24 @@ impl PendingRowsBatcher { Ok(aligned_batches) } - - fn get_or_spawn_worker(&self, key: BatchKey) -> PendingWorker { - if let Some(worker) = self.workers.get(&key) - && !worker.tx.is_closed() - { - return worker.clone(); - } - - let entry = self.workers.entry(key.clone()); - match entry { - Entry::Occupied(mut worker) => { - if worker.get().tx.is_closed() { - let new_worker = self.spawn_worker(key); - worker.insert(new_worker.clone()); - PENDING_WORKERS.set(self.workers.len() as i64); - new_worker - } else { - worker.get().clone() - } - } - Entry::Vacant(vacant) => { - let worker = self.spawn_worker(key); - - vacant.insert(worker.clone()); - PENDING_WORKERS.set(self.workers.len() as i64); - worker - } + async fn get_or_spawn_worker(&self, key: BatchKey) -> PendingWorker { + let (tx, receiver) = self + .workers + .get_or_create(key.clone(), self.worker_channel_capacity) + .await; + if let Some(rx) = receiver { + self.spawn_worker(key, tx.clone(), rx); + PENDING_WORKERS.set(self.workers.len().await as i64); } + PendingWorker { tx } } - fn spawn_worker(&self, key: BatchKey) -> PendingWorker { - let (tx, rx) = mpsc::channel(self.worker_channel_capacity); - let worker = PendingWorker { tx: tx.clone() }; + fn spawn_worker( + &self, + key: BatchKey, + tx: mpsc::Sender, + rx: mpsc::Receiver, + ) { let worker_idle_timeout = self .flush_interval .checked_mul(WORKER_IDLE_TIMEOUT_MULTIPLIER) @@ -872,7 +857,7 @@ impl PendingRowsBatcher { start_worker( key, - worker.tx.clone(), + tx, self.workers.clone(), rx, self.shutdown.clone(), @@ -880,13 +865,10 @@ impl PendingRowsBatcher { self.node_manager.clone(), self.catalog_manager.clone(), self.flow_notification_tx.clone(), - self.flush_interval, worker_idle_timeout, - self.max_batch_rows, - self.flush_semaphore.clone(), + self.flush_policy, + self.flush_limiter.clone(), ); - - worker } } @@ -901,10 +883,8 @@ impl PendingBatch { let db_string = ctx.get_db_string(); Self { tables: HashMap::new(), - total_row_count: 0, db_string, ctx, - waiters: Vec::new(), } } @@ -929,22 +909,21 @@ impl PendingBatch { fn start_worker( key: BatchKey, worker_tx: mpsc::Sender, - workers: Arc>, + workers: Arc>, mut rx: mpsc::Receiver, shutdown: broadcast::Sender<()>, partition_manager: PartitionRuleManagerRef, node_manager: NodeManagerRef, catalog_manager: CatalogManagerRef, - flow_notification_tx: mpsc::Sender, - flush_interval: Duration, + flow_notification_tx: Notifier, worker_idle_timeout: Duration, - max_batch_rows: usize, - flush_semaphore: Arc, + flush_policy: TimingFlushPolicy, + flush_limiter: FlushLimiter, ) { tokio::spawn(async move { + // The business batch and pending flush are populated and drained together. let mut batch = None; - let flush_timer = tokio::time::sleep(flush_interval); - tokio::pin!(flush_timer); + let mut pending_flush = PendingCore::new(flush_policy); let mut shutdown_rx = shutdown.subscribe(); let idle_deadline = tokio::time::Instant::now() + worker_idle_timeout; let idle_timer = tokio::time::sleep_until(idle_deadline); @@ -955,43 +934,37 @@ fn start_worker( cmd = rx.recv() => { match cmd { Some(WorkerCommand::Submit { table_batches, total_rows, ctx, response_tx, _permit }) => { - idle_timer.as_mut().reset(tokio::time::Instant::now() + worker_idle_timeout); + let submitted_at = tokio::time::Instant::now(); + idle_timer.as_mut().reset(submitted_at + worker_idle_timeout); - if batch.is_none() { - // Anchor the flush deadline to this batch's first submission, - // rather than to the worker's creation time. - flush_timer - .as_mut() - .reset(tokio::time::Instant::now() + flush_interval); - } + pending_flush.submit( + FlushWaiter { response_tx, _permit }, + total_rows, + ); let pending_batch = batch.get_or_insert_with(||{ PENDING_BATCHES.inc(); PendingBatch::new(ctx) }); - pending_batch.waiters.push(FlushWaiter { response_tx, _permit }); - for (table_name, table_id, record_batch) in table_batches { pending_batch.add_table_batch(table_name, table_id, record_batch); } - pending_batch.total_row_count += total_rows; PENDING_ROWS.add(total_rows as i64); - if pending_batch.total_row_count >= max_batch_rows - && let Some(flush) = drain_batch(&mut batch) { - spawn_flush( - flush, - partition_manager.clone(), - node_manager.clone(), - catalog_manager.clone(), - flow_notification_tx.clone(), - flush_semaphore.clone(), - ).await; + if let Some(flush) = drain_batch(&mut batch, &mut pending_flush, Some(FlushTrigger::Submission)) { + spawn_flush( + flush, + partition_manager.clone(), + node_manager.clone(), + catalog_manager.clone(), + flow_notification_tx.clone(), + flush_limiter.clone(), + ).await; } } None => { - if let Some(flush) = drain_batch(&mut batch) { + if let Some(flush) = drain_batch(&mut batch, &mut pending_flush, None) { flush_batch_with_managers( flush, partition_manager.clone(), @@ -1010,7 +983,7 @@ fn start_worker( } _ = &mut idle_timer => { if !should_close_worker_on_idle_timeout( - batch.as_ref().map_or(0, |batch| batch.total_row_count), + pending_flush.total_rows(), rx.len(), ) { idle_timer @@ -1027,20 +1000,20 @@ fn start_worker( ); break; } - _ = &mut flush_timer, if batch.is_some() => { - if let Some(flush) = drain_batch(&mut batch) { + _ = pending_flush.wait_flush() => { + if let Some(flush) = drain_batch(&mut batch, &mut pending_flush, Some(FlushTrigger::Deadline)) { spawn_flush( flush, partition_manager.clone(), node_manager.clone(), catalog_manager.clone(), flow_notification_tx.clone(), - flush_semaphore.clone(), + flush_limiter.clone(), ).await; } } _ = shutdown_rx.recv() => { - if let Some(flush) = drain_batch(&mut batch) { + if let Some(flush) = drain_batch(&mut batch, &mut pending_flush, None) { flush_batch_with_managers( flush, partition_manager.clone(), @@ -1054,41 +1027,46 @@ fn start_worker( } } - remove_worker_if_same_channel(workers.as_ref(), &key, &worker_tx); + remove_worker_if_same_channel(workers.as_ref(), &key, &worker_tx).await; }); } -fn remove_worker_if_same_channel( - workers: &DashMap, +async fn remove_worker_if_same_channel( + workers: &WorkerRegistry, key: &BatchKey, worker_tx: &mpsc::Sender, ) -> bool { - if let Some(worker) = workers.get(key) - && worker.tx.same_channel(worker_tx) - { - drop(worker); - workers.remove(key); - PENDING_WORKERS.set(workers.len() as i64); - return true; + if workers.remove_if_same(key, worker_tx).await { + PENDING_WORKERS.set(workers.len().await as i64); + true + } else { + false } - - false } fn should_close_worker_on_idle_timeout(total_row_count: usize, queued_requests: usize) -> bool { total_row_count == 0 && queued_requests == 0 } -fn drain_batch(batch: &mut Option) -> Option { +/// Transfers the ready batch, or drains unconditionally when no trigger is given. +/// Execution and flush permits remain owned by the caller. +fn drain_batch( + batch: &mut Option, + pending_flush: &mut PendingCore, + trigger: Option, +) -> Option { + let total_row_count = pending_flush.total_rows(); + let waiters = match trigger { + Some(trigger) => pending_flush.take_ready(trigger)?, + None => pending_flush.take_pending()?, + }; let batch = batch.take()?; - let total_row_count = batch.total_row_count; if total_row_count == 0 { return None; } let table_batches = batch.tables.into_values().collect(); - let waiters = batch.waiters; PENDING_ROWS.sub(total_row_count as i64); PENDING_BATCHES.dec(); @@ -1107,10 +1085,10 @@ async fn spawn_flush( partition_manager: PartitionRuleManagerRef, node_manager: NodeManagerRef, catalog_manager: CatalogManagerRef, - flow_notification_tx: mpsc::Sender, - semaphore: Arc, + flow_notification_tx: Notifier, + flush_limiter: FlushLimiter, ) { - match semaphore.acquire_owned().await { + match flush_limiter.acquire().await { Ok(permit) => { tokio::spawn(async move { let _permit = permit; @@ -1323,7 +1301,7 @@ async fn flush_batch_with_managers( partition_manager: PartitionRuleManagerRef, node_manager: NodeManagerRef, catalog_manager: CatalogManagerRef, - flow_notification_tx: mpsc::Sender, + flow_notification_tx: Notifier, ) { let partition_provider = PartitionManagerPhysicalFlushAdapter { partition_manager }; let node_requester = NodeManagerPhysicalFlushAdapter { @@ -1345,7 +1323,7 @@ async fn flush_batch( partition_manager: &(impl PhysicalFlushPartitionProvider + ?Sized), node_manager: &(impl PhysicalFlushNodeRequester + ?Sized), catalog_manager: &(impl PhysicalFlushCatalogProvider + ?Sized), - flow_notification_tx: mpsc::Sender, + flow_notification_tx: Notifier, ) { let FlushBatch { table_batches, @@ -1403,21 +1381,13 @@ fn extract_timestamps(table_batch: &TableBatch) -> Vec { let mut timestamps = Vec::with_capacity(table_batch.row_count); for batch in &table_batch.batches { let timestamp_column = batch.batch.column(batch.timestamp_index); - let Some((timestamp_values, _)) = - datatypes::timestamp::timestamp_array_to_primitive(timestamp_column) - else { + let Some(()) = append_timestamps(timestamp_column, &mut timestamps) else { error!( "Failed to extract timestamps from record batch, table_id: {}, timestamp_index: {}", table_batch.table_id, batch.timestamp_index ); continue; }; - - if timestamp_values.null_count() == 0 { - timestamps.extend_from_slice(timestamp_values.values()); - } else { - timestamps.extend(timestamp_values.iter().flatten()); - } } timestamps } @@ -1428,10 +1398,10 @@ struct FlowNotification { } fn try_enqueue_flow_notification( - tx: &mpsc::Sender, + tx: &Notifier, notification: FlowNotification, ) -> bool { - match tx.try_send(notification) { + match tx.try_notify(notification) { Ok(()) => true, Err(mpsc::error::TrySendError::Full(notification)) => { FLOW_NOTIFICATION_DROPPED.with_label_values(&["full"]).inc(); @@ -1456,7 +1426,7 @@ fn try_enqueue_flow_notification( } } -fn enqueue_flow_notifications(table_batches: Vec, tx: &mpsc::Sender) { +fn enqueue_flow_notifications(table_batches: Vec, tx: &Notifier) { for table_batch in table_batches { let timestamps = extract_timestamps(&table_batch); if timestamps.is_empty() { @@ -1517,14 +1487,17 @@ fn start_flow_notification_worker( table_flownode_set_cache: TableFlownodeSetCacheRef, node_manager: NodeManagerRef, ) { - common_runtime::spawn_global(async move { - tokio_stream::wrappers::ReceiverStream::new(notification_rx) - .for_each_concurrent(MAX_CONCURRENT_FLOW_NOTIFICATIONS, |notification| { + spawn_global(async move { + run_notifier( + notification_rx, + MAX_CONCURRENT_FLOW_NOTIFICATIONS, + |notification| { let table_flownode_set_cache = table_flownode_set_cache.clone(); let node_manager = node_manager.clone(); handle_flow_notification(notification, table_flownode_set_cache, node_manager) - }) - .await; + }, + ) + .await; }); } @@ -1534,7 +1507,7 @@ fn notify_flow_dirty_windows_after_flush( table_flownode_set_cache: TableFlownodeSetCacheRef, node_manager: NodeManagerRef, ) { - let (tx, rx) = mpsc::channel(table_batches.len().max(1)); + let (tx, rx) = Notifier::try_new(table_batches.len().max(1)).unwrap(); start_flow_notification_worker(rx, table_flownode_set_cache, node_manager); enqueue_flow_notifications(table_batches, &tx); } @@ -1804,7 +1777,7 @@ fn encode_region_write_requests( let _timer = PENDING_ROWS_BATCH_FLUSH_STAGE_ELAPSED .with_label_values(&["flush_physical_encode_ipc"]) .start_timer(); - record_batch_to_ipc(resolved.planned.batch)? + record_batch_to_ipc(resolved.planned.batch).map_err(map_ipc_error)? }; let request = RegionRequest { @@ -1847,36 +1820,24 @@ fn notify_waiters(waiters: Vec, result: Result<()>) { } } -fn record_batch_to_ipc(record_batch: RecordBatch) -> Result<(Bytes, Bytes, Bytes)> { - let mut encoder = FlightEncoder::default(); - let schema = encoder.encode_schema(record_batch.schema().as_ref()); - let mut iter = encoder - .encode(FlightMessage::RecordBatch(record_batch)) - .into_iter(); - let Some(flight_data) = iter.next() else { - return Err(Error::Internal { - err_msg: "Failed to encode empty flight data".to_string(), - }); - }; - if iter.next().is_some() { - return Err(Error::NotSupported { - feat: "bulk insert RecordBatch with dictionary arrays".to_string(), - }); +fn map_ipc_error(error: GrpcError) -> Error { + match error { + GrpcError::NotSupported { feat } => Error::NotSupported { feat }, + GrpcError::InvalidFlightData { reason, .. } => Error::Internal { err_msg: reason }, + error => Error::Internal { + err_msg: error.to_string(), + }, } - - Ok(( - schema.data_header, - flight_data.data_header, - flight_data.data_body, - )) } #[cfg(test)] mod tests { use std::any::Any; use std::collections::{HashMap, HashSet}; + use std::future::{Future, poll_fn}; use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::{Arc, Mutex}; + use std::task::Poll; use std::time::Duration; use api::region::RegionResponse; @@ -1912,7 +1873,6 @@ mod tests { }; use common_query::request::QueryRequest; use common_recordbatch::SendableRecordBatchStream; - use dashmap::DashMap; use datatypes::schema::{ColumnSchema as DtColumnSchema, Schema as DtSchema}; use moka::future::CacheBuilder; use partition::cache::new_partition_info_cache; @@ -1927,21 +1887,22 @@ mod tests { use tokio::sync::{Notify, Semaphore, broadcast, mpsc, oneshot}; use tokio::time::{advance, sleep}; - use super::{ - BatchKey, Error, FlushBatch, FlushRegionWrite, FlushWaiter, PendingBatch, - PendingRowsBatcher, PendingWorker, PhysicalFlushCatalogProvider, + use crate::error; + use crate::metrics::FLOW_NOTIFICATION_DROPPED; + use crate::pending_rows_batcher::{ + BatchKey, Error, FlushBatch, FlushLimiter, FlushRegionWrite, FlushTrigger, FlushWaiter, + Notifier, PendingBatch, PendingCore, PendingRowsBatcher, PhysicalFlushCatalogProvider, PhysicalFlushNodeRequester, PhysicalFlushPartitionProvider, PhysicalTableMetadata, - PlannedRegionBatch, RecordBatchWithTsIdx, ResolvedRegionBatch, TableBatch, WorkerCommand, - batch_key_from_ctx, columns_taxonomy, drain_batch, encode_region_write_requests, - extract_timestamps, flush_batch, flush_batch_physical, flush_region_writes_concurrently, - greptime_timestamp, notify_flow_dirty_windows_after_flush, plan_region_batches, + PlannedRegionBatch, RecordBatchWithTsIdx, RequestLimiter, ResolvedRegionBatch, TableBatch, + TimingFlushPolicy, WorkerCommand, WorkerRegistry, batch_key_from_ctx, columns_taxonomy, + drain_batch, encode_region_write_requests, extract_timestamps, flush_batch, + flush_batch_physical, flush_region_writes_concurrently, greptime_timestamp, + notify_flow_dirty_windows_after_flush, notify_waiters, plan_region_batches, remove_worker_if_same_channel, should_close_worker_on_idle_timeout, should_dispatch_concurrently, start_flow_notification_worker, start_worker, strip_partition_columns_from_batch, transform_logical_batches_to_physical, try_enqueue_flow_notification, }; - use crate::error; - use crate::metrics::FLOW_NOTIFICATION_DROPPED; use crate::prom_row_builder::rows_to_aligned_record_batch; fn mock_rows(row_count: usize, schema_name: &str) -> Rows { @@ -2165,8 +2126,8 @@ mod tests { #[test] fn test_flow_notification_queue_drops_when_full() { - let (tx, mut rx) = mpsc::channel(1); - let notification = |table_id| super::FlowNotification { + let (tx, mut rx) = Notifier::try_new(1).unwrap(); + let notification = |table_id| crate::pending_rows_batcher::FlowNotification { table_id, timestamps: vec![table_id as i64], }; @@ -2178,6 +2139,12 @@ mod tests { assert_eq!(1, rx.try_recv().unwrap().table_id); assert_eq!(dropped_before + 1, dropped.get()); + + let closed = FLOW_NOTIFICATION_DROPPED.with_label_values(&["closed"]); + let closed_before = closed.get(); + drop(rx); + assert!(!try_enqueue_flow_notification(&tx, notification(3))); + assert_eq!(closed_before + 1, closed.get()); } fn mock_physical_table_metadata(table_id: TableId) -> PhysicalTableMetadata { @@ -2375,11 +2342,20 @@ mod tests { assert_eq!(2, table_rows[0].1.rows.len()); } - #[test] - fn test_drain_batch_takes_initialized_pending_batch_from_option() { + #[tokio::test] + async fn test_drain_batch_takes_initialized_pending_batch_from_option() { let ctx = session::context::QueryContext::arc(); let (response_tx, _response_rx) = oneshot::channel(); let permit = Arc::new(Semaphore::new(1)).try_acquire_owned().unwrap(); + let mut pending_flush = + PendingCore::new(TimingFlushPolicy::try_new(Duration::from_secs(10), 1).unwrap()); + pending_flush.submit( + FlushWaiter { + response_tx, + _permit: Arc::new(permit), + }, + 1, + ); let mut batch = Some(PendingBatch { tables: HashMap::from([( 42, @@ -2390,24 +2366,77 @@ mod tests { row_count: 1, }, )]), - total_row_count: 1, db_string: ctx.get_db_string(), ctx: ctx.clone(), - waiters: vec![FlushWaiter { - response_tx, - _permit: permit, - }], }); - let flush = drain_batch(&mut batch).unwrap(); + let flush = drain_batch( + &mut batch, + &mut pending_flush, + Some(FlushTrigger::Submission), + ) + .unwrap(); assert!(batch.is_none()); + assert!(pending_flush.is_empty()); + assert_eq!(0, pending_flush.total_rows()); + assert_eq!(1, flush.waiters.len()); assert_eq!(1, flush.total_row_count); assert_eq!(1, flush.table_batches.len()); assert_eq!(ctx.get_db_string(), flush.db_string); assert_eq!(ctx.current_catalog(), flush.ctx.current_catalog()); } + #[tokio::test] + async fn test_drain_batch_preserves_unready_state_and_clears_zero_rows() { + for total_rows in [0, 1] { + let mut pending_flush = + PendingCore::new(TimingFlushPolicy::try_new(Duration::from_secs(10), 2).unwrap()); + let mut batch = Some(PendingBatch::new(session::context::QueryContext::arc())); + let semaphore = Arc::new(Semaphore::new(1)); + let (response_tx, mut response_rx) = oneshot::channel(); + pending_flush.submit( + FlushWaiter { + response_tx, + _permit: Arc::new(semaphore.clone().acquire_owned().await.unwrap()), + }, + total_rows, + ); + assert!( + drain_batch( + &mut batch, + &mut pending_flush, + Some(FlushTrigger::Submission) + ) + .is_none() + ); + assert!(batch.is_some()); + assert!(!pending_flush.is_empty()); + assert_eq!(total_rows, pending_flush.total_rows()); + assert_eq!(0, semaphore.available_permits()); + + let drained = drain_batch(&mut batch, &mut pending_flush, None); + assert!(batch.is_none()); + assert!(pending_flush.is_empty()); + assert_eq!(0, pending_flush.total_rows()); + if total_rows == 0 { + assert!(drained.is_none()); + assert_eq!(1, semaphore.available_permits()); + assert!(matches!( + response_rx.try_recv(), + Err(oneshot::error::TryRecvError::Closed) + )); + } else { + let drained = drained.unwrap(); + assert_eq!(1, drained.total_row_count); + assert_eq!(1, drained.waiters.len()); + assert_eq!(0, semaphore.available_permits()); + drop(drained); + assert_eq!(1, semaphore.available_permits()); + } + } + } + #[test] fn test_pending_batch_keeps_same_name_batches_with_distinct_table_ids() { let ctx = session::context::QueryContext::arc(); @@ -2572,8 +2601,8 @@ mod tests { fn mock_flow_notification_sender( cache: TableFlownodeSetCacheRef, node_manager: NodeManagerRef, - ) -> mpsc::Sender { - let (tx, rx) = mpsc::channel(16); + ) -> Notifier { + let (tx, rx) = Notifier::try_new(16).unwrap(); start_flow_notification_worker(rx, cache, node_manager); tx } @@ -2930,9 +2959,25 @@ mod tests { assert_eq!(batches[&batch_key_from_ctx(&skip_wal_ctx)], 2); } - #[test] - fn test_remove_worker_if_same_channel_removes_matching_entry() { - let workers = DashMap::new(); + #[tokio::test] + async fn test_cancelled_waiter_retains_request_slot_until_notification() { + let limiter = RequestLimiter::try_new(1).unwrap(); + let (response_tx, response_rx) = oneshot::channel(); + let waiter = FlushWaiter { + response_tx, + _permit: limiter.acquire().await.unwrap(), + }; + drop(response_rx); + let next = limiter.acquire(); + tokio::pin!(next); + assert!(poll_fn(|cx| Poll::Ready(next.as_mut().poll(cx).is_pending())).await); + notify_waiters(vec![waiter], Ok(())); + let _permit = next.await.unwrap(); + } + + #[tokio::test] + async fn test_remove_worker_if_same_channel_removes_matching_entry() { + let workers = WorkerRegistry::new(); let key = BatchKey { catalog: "greptime".to_string(), schema: "public".to_string(), @@ -2941,15 +2986,15 @@ mod tests { }; let (tx, _rx) = mpsc::channel::(1); - workers.insert(key.clone(), PendingWorker { tx: tx.clone() }); + workers.get_or_insert_with(key.clone(), || tx.clone()).await; - assert!(remove_worker_if_same_channel(&workers, &key, &tx)); - assert!(!workers.contains_key(&key)); + assert!(remove_worker_if_same_channel(&workers, &key, &tx).await); + assert!(workers.is_empty().await); } - #[test] - fn test_remove_worker_if_same_channel_keeps_newer_entry() { - let workers = DashMap::new(); + #[tokio::test] + async fn test_remove_worker_if_same_channel_keeps_newer_entry() { + let workers = WorkerRegistry::new(); let key = BatchKey { catalog: "greptime".to_string(), schema: "public".to_string(), @@ -2959,16 +3004,13 @@ mod tests { let (stale_tx, _stale_rx) = mpsc::channel::(1); let (fresh_tx, _fresh_rx) = mpsc::channel::(1); - workers.insert( - key.clone(), - PendingWorker { - tx: fresh_tx.clone(), - }, - ); + workers + .get_or_insert_with(key.clone(), || fresh_tx.clone()) + .await; - assert!(!remove_worker_if_same_channel(&workers, &key, &stale_tx)); - assert!(workers.contains_key(&key)); - assert!(workers.get(&key).unwrap().tx.same_channel(&fresh_tx)); + assert!(!remove_worker_if_same_channel(&workers, &key, &stale_tx).await); + assert!(workers.get(&key).await.is_some()); + assert!(workers.get(&key).await.unwrap().same_channel(&fresh_tx)); } #[test] @@ -2997,7 +3039,7 @@ mod tests { total_rows, ctx: session::context::QueryContext::arc(), response_tx, - _permit: permit, + _permit: Arc::new(permit), }) .await .unwrap(); @@ -3039,7 +3081,7 @@ mod tests { } #[tokio::test(start_paused = true)] - async fn test_worker_rearms_creation_relative_deadline_after_size_flush() { + async fn test_worker_preserves_first_deadline_and_inline_shutdown() { let flush_interval = Duration::from_secs(10); let worker_idle_timeout = Duration::from_secs(30); let key = BatchKey { @@ -3048,14 +3090,11 @@ mod tests { physical_table: "phy".to_string(), skip_wal: false, }; - let workers = Arc::new(DashMap::new()); + let workers = Arc::new(WorkerRegistry::new()); let (worker_tx, worker_rx) = mpsc::channel(1); - workers.insert( - key.clone(), - PendingWorker { - tx: worker_tx.clone(), - }, - ); + workers + .get_or_insert_with(key.clone(), || worker_tx.clone()) + .await; let backend = Arc::new(MemoryKvBackend::default()); let table_route_cache = Arc::new(new_table_route_cache( @@ -3077,9 +3116,10 @@ mod tests { datanodes: Arc::new(HashMap::new()), }); let catalog_manager = MemoryCatalogManager::with_default_setup(); - let (flow_notification_tx, _flow_notification_rx) = mpsc::channel(1); + let (flow_notification_tx, _flow_notification_rx) = Notifier::try_new(1).unwrap(); let (shutdown, _) = broadcast::channel(1); + let flush_limiter = FlushLimiter::try_new(1).unwrap(); start_worker( key.clone(), worker_tx.clone(), @@ -3090,10 +3130,9 @@ mod tests { node_manager, catalog_manager, flow_notification_tx, - flush_interval, worker_idle_timeout, - 2, - Arc::new(Semaphore::new(1)), + TimingFlushPolicy::try_new(flush_interval, 3).unwrap(), + flush_limiter.clone(), ); // Start the worker, then size-flush a batch halfway to the first @@ -3101,7 +3140,7 @@ mod tests { // drains the batch before that deadline is reached. tokio::task::yield_now().await; advance(flush_interval / 2).await; - let size_flush_rx = submit_mock_worker_batch(&worker_tx, 2, 1000).await; + let size_flush_rx = submit_mock_worker_batch(&worker_tx, 3, 1000).await; let size_flush_result = receive_mock_flush_result(size_flush_rx, "row threshold did not flush the first batch") .await; @@ -3112,7 +3151,10 @@ mod tests { advance(flush_interval / 5).await; let mut timed_flush_rx = submit_mock_worker_batch(&worker_tx, 1, 2000).await; - advance(flush_interval - Duration::from_millis(1)).await; + let first_submission = tokio::time::Instant::now(); + advance(flush_interval / 2).await; + let later_flush_rx = submit_mock_worker_batch(&worker_tx, 1, 3000).await; + advance(flush_interval / 2 - Duration::from_millis(1)).await; for _ in 0..10 { tokio::task::yield_now().await; } @@ -3129,15 +3171,39 @@ mod tests { .await; assert_missing_physical_table(timed_flush_result); + assert_eq!( + first_submission + flush_interval, + tokio::time::Instant::now() + ); + assert_missing_physical_table( + receive_mock_flush_result( + later_flush_rx, + "later submission was not included in the timed flush", + ) + .await, + ); + + // Shutdown retains the existing inline path even with no flush permits. + let _held_permit = flush_limiter.acquire().await.unwrap(); + let shutdown_flush_rx = submit_mock_worker_batch(&worker_tx, 1, 4000).await; + let shutdown_at = tokio::time::Instant::now(); let _ = shutdown.send(()); + assert_missing_physical_table( + receive_mock_flush_result( + shutdown_flush_rx, + "shutdown incorrectly waited for a flush permit", + ) + .await, + ); + assert_eq!(shutdown_at, tokio::time::Instant::now()); for _ in 0..10 { - if !workers.contains_key(&key) { + if workers.is_empty().await { break; } tokio::task::yield_now().await; } assert!( - !workers.contains_key(&key), + workers.is_empty().await, "worker did not exit after shutdown" ); }