mirror of
https://github.com/GreptimeTeam/greptimedb.git
synced 2026-09-23 21:55:38 +00:00
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:
Generated
+11
@@ -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",
|
||||
|
||||
@@ -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" }
|
||||
|
||||
@@ -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"] }
|
||||
@@ -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();
|
||||
}
|
||||
}
|
||||
@@ -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());
|
||||
}
|
||||
}
|
||||
@@ -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);
|
||||
}
|
||||
}
|
||||
@@ -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;
|
||||
@@ -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);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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]);
|
||||
}
|
||||
}
|
||||
@@ -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());
|
||||
}
|
||||
}
|
||||
@@ -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);
|
||||
}
|
||||
}
|
||||
@@ -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);
|
||||
}
|
||||
}
|
||||
@@ -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<()> {
|
||||
|
||||
@@ -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]);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
}),
|
||||
|
||||
@@ -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();
|
||||
|
||||
@@ -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())
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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
|
||||
|
||||
@@ -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"
|
||||
);
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user