refactor: reuse common batching components in Prom ingestion (#9114)

* refactor: compose Prom ingestion with common batcher components

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

* refactor: share bulk insert IPC encoding

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

* refactor: centralize batcher worker channel creation

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

* refactor: share timestamp extraction for batched writes

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

* refactor: preserve sharded batcher worker lookup

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

* fix: align batcher worker lookup and configuration validation

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

* refactor: validate batcher limits inside fallible constructors

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

---------

Signed-off-by: WenyXu <wenymedia@gmail.com>
This commit is contained in:
Weny Xu
2026-09-16 09:21:49 +00:00
committed by GitHub
parent d185a8bf76
commit e0915a8380
23 changed files with 1907 additions and 258 deletions
Generated
+11
View File
@@ -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",
+2
View File
@@ -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" }
+19
View File
@@ -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"] }
+73
View File
@@ -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<Semaphore>,
}
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<Self> {
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<OwnedSemaphorePermit, AcquireError> {
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();
}
}
+43
View File
@@ -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;
}
@@ -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<Self> {
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<T>(&self, batch: &PendingBatch<T>) -> 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<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_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());
}
}
+141
View File
@@ -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);
}
}
+27
View File
@@ -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;
+145
View File
@@ -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 capacity is zero or exceeds Tokio's supported capacity.
pub fn try_new(capacity: usize) -> Option<(Self, Receiver<T>)> {
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<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(0).is_none());
assert!(Notifier::<usize>::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);
}
}
}
}
+115
View File
@@ -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]);
}
}
+260
View File
@@ -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<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::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), rows).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());
}
}
+115
View File
@@ -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<Semaphore>,
}
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<Self> {
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<Arc<OwnedSemaphorePermit>, 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<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(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);
}
}
+265
View File
@@ -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<K, T> {
workers: DashMap<K, Sender<T>>,
}
impl<K: Eq + Hash, T> Default for WorkerRegistry<K, T> {
fn default() -> Self {
Self {
workers: DashMap::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
.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<T>, Option<Receiver<T>>) {
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<F>(&self, key: K, create: F) -> Sender<T>
where
F: FnOnce() -> Sender<T>,
{
if let Some(tx) = self.get(&key).await {
return tx;
}
match self.workers.entry(key) {
Entry::Occupied(mut entry) => {
if entry.get().is_closed() {
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<T>) -> 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);
}
}
+74 -1
View File
@@ -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<()> {
+59 -1
View File
@@ -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<i64>) -> 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]);
}
}
}
+1
View File
@@ -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
+17 -1
View File
@@ -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
}),
+48 -8
View File
@@ -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();
+2 -21
View File
@@ -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<Vec<i64>> {
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())
}
+2
View File
@@ -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};
@@ -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<Vec<i64>> {
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)
}
+1
View File
@@ -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
+292 -226
View File
@@ -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<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,
@@ -376,24 +379,21 @@ impl PendingRowsBatcher {
max_inflight_requests: usize,
flow_notification_queue_capacity: NonZeroUsize,
) -> Option<Arc<Self>> {
// Disable the batcher if flush is disabled or configuration is invalid.
// Zero values for these knobs either cause panics (e.g., zero-capacity channels)
// or deadlocks (e.g., semaphores with no permits).
if flush_interval.is_zero()
|| max_batch_rows == 0
|| max_concurrent_flushes == 0
|| worker_channel_capacity == 0
|| max_inflight_requests == 0
{
if worker_channel_capacity == 0 || worker_channel_capacity > Semaphore::MAX_PERMITS {
return None;
}
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<WorkerCommand>,
rx: mpsc::Receiver<WorkerCommand>,
) {
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<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 +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<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 +1085,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 +1301,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 +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<FlowNotification>,
flow_notification_tx: Notifier<FlowNotification>,
) {
let FlushBatch {
table_batches,
@@ -1403,21 +1381,13 @@ fn extract_timestamps(table_batch: &TableBatch) -> Vec<i64> {
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<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 +1426,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 +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<FlushWaiter>, 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<super::FlowNotification> {
let (tx, rx) = mpsc::channel(16);
) -> Notifier<crate::pending_rows_batcher::FlowNotification> {
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::<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 +3004,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 +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"
);
}