refactor: compose Prom ingestion with common batcher components

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

Cargo.lock

@@ -2316,6 +2316,14 @@ dependencies = [
  "zeroize",
 ]
 
+[[package]]
+name = "common-batcher"
+version = "1.3.0-alpha.1"
+dependencies = [
+ "futures",
+ "tokio",
+]
+
 [[package]]
 name = "common-catalog"
 version = "1.3.0-alpha.1"
@@ -13556,6 +13564,7 @@ dependencies = [
  "chrono",
  "client",
  "common-base",
+ "common-batcher",
  "common-catalog",
  "common-decimal",
  "common-error",

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" }

src/common/batcher/Cargo.toml

@@ -0,0 +1,18 @@
+[package]
+name = "common-batcher"
+version.workspace = true
+edition.workspace = true
+license.workspace = true
+
+[lib]
+path = "src/lib.rs"
+
+[lints]
+workspace = true
+
+[dependencies]
+futures.workspace = true
+tokio.workspace = true
+
+[dev-dependencies]
+tokio = { workspace = true, features = ["test-util"] }

src/common/batcher/src/flush_limiter.rs

@@ -0,0 +1,69 @@
+// 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::num::NonZeroUsize;
+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<Semaphore>,
+}
+
+impl FlushLimiter {
+    /// Returns `None` if the limit exceeds Tokio's supported semaphore capacity.
+    pub fn try_new(max_concurrent_flushes: NonZeroUsize) -> Option<Self> {
+        let permits = max_concurrent_flushes.get();
+        (permits <= Semaphore::MAX_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<OwnedSemaphorePermit, AcquireError> {
+        self.semaphore.clone().acquire_owned().await
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use std::future::Future;
+
+    use super::*;
+
+    #[test]
+    fn test_invalid_capacity() {
+        assert!(FlushLimiter::try_new(NonZeroUsize::new(usize::MAX).unwrap()).is_none());
+    }
+
+    #[tokio::test]
+    async fn test_clones_share_flush_budget() {
+        let limiter = FlushLimiter::try_new(NonZeroUsize::new(1).unwrap()).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();
+    }
+}

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<T>(&self, batch: &PendingBatch<T>) -> Option<Instant>;
+
+    /// Returns whether the batch should be flushed for this event at `now`.
+    ///
+    /// Workers never flush empty batches, regardless of this decision.
+    fn should_flush<T>(&self, batch: &PendingBatch<T>, now: Instant, trigger: FlushTrigger)
+    -> bool;
+}

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

@@ -0,0 +1,147 @@
+// 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::num::NonZeroUsize;
+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: NonZeroUsize,
+}
+
+impl TimingFlushPolicy {
+    /// Creates a policy, rejecting zero intervals and unrepresentable deadlines.
+    pub fn try_new(flush_interval: Duration, max_batch_rows: NonZeroUsize) -> Option<Self> {
+        if flush_interval.is_zero() {
+            return None;
+        }
+        Instant::now().checked_add(flush_interval)?;
+        Some(Self {
+            flush_interval,
+            max_batch_rows,
+        })
+    }
+
+    /// Checks only the row trigger when a caller drives deadline events separately.
+    pub fn reached_row_threshold<T>(&self, batch: &PendingBatch<T>) -> bool {
+        !batch.is_empty() && batch.total_rows() >= self.max_batch_rows.get()
+    }
+}
+
+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<T>(&self, batch: &PendingBatch<T>) -> Option<Instant> {
+        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<T>(
+        &self,
+        batch: &PendingBatch<T>,
+        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_first_submission_deadline() {
+        let policy =
+            TimingFlushPolicy::try_new(Duration::from_millis(10), NonZeroUsize::new(100).unwrap())
+                .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), NonZeroUsize::new(3).unwrap())
+                .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() {
+        let max_batch_rows = NonZeroUsize::new(1).unwrap();
+        assert!(TimingFlushPolicy::try_new(Duration::ZERO, max_batch_rows).is_none());
+        assert!(TimingFlushPolicy::try_new(Duration::MAX, max_batch_rows).is_none());
+    }
+}

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<Instant>,
+    sleep: Option<Pin<Box<Sleep>>>,
+}
+
+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<Instant>) {
+        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);
+    }
+}

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;

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<T> {
+    sender: Sender<T>,
+}
+
+impl<T> Clone for Notifier<T> {
+    fn clone(&self) -> Self {
+        Self {
+            sender: self.sender.clone(),
+        }
+    }
+}
+
+impl<T> Notifier<T> {
+    /// Returns `None` if the queue exceeds Tokio's supported capacity.
+    pub fn try_new(capacity: NonZeroUsize) -> Option<(Self, Receiver<T>)> {
+        if capacity.get() > Semaphore::MAX_PERMITS {
+            return None;
+        }
+        let (sender, receiver) = mpsc::channel(capacity.get());
+        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<T>> {
+        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<T, F, Fut>(
+    receiver: Receiver<T>,
+    max_concurrent_notifications: NonZeroUsize,
+    handler: F,
+) where
+    F: FnMut(T) -> Fut,
+    Fut: Future<Output = ()>,
+{
+    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::<usize>::try_new(NonZeroUsize::new(usize::MAX).unwrap()).is_none());
+        let (notifier, receiver) = Notifier::try_new(NonZeroUsize::new(1).unwrap()).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(NonZeroUsize::new(3).unwrap()).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);
+            }
+        }
+    }
+}

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<T> {
+    items: Vec<T>,
+    total_rows: usize,
+    first_submitted_at: Option<Instant>,
+}
+
+impl<T> Default for PendingBatch<T> {
+    fn default() -> Self {
+        Self {
+            items: Vec::new(),
+            total_rows: 0,
+            first_submitted_at: None,
+        }
+    }
+}
+
+impl<T> PendingBatch<T> {
+    /// 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<Instant> {
+        self.first_submitted_at
+    }
+
+    /// Takes all submissions in arrival order and resets the batch's accounting.
+    pub fn take(&mut self) -> Vec<T> {
+        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]);
+    }
+}

