1use std::collections::HashMap;
16use std::sync::Arc;
17use std::time::Duration;
18
19use catalog::CatalogManagerRef;
20use common_batcher::flush_limiter::FlushLimiter;
21use common_batcher::flush_policy::FlushTrigger;
22use common_batcher::flush_policy::timing::TimingFlushPolicy;
23use common_batcher::pending_worker::PendingWorker as PendingCore;
24use common_batcher::worker_registry::WorkerRegistry;
25use common_meta::node_manager::NodeManagerRef;
26use common_telemetry::debug;
27use partition::manager::PartitionRuleManagerRef;
28use session::context::QueryContextRef;
29use table::metadata::TableId;
30use tokio::sync::{OwnedSemaphorePermit, broadcast, mpsc, oneshot};
31
32use crate::batcher::flow_notifier::FlowNotifier;
33use crate::batcher::logical_table::BatchKey;
34use crate::batcher::logical_table::batch::{Batch, flush_batch_with_managers, spawn_flush};
35use crate::batcher::logical_table::batch_convert::{RecordBatchWithTsIdx, TableBatch};
36use crate::error::Error;
37use crate::metrics::{PENDING_BATCHES, PENDING_ROWS, PENDING_WORKERS};
38
39pub(in crate::batcher::logical_table) struct PendingBatch {
40 pub(in crate::batcher::logical_table) tables: HashMap<TableId, TableBatch>,
41 pub(in crate::batcher::logical_table) db_string: String,
42 pub(in crate::batcher::logical_table) ctx: QueryContextRef,
43}
44
45pub(in crate::batcher::logical_table) struct FlushWaiter {
46 pub(in crate::batcher::logical_table) response_tx:
47 oneshot::Sender<std::result::Result<(), Arc<Error>>>,
48 pub(in crate::batcher::logical_table) _permit: Arc<OwnedSemaphorePermit>,
49}
50
51#[derive(Clone)]
52pub(in crate::batcher::logical_table) struct PendingWorker {
53 pub(in crate::batcher::logical_table) tx: mpsc::Sender<WorkerCommand>,
54}
55
56pub(in crate::batcher::logical_table) enum WorkerCommand {
57 Submit {
58 table_batches: Vec<(String, u32, RecordBatchWithTsIdx)>,
59 total_rows: usize,
60 ctx: QueryContextRef,
61 response_tx: oneshot::Sender<std::result::Result<(), Arc<Error>>>,
62 _permit: Arc<OwnedSemaphorePermit>,
63 },
64 #[cfg(test)]
65 Ack { ack_tx: oneshot::Sender<()> },
66}
67
68impl PendingBatch {
69 pub(in crate::batcher::logical_table) fn new(ctx: QueryContextRef) -> Self {
70 let db_string = ctx.get_db_string();
71 Self {
72 tables: HashMap::new(),
73 db_string,
74 ctx,
75 }
76 }
77
78 pub(in crate::batcher::logical_table) fn add_table_batch(
79 &mut self,
80 table_name: String,
81 table_id: TableId,
82 record_batch: RecordBatchWithTsIdx,
83 ) {
84 let entry = self.tables.entry(table_id).or_insert_with(|| TableBatch {
85 table_name,
86 table_id,
87 batches: Vec::new(),
88 row_count: 0,
89 });
90 entry.row_count += record_batch.batch.num_rows();
91 entry.batches.push(record_batch);
92 }
93}
94
95#[allow(clippy::too_many_arguments)]
96pub(in crate::batcher::logical_table) fn start_worker(
97 key: BatchKey,
98 worker_tx: mpsc::Sender<WorkerCommand>,
99 workers: Arc<WorkerRegistry<BatchKey, WorkerCommand>>,
100 mut rx: mpsc::Receiver<WorkerCommand>,
101 shutdown: broadcast::Sender<()>,
102 partition_manager: PartitionRuleManagerRef,
103 node_manager: NodeManagerRef,
104 catalog_manager: CatalogManagerRef,
105 flow_notification_tx: FlowNotifier,
106 worker_idle_timeout: Duration,
107 flush_policy: TimingFlushPolicy,
108 flush_limiter: FlushLimiter,
109) {
110 tokio::spawn(async move {
111 let mut batch = None;
113 let mut pending_flush = PendingCore::new(flush_policy);
114 let mut shutdown_rx = shutdown.subscribe();
115 let idle_deadline = tokio::time::Instant::now() + worker_idle_timeout;
116 let idle_timer = tokio::time::sleep_until(idle_deadline);
117 tokio::pin!(idle_timer);
118
119 loop {
120 tokio::select! {
121 cmd = rx.recv() => {
122 match cmd {
123 Some(WorkerCommand::Submit { table_batches, total_rows, ctx, response_tx, _permit }) => {
124 let submitted_at = tokio::time::Instant::now();
125 idle_timer.as_mut().reset(submitted_at + worker_idle_timeout);
126
127 pending_flush.submit(
128 FlushWaiter { response_tx, _permit },
129 total_rows,
130 );
131 let pending_batch = batch.get_or_insert_with(||{
132 PENDING_BATCHES.inc();
133 PendingBatch::new(ctx)
134 });
135
136 for (table_name, table_id, record_batch) in table_batches {
137 pending_batch.add_table_batch(table_name, table_id, record_batch);
138 }
139
140 PENDING_ROWS.add(total_rows as i64);
141
142 if let Some(flush) = drain_batch(&mut batch, &mut pending_flush, Some(FlushTrigger::Submission)) {
143 spawn_flush(
144 flush,
145 partition_manager.clone(),
146 node_manager.clone(),
147 catalog_manager.clone(),
148 flow_notification_tx.clone(),
149 flush_limiter.clone(),
150 ).await;
151 }
152 }
153 None => {
154 if let Some(flush) = drain_batch(&mut batch, &mut pending_flush, None) {
155 flush_batch_with_managers(
156 flush,
157 partition_manager.clone(),
158 node_manager.clone(),
159 catalog_manager.clone(),
160 flow_notification_tx.clone(),
161 ).await;
162 }
163 break;
164 }
165 #[cfg(test)]
166 Some(WorkerCommand::Ack { ack_tx }) => {
167 let _ = ack_tx.send(());
168 }
169 }
170 }
171 _ = &mut idle_timer => {
172 if !should_close_worker_on_idle_timeout(
173 pending_flush.total_rows(),
174 rx.len(),
175 ) {
176 idle_timer
177 .as_mut()
178 .reset(tokio::time::Instant::now() + worker_idle_timeout);
179 continue;
180 }
181
182 debug!(
183 "Closing idle pending rows worker due to timeout: catalog={}, schema={}, physical_table={}",
184 key.catalog,
185 key.schema,
186 key.physical_table
187 );
188 break;
189 }
190 _ = pending_flush.wait_flush() => {
191 if let Some(flush) = drain_batch(&mut batch, &mut pending_flush, Some(FlushTrigger::Deadline)) {
192 spawn_flush(
193 flush,
194 partition_manager.clone(),
195 node_manager.clone(),
196 catalog_manager.clone(),
197 flow_notification_tx.clone(),
198 flush_limiter.clone(),
199 ).await;
200 }
201 }
202 _ = shutdown_rx.recv() => {
203 if let Some(flush) = drain_batch(&mut batch, &mut pending_flush, None) {
204 flush_batch_with_managers(
205 flush,
206 partition_manager.clone(),
207 node_manager.clone(),
208 catalog_manager.clone(),
209 flow_notification_tx.clone(),
210 ).await;
211 }
212 break;
213 }
214 }
215 }
216
217 remove_worker_if_same_channel(workers.as_ref(), &key, &worker_tx).await;
218 });
219}
220
221pub(in crate::batcher::logical_table) async fn remove_worker_if_same_channel(
222 workers: &WorkerRegistry<BatchKey, WorkerCommand>,
223 key: &BatchKey,
224 worker_tx: &mpsc::Sender<WorkerCommand>,
225) -> bool {
226 if workers.remove_if_same(key, worker_tx).await {
227 PENDING_WORKERS.set(workers.len().await as i64);
228 true
229 } else {
230 false
231 }
232}
233
234pub(in crate::batcher::logical_table) fn should_close_worker_on_idle_timeout(
235 total_row_count: usize,
236 queued_requests: usize,
237) -> bool {
238 total_row_count == 0 && queued_requests == 0
239}
240
241pub(in crate::batcher::logical_table) fn drain_batch(
244 batch: &mut Option<PendingBatch>,
245 pending_flush: &mut PendingCore<FlushWaiter, TimingFlushPolicy>,
246 trigger: Option<FlushTrigger>,
247) -> Option<Batch> {
248 let total_row_count = pending_flush.total_rows();
249 let waiters = match trigger {
250 Some(trigger) => pending_flush.take_ready(trigger)?,
251 None => pending_flush.take_pending()?,
252 };
253 let batch = batch.take()?;
254
255 if total_row_count == 0 {
256 return None;
257 }
258
259 let table_batches = batch.tables.into_values().collect();
260
261 PENDING_ROWS.sub(total_row_count as i64);
262 PENDING_BATCHES.dec();
263
264 Some(Batch {
265 table_batches,
266 total_row_count,
267 db_string: batch.db_string,
268 ctx: batch.ctx,
269 waiters,
270 })
271}
272
273#[cfg(test)]
274mod tests {
275 use std::collections::HashMap;
276 use std::sync::Arc;
277 use std::time::Duration;
278
279 use catalog::memory::MemoryCatalogManager;
280 use common_batcher::flush_limiter::FlushLimiter;
281 use common_batcher::flush_policy::FlushTrigger;
282 use common_batcher::flush_policy::timing::TimingFlushPolicy;
283 use common_batcher::pending_worker::PendingWorker as PendingCore;
284 use common_batcher::worker_registry::WorkerRegistry;
285 use common_meta::cache::new_table_route_cache;
286 use common_meta::kv_backend::memory::MemoryKvBackend;
287 use common_meta::node_manager::NodeManagerRef;
288 use moka::future::CacheBuilder;
289 use partition::cache::new_partition_info_cache;
290 use partition::manager::PartitionRuleManager;
291 use tokio::sync::{Semaphore, broadcast, mpsc, oneshot};
292 use tokio::time::advance;
293
294 use crate::batcher::flow_notifier::FlowNotifier;
295 use crate::batcher::logical_table::BatchKey;
296 use crate::batcher::logical_table::batch_convert::TableBatch;
297 use crate::batcher::logical_table::pending_worker::{
298 FlushWaiter, PendingBatch, WorkerCommand, drain_batch, remove_worker_if_same_channel,
299 should_close_worker_on_idle_timeout, start_worker,
300 };
301 use crate::batcher::logical_table::test_util::{
302 ConcurrentMockNodeManager, mock_aligned_tag_batch,
303 };
304 use crate::error::Error;
305 use crate::metrics::FLOW_NOTIFICATION_DROPPED;
306
307 #[tokio::test]
308 async fn test_drain_batch_takes_initialized_pending_batch_from_option() {
309 let ctx = session::context::QueryContext::arc();
310 let (response_tx, _response_rx) = oneshot::channel();
311 let permit = Arc::new(Semaphore::new(1)).try_acquire_owned().unwrap();
312 let mut pending_flush =
313 PendingCore::new(TimingFlushPolicy::try_new(Duration::from_secs(10), 1).unwrap());
314 pending_flush.submit(
315 FlushWaiter {
316 response_tx,
317 _permit: Arc::new(permit),
318 },
319 1,
320 );
321 let mut batch = Some(PendingBatch {
322 tables: HashMap::from([(
323 42,
324 TableBatch {
325 table_name: "cpu".to_string(),
326 table_id: 42,
327 batches: vec![mock_aligned_tag_batch("tag1", "host-1", 1000, 1.0)],
328 row_count: 1,
329 },
330 )]),
331 db_string: ctx.get_db_string(),
332 ctx: ctx.clone(),
333 });
334
335 let flush = drain_batch(
336 &mut batch,
337 &mut pending_flush,
338 Some(FlushTrigger::Submission),
339 )
340 .unwrap();
341
342 assert!(batch.is_none());
343 assert!(pending_flush.is_empty());
344 assert_eq!(0, pending_flush.total_rows());
345 assert_eq!(1, flush.waiters.len());
346 assert_eq!(1, flush.total_row_count);
347 assert_eq!(1, flush.table_batches.len());
348 assert_eq!(ctx.get_db_string(), flush.db_string);
349 assert_eq!(ctx.current_catalog(), flush.ctx.current_catalog());
350 }
351
352 #[tokio::test]
353 async fn test_drain_batch_preserves_unready_state_and_clears_zero_rows() {
354 for total_rows in [0, 1] {
355 let mut pending_flush =
356 PendingCore::new(TimingFlushPolicy::try_new(Duration::from_secs(10), 2).unwrap());
357 let mut batch = Some(PendingBatch::new(session::context::QueryContext::arc()));
358 let semaphore = Arc::new(Semaphore::new(1));
359 let (response_tx, mut response_rx) = oneshot::channel();
360 pending_flush.submit(
361 FlushWaiter {
362 response_tx,
363 _permit: Arc::new(semaphore.clone().acquire_owned().await.unwrap()),
364 },
365 total_rows,
366 );
367 assert!(
368 drain_batch(
369 &mut batch,
370 &mut pending_flush,
371 Some(FlushTrigger::Submission)
372 )
373 .is_none()
374 );
375 assert!(batch.is_some());
376 assert!(!pending_flush.is_empty());
377 assert_eq!(total_rows, pending_flush.total_rows());
378 assert_eq!(0, semaphore.available_permits());
379
380 let drained = drain_batch(&mut batch, &mut pending_flush, None);
381 assert!(batch.is_none());
382 assert!(pending_flush.is_empty());
383 assert_eq!(0, pending_flush.total_rows());
384 if total_rows == 0 {
385 assert!(drained.is_none());
386 assert_eq!(1, semaphore.available_permits());
387 assert!(matches!(
388 response_rx.try_recv(),
389 Err(oneshot::error::TryRecvError::Closed)
390 ));
391 } else {
392 let drained = drained.unwrap();
393 assert_eq!(1, drained.total_row_count);
394 assert_eq!(1, drained.waiters.len());
395 assert_eq!(0, semaphore.available_permits());
396 drop(drained);
397 assert_eq!(1, semaphore.available_permits());
398 }
399 }
400 }
401
402 #[test]
403 fn test_pending_batch_keeps_same_name_batches_with_distinct_table_ids() {
404 let ctx = session::context::QueryContext::arc();
405 let mut pending_batch = PendingBatch::new(ctx);
406
407 pending_batch.add_table_batch(
408 "cpu".to_string(),
409 42,
410 mock_aligned_tag_batch("tag1", "host-1", 1000, 1.0),
411 );
412 pending_batch.add_table_batch(
413 "cpu".to_string(),
414 43,
415 mock_aligned_tag_batch("tag1", "host-1", 2000, 2.0),
416 );
417
418 assert_eq!(2, pending_batch.tables.len());
419 assert_eq!(42, pending_batch.tables[&42].table_id);
420 assert_eq!(43, pending_batch.tables[&43].table_id);
421 assert_eq!("cpu", pending_batch.tables[&42].table_name);
422 assert_eq!("cpu", pending_batch.tables[&43].table_name);
423 }
424
425 #[tokio::test]
426 async fn test_remove_worker_if_same_channel_removes_matching_entry() {
427 let workers = WorkerRegistry::new();
428 let key = BatchKey {
429 catalog: "greptime".to_string(),
430 schema: "public".to_string(),
431 physical_table: "phy".to_string(),
432 skip_wal: false,
433 };
434
435 let (tx, _rx) = mpsc::channel::<WorkerCommand>(1);
436 workers.get_or_insert_with(key.clone(), || tx.clone()).await;
437
438 assert!(remove_worker_if_same_channel(&workers, &key, &tx).await);
439 assert!(workers.is_empty().await);
440 }
441
442 #[tokio::test]
443 async fn test_remove_worker_if_same_channel_keeps_newer_entry() {
444 let workers = WorkerRegistry::new();
445 let key = BatchKey {
446 catalog: "greptime".to_string(),
447 schema: "public".to_string(),
448 physical_table: "phy".to_string(),
449 skip_wal: false,
450 };
451
452 let (stale_tx, _stale_rx) = mpsc::channel::<WorkerCommand>(1);
453 let (fresh_tx, _fresh_rx) = mpsc::channel::<WorkerCommand>(1);
454 workers
455 .get_or_insert_with(key.clone(), || fresh_tx.clone())
456 .await;
457
458 assert!(!remove_worker_if_same_channel(&workers, &key, &stale_tx).await);
459 assert!(workers.get(&key).await.is_some());
460 assert!(workers.get(&key).await.unwrap().same_channel(&fresh_tx));
461 }
462
463 #[test]
464 fn test_worker_idle_timeout_close_decision() {
465 assert!(should_close_worker_on_idle_timeout(0, 0));
466 assert!(!should_close_worker_on_idle_timeout(1, 0));
467 assert!(!should_close_worker_on_idle_timeout(0, 1));
468 }
469
470 #[tokio::test(start_paused = true)]
471 async fn test_worker_preserves_first_deadline_and_inline_shutdown() {
472 let flush_interval = Duration::from_secs(10);
473 let worker_idle_timeout = Duration::from_secs(30);
474 let key = BatchKey {
475 catalog: "greptime".to_string(),
476 schema: "public".to_string(),
477 physical_table: "phy".to_string(),
478 skip_wal: false,
479 };
480 let workers = Arc::new(WorkerRegistry::new());
481 let (worker_tx, worker_rx) = mpsc::channel(1);
482 workers
483 .get_or_insert_with(key.clone(), || worker_tx.clone())
484 .await;
485
486 let backend = Arc::new(MemoryKvBackend::default());
487 let table_route_cache = Arc::new(new_table_route_cache(
488 "pending-rows-flush-deadline-routes".to_string(),
489 CacheBuilder::new(1).build(),
490 backend.clone(),
491 ));
492 let partition_info_cache = Arc::new(new_partition_info_cache(
493 "pending-rows-flush-deadline-partitions".to_string(),
494 CacheBuilder::new(1).build(),
495 table_route_cache.clone(),
496 ));
497 let partition_manager = Arc::new(PartitionRuleManager::new(
498 backend,
499 table_route_cache,
500 partition_info_cache,
501 ));
502 let node_manager: NodeManagerRef = Arc::new(ConcurrentMockNodeManager {
503 datanodes: Arc::new(HashMap::new()),
504 });
505 let catalog_manager = MemoryCatalogManager::with_default_setup();
506 let (flow_notification_tx, _flow_notification_rx) =
507 FlowNotifier::try_new(1, FLOW_NOTIFICATION_DROPPED.clone()).unwrap();
508 let (shutdown, _) = broadcast::channel(1);
509
510 let flush_limiter = FlushLimiter::try_new(1).unwrap();
511 start_worker(
512 key.clone(),
513 worker_tx.clone(),
514 workers.clone(),
515 worker_rx,
516 shutdown.clone(),
517 partition_manager,
518 node_manager,
519 catalog_manager,
520 flow_notification_tx,
521 worker_idle_timeout,
522 TimingFlushPolicy::try_new(flush_interval, 3).unwrap(),
523 flush_limiter.clone(),
524 );
525
526 tokio::task::yield_now().await;
530 advance(flush_interval / 2).await;
531 let size_flush_rx = submit_mock_worker_batch(&worker_tx, 3, 1000).await;
532 let size_flush_result =
533 receive_mock_flush_result(size_flush_rx, "row threshold did not flush the first batch")
534 .await;
535 assert_missing_physical_table(size_flush_result);
536
537 advance(flush_interval / 5).await;
540 let mut timed_flush_rx = submit_mock_worker_batch(&worker_tx, 1, 2000).await;
541
542 let first_submission = tokio::time::Instant::now();
543 advance(flush_interval / 2).await;
544 let later_flush_rx = submit_mock_worker_batch(&worker_tx, 1, 3000).await;
545 advance(flush_interval / 2 - Duration::from_millis(1)).await;
546 for _ in 0..10 {
547 tokio::task::yield_now().await;
548 }
549 assert!(matches!(
550 timed_flush_rx.try_recv(),
551 Err(oneshot::error::TryRecvError::Empty)
552 ));
553
554 advance(Duration::from_millis(1)).await;
555 let timed_flush_result = receive_mock_flush_result(
556 timed_flush_rx,
557 "batch was not flushed one interval after its creation",
558 )
559 .await;
560 assert_missing_physical_table(timed_flush_result);
561
562 assert_eq!(
563 first_submission + flush_interval,
564 tokio::time::Instant::now()
565 );
566 assert_missing_physical_table(
567 receive_mock_flush_result(
568 later_flush_rx,
569 "later submission was not included in the timed flush",
570 )
571 .await,
572 );
573
574 let _held_permit = flush_limiter.acquire().await.unwrap();
576 let shutdown_flush_rx = submit_mock_worker_batch(&worker_tx, 1, 4000).await;
577 let shutdown_at = tokio::time::Instant::now();
578 let _ = shutdown.send(());
579 assert_missing_physical_table(
580 receive_mock_flush_result(
581 shutdown_flush_rx,
582 "shutdown incorrectly waited for a flush permit",
583 )
584 .await,
585 );
586 assert_eq!(shutdown_at, tokio::time::Instant::now());
587 for _ in 0..10 {
588 if workers.is_empty().await {
589 break;
590 }
591 tokio::task::yield_now().await;
592 }
593 assert!(
594 workers.is_empty().await,
595 "worker did not exit after shutdown"
596 );
597 }
598
599 const WORKER_TEST_TIMEOUT: Duration = Duration::from_secs(30);
600
601 async fn submit_mock_worker_batch(
602 worker_tx: &mpsc::Sender<WorkerCommand>,
603 total_rows: usize,
604 timestamp: i64,
605 ) -> oneshot::Receiver<std::result::Result<(), Arc<Error>>> {
606 let (response_tx, response_rx) = oneshot::channel();
607 let permit = Arc::new(Semaphore::new(1)).acquire_owned().await.unwrap();
608 worker_tx
609 .send(WorkerCommand::Submit {
610 table_batches: vec![(
611 "cpu".to_string(),
612 42,
613 mock_aligned_tag_batch("tag1", "host-1", timestamp, 1.0),
614 )],
615 total_rows,
616 ctx: session::context::QueryContext::arc(),
617 response_tx,
618 _permit: Arc::new(permit),
619 })
620 .await
621 .unwrap();
622
623 let (ack_tx, ack_rx) = oneshot::channel();
627 worker_tx.send(WorkerCommand::Ack { ack_tx }).await.unwrap();
628 ack_rx
629 .await
630 .expect("worker exited before acking the submitted batch");
631
632 response_rx
633 }
634
635 async fn receive_mock_flush_result(
636 response_rx: oneshot::Receiver<std::result::Result<(), Arc<Error>>>,
637 context: &str,
638 ) -> std::result::Result<(), Arc<Error>> {
639 tokio::time::timeout(WORKER_TEST_TIMEOUT, response_rx)
642 .await
643 .unwrap_or_else(|_| panic!("{context}"))
644 .expect("flush result channel closed without a result")
645 }
646
647 fn assert_missing_physical_table(result: std::result::Result<(), Arc<Error>>) {
648 let err = result.expect_err("the empty catalog should make the flush fail");
649 assert!(
650 matches!(
651 err.as_ref(),
652 Error::Internal { err_msg }
653 if err_msg.contains("not found during pending flush")
654 ),
655 "unexpected flush error: {err}"
656 );
657 }
658}