src/common/batcher/src/pending_worker.rs

@@ -0,0 +1,264 @@
+// 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<T, P> {
+    batch: PendingBatch<T>,
+    flush_policy: P,
+    flush_timer: FlushTimer,
+}
+
+impl<T, P> PendingWorker<T, P> {
+    /// 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<Vec<T>> {
+        self.flush_timer.set_deadline(None);
+        if self.batch.is_empty() {
+            None
+        } else {
+            Some(self.batch.take())
+        }
+    }
+}
+
+impl<T, P: FlushPolicy> PendingWorker<T, P> {
+    /// 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<Vec<T>> {
+        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::num::NonZeroUsize;
+    use std::task::Poll;
+    use std::time::Duration;
+
+    use super::*;
+    use crate::flush_policy::timing::TimingFlushPolicy;
+
+    fn worker(rows: usize) -> PendingWorker<i32, TimingFlushPolicy> {
+        PendingWorker::new(
+            TimingFlushPolicy::try_new(Duration::from_millis(10), NonZeroUsize::new(rows).unwrap())
+                .unwrap(),
+        )
+    }
+
+    async fn assert_wait_pending<T, P>(worker: &mut PendingWorker<T, P>) {
+        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<T>(&self, _: &PendingBatch<T>) -> Option<Instant> {
+                None
+            }
+            fn should_flush<T>(
+                &self,
+                batch: &PendingBatch<T>,
+                _: 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());
+    }
+}

src/common/batcher/src/request_limiter.rs

@@ -0,0 +1,117 @@
+// 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::num::NonZeroUsize;
+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<Semaphore>,
+}
+
+impl RequestLimiter {
+    /// Returns `None` if the limit exceeds Tokio's supported semaphore capacity.
+    pub fn try_new(max_inflight_requests: NonZeroUsize) -> Option<Self> {
+        let permits = max_inflight_requests.get();
+        (permits <= Semaphore::MAX_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<Arc<OwnedSemaphorePermit>, AcquireError> {
+        self.semaphore.clone().acquire_owned().await.map(Arc::new)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use std::future::{Future, poll_fn};
+    use std::num::NonZeroUsize;
+    use std::pin::Pin;
+    use std::task::Poll;
+
+    use tokio::sync::Semaphore;
+
+    use crate::request_limiter::RequestLimiter;
+
+    async fn is_pending<F: Future>(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(NonZeroUsize::new(capacity).unwrap()).is_some());
+        }
+        for capacity in [Semaphore::MAX_PERMITS + 1, usize::MAX] {
+            assert!(RequestLimiter::try_new(NonZeroUsize::new(capacity).unwrap()).is_none());
+        }
+    }
+
+    #[tokio::test]
+    async fn test_last_submission_releases_request_slot() {
+        let limiter = RequestLimiter::try_new(NonZeroUsize::new(1).unwrap()).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(NonZeroUsize::new(1).unwrap()).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(NonZeroUsize::new(2).unwrap()).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);
+    }
+}

src/common/batcher/src/worker_registry.rs

@@ -0,0 +1,188 @@
+// 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::collections::HashMap;
+use std::hash::Hash;
+
+use tokio::sync::Mutex;
+use tokio::sync::mpsc::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<K, T> {
+    workers: Mutex<HashMap<K, Sender<T>>>,
+}
+
+impl<K, T> Default for WorkerRegistry<K, T> {
+    fn default() -> Self {
+        Self {
+            workers: Mutex::new(HashMap::new()),
+        }
+    }
+}
+
+impl<K: Eq + Hash, T> WorkerRegistry<K, T> {
+    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<Sender<T>> {
+        self.workers
+            .lock()
+            .await
+            .get(key)
+            .filter(|tx| !tx.is_closed())
+            .cloned()
+    }
+
+    /// Reuses a live sender or atomically creates its replacement.
+    ///
+    /// `create` runs synchronously under the registry 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<F>(&self, key: K, create: F) -> Sender<T>
+    where
+        F: FnOnce() -> Sender<T>,
+    {
+        let mut workers = self.workers.lock().await;
+        if let Some(tx) = workers.get(&key)
+            && !tx.is_closed()
+        {
+            return tx.clone();
+        }
+        let tx = create();
+        workers.insert(key, tx.clone());
+        tx
+    }
+
+    /// Removes the key only if it still points to this worker's channel.
+    pub async fn remove_if_same(&self, key: &K, tx: &Sender<T>) -> bool {
+        let mut workers = self.workers.lock().await;
+        if workers
+            .get(key)
+            .is_some_and(|current| current.same_channel(tx))
+        {
+            workers.remove(key);
+            true
+        } else {
+            false
+        }
+    }
+
+    /// Number of registered entries, including senders whose receivers closed.
+    pub async fn len(&self) -> usize {
+        self.workers.lock().await.len()
+    }
+
+    pub async fn is_empty(&self) -> bool {
+        self.workers.lock().await.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_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]
+    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);
+    }
+}

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

src/servers/src/pending_rows_batcher.rs

@@ -30,15 +30,20 @@ use arrow::record_batch::RecordBatch;
 use async_trait::async_trait;
 use bytes::Bytes;
 use catalog::CatalogManagerRef;
+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::flight::{FlightEncoder, FlightMessage};
 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 metric_engine::batch_modifier::{TagColumnInfo, modify_batch_sparse};
 use partition::manager::PartitionRuleManagerRef;
 use partition::partition::PartitionRuleRef;
@@ -47,7 +52,7 @@ use smallvec::SmallVec;
 use snafu::{OptionExt, ResultExt, ensure};
 use store_api::storage::{RegionId, TableId};
 use table::metadata::{TableInfo, TableInfoRef};
-use tokio::sync::{OwnedSemaphorePermit, Semaphore, broadcast, mpsc, oneshot};
+use tokio::sync::{OwnedSemaphorePermit, broadcast, mpsc, oneshot};
 
 use crate::error;
 use crate::error::{Error, Result};
@@ -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<TableId, TableBatch>,
-    total_row_count: usize,
     db_string: String,
     ctx: QueryContextRef,
-    waiters: Vec<FlushWaiter>,
 }
 
 struct FlushWaiter {
     response_tx: oneshot::Sender<std::result::Result<(), Arc<Error>>>,
-    _permit: OwnedSemaphorePermit,
+    _permit: Arc<OwnedSemaphorePermit>,
 }
 
 struct FlushBatch {
@@ -322,7 +325,7 @@ enum WorkerCommand {
         total_rows: usize,
         ctx: QueryContextRef,
         response_tx: oneshot::Sender<std::result::Result<(), Arc<Error>>>,
-        _permit: OwnedSemaphorePermit,
+        _permit: Arc<OwnedSemaphorePermit>,
     },
     #[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<DashMap<BatchKey, PendingWorker>>,
+    workers: Arc<WorkerRegistry<BatchKey, WorkerCommand>>,
     flush_interval: Duration,
-    max_batch_rows: usize,
+    flush_policy: TimingFlushPolicy,
     partition_manager: PartitionRuleManagerRef,
     node_manager: NodeManagerRef,
     catalog_manager: CatalogManagerRef,
-    flow_notification_tx: mpsc::Sender<FlowNotification>,
-    flush_semaphore: Arc<Semaphore>,
-    inflight_semaphore: Arc<Semaphore>,
+    flow_notification_tx: Notifier<FlowNotification>,
+    flush_limiter: FlushLimiter,
+    request_limiter: RequestLimiter,
     worker_channel_capacity: usize,
     prom_store_with_metric_engine: bool,
     schema_alterer: PendingRowsSchemaAltererRef,
@@ -388,12 +391,18 @@ impl PendingRowsBatcher {
             return None;
         }
 
+        let flush_policy =
+            TimingFlushPolicy::try_new(flush_interval, NonZeroUsize::new(max_batch_rows)?)?;
+        let flush_limiter = FlushLimiter::try_new(NonZeroUsize::new(max_concurrent_flushes)?)?;
+
+        let request_limiter = RequestLimiter::try_new(NonZeroUsize::new(max_inflight_requests)?)?;
+        let (flow_notification_tx, flow_notification_rx) =
+            Notifier::try_new(flow_notification_queue_capacity)?;
+
         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 +412,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 +442,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 +465,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 +478,8 @@ impl PendingRowsBatcher {
                             self.workers.as_ref(),
                             &batch_key,
                             &worker.tx,
-                        );
+                        )
+                        .await;
                     }
                 }
             }
@@ -833,38 +842,29 @@ 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 mut receiver = None;
+        let tx = self
+            .workers
+            .get_or_insert_with(key.clone(), || {
+                let (tx, rx) = mpsc::channel(self.worker_channel_capacity);
+                receiver = Some(rx);
+                tx
+            })
+            .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<WorkerCommand>,
+        rx: mpsc::Receiver<WorkerCommand>,
+    ) {
         let worker_idle_timeout = self
             .flush_interval
             .checked_mul(WORKER_IDLE_TIMEOUT_MULTIPLIER)
@@ -872,7 +872,7 @@ impl PendingRowsBatcher {
 
         start_worker(
             key,
-            worker.tx.clone(),
+            tx,
             self.workers.clone(),
             rx,
             self.shutdown.clone(),
@@ -880,13 +880,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 +898,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 +924,21 @@ impl PendingBatch {
 fn start_worker(
     key: BatchKey,
     worker_tx: mpsc::Sender<WorkerCommand>,
-    workers: Arc<DashMap<BatchKey, PendingWorker>>,
+    workers: Arc<WorkerRegistry<BatchKey, WorkerCommand>>,
     mut rx: mpsc::Receiver<WorkerCommand>,
     shutdown: broadcast::Sender<()>,
     partition_manager: PartitionRuleManagerRef,
     node_manager: NodeManagerRef,
     catalog_manager: CatalogManagerRef,
-    flow_notification_tx: mpsc::Sender<FlowNotification>,
-    flush_interval: Duration,
+    flow_notification_tx: Notifier<FlowNotification>,
     worker_idle_timeout: Duration,
-    max_batch_rows: usize,
-    flush_semaphore: Arc<Semaphore>,
+    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 +949,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 +998,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 +1015,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 +1042,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<BatchKey, PendingWorker>,
+async fn remove_worker_if_same_channel(
+    workers: &WorkerRegistry<BatchKey, WorkerCommand>,
     key: &BatchKey,
     worker_tx: &mpsc::Sender<WorkerCommand>,
 ) -> 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<PendingBatch>) -> Option<FlushBatch> {
+/// 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<PendingBatch>,
+    pending_flush: &mut PendingCore<FlushWaiter, TimingFlushPolicy>,
+    trigger: Option<FlushTrigger>,
+) -> Option<FlushBatch> {
+    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 +1100,10 @@ async fn spawn_flush(
     partition_manager: PartitionRuleManagerRef,
     node_manager: NodeManagerRef,
     catalog_manager: CatalogManagerRef,
-    flow_notification_tx: mpsc::Sender<FlowNotification>,
-    semaphore: Arc<Semaphore>,
+    flow_notification_tx: Notifier<FlowNotification>,
+    flush_limiter: FlushLimiter,
 ) {
-    match semaphore.acquire_owned().await {
+    match flush_limiter.acquire().await {
         Ok(permit) => {
             tokio::spawn(async move {
                 let _permit = permit;
@@ -1323,7 +1316,7 @@ async fn flush_batch_with_managers(
     partition_manager: PartitionRuleManagerRef,
     node_manager: NodeManagerRef,
     catalog_manager: CatalogManagerRef,
-    flow_notification_tx: mpsc::Sender<FlowNotification>,
+    flow_notification_tx: Notifier<FlowNotification>,
 ) {
     let partition_provider = PartitionManagerPhysicalFlushAdapter { partition_manager };
     let node_requester = NodeManagerPhysicalFlushAdapter {
@@ -1345,7 +1338,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<FlowNotification>,
+    flow_notification_tx: Notifier<FlowNotification>,
 ) {
     let FlushBatch {
         table_batches,
@@ -1428,10 +1421,10 @@ struct FlowNotification {
 }
 
 fn try_enqueue_flow_notification(
-    tx: &mpsc::Sender<FlowNotification>,
+    tx: &Notifier<FlowNotification>,
     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 +1449,7 @@ fn try_enqueue_flow_notification(
     }
 }
 
-fn enqueue_flow_notifications(table_batches: Vec<TableBatch>, tx: &mpsc::Sender<FlowNotification>) {
+fn enqueue_flow_notifications(table_batches: Vec<TableBatch>, tx: &Notifier<FlowNotification>) {
     for table_batch in table_batches {
         let timestamps = extract_timestamps(&table_batch);
         if timestamps.is_empty() {
@@ -1517,14 +1510,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 +1530,8 @@ 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(NonZeroUsize::new(table_batches.len().max(1)).unwrap()).unwrap();
     start_flow_notification_worker(rx, table_flownode_set_cache, node_manager);
     enqueue_flow_notifications(table_batches, &tx);
 }
@@ -1875,8 +1872,11 @@ fn record_batch_to_ipc(record_batch: RecordBatch) -> Result<(Bytes, Bytes, Bytes
 mod tests {
     use std::any::Any;
     use std::collections::{HashMap, HashSet};
+    use std::future::{Future, poll_fn};
+    use std::num::NonZeroUsize;
     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 +1912,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 +1926,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 +2165,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(NonZeroUsize::new(1).unwrap()).unwrap();
+        let notification = |table_id| crate::pending_rows_batcher::FlowNotification {
             table_id,
             timestamps: vec![table_id as i64],
         };
@@ -2178,6 +2178,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 +2381,22 @@ 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), NonZeroUsize::new(1).unwrap())
+                .unwrap(),
+        );
+        pending_flush.submit(
+            FlushWaiter {
+                response_tx,
+                _permit: Arc::new(permit),
+            },
+            1,
+        );
         let mut batch = Some(PendingBatch {
             tables: HashMap::from([(
                 42,
@@ -2390,24 +2407,79 @@ 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), NonZeroUsize::new(2).unwrap())
+                    .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 +2644,8 @@ mod tests {
     fn mock_flow_notification_sender(
         cache: TableFlownodeSetCacheRef,
         node_manager: NodeManagerRef,
-    ) -> mpsc::Sender<super::FlowNotification> {
-        let (tx, rx) = mpsc::channel(16);
+    ) -> Notifier<crate::pending_rows_batcher::FlowNotification> {
+        let (tx, rx) = Notifier::try_new(NonZeroUsize::new(16).unwrap()).unwrap();
         start_flow_notification_worker(rx, cache, node_manager);
         tx
     }
@@ -2930,9 +3002,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(NonZeroUsize::new(1).unwrap()).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 +3029,15 @@ mod tests {
         };
 
         let (tx, _rx) = mpsc::channel::<WorkerCommand>(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 +3047,13 @@ mod tests {
 
         let (stale_tx, _stale_rx) = mpsc::channel::<WorkerCommand>(1);
         let (fresh_tx, _fresh_rx) = mpsc::channel::<WorkerCommand>(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 +3082,7 @@ mod tests {
                 total_rows,
                 ctx: session::context::QueryContext::arc(),
                 response_tx,
-                _permit: permit,
+                _permit: Arc::new(permit),
             })
             .await
             .unwrap();
@@ -3039,7 +3124,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 +3133,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 +3159,11 @@ 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(NonZeroUsize::new(1).unwrap()).unwrap();
         let (shutdown, _) = broadcast::channel(1);
 
+        let flush_limiter = FlushLimiter::try_new(NonZeroUsize::new(1).unwrap()).unwrap();
         start_worker(
             key.clone(),
             worker_tx.clone(),
@@ -3090,10 +3174,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, NonZeroUsize::new(3).unwrap()).unwrap(),
+            flush_limiter.clone(),
         );
 
         // Start the worker, then size-flush a batch halfway to the first
@@ -3101,7 +3184,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 +3195,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 +3215,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"
         );
     }