servers/batcher/
table.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15//! Table-scoped RecordBatch accumulation before bulk routing and encoding.
16
17mod batch;
18mod flow_notifier;
19mod metrics;
20mod pending_batch;
21mod pending_worker;
22
23use std::num::NonZeroUsize;
24use std::sync::Arc;
25use std::time::Duration;
26
27use arrow::record_batch::RecordBatch;
28use async_trait::async_trait;
29use common_batcher::flush_limiter::FlushLimiter;
30use common_batcher::flush_policy::timing::TimingFlushPolicy;
31use common_batcher::request_limiter::RequestLimiter;
32use common_batcher::worker_registry::WorkerRegistry;
33use common_error::ext::BoxedError;
34use common_query::{Output, WriteCompletion, WriteProgress};
35use operator::batcher::PendingRowsBatcher;
36use operator::error::{ExternalSnafu, Result, UnexpectedSnafu};
37use operator::insert::Inserter;
38use session::context::QueryContextRef;
39use snafu::ResultExt;
40use table::metadata::TableInfoRef;
41use tokio::sync::{OwnedSemaphorePermit, Semaphore, broadcast, oneshot};
42
43use crate::batcher::pending_rows_batch_sync_enabled;
44use crate::batcher::table::flow_notifier::FlowNotifier;
45use crate::batcher::table::metrics::PENDING_WORKERS;
46use crate::batcher::table::pending_batch::PendingBatch;
47use crate::batcher::table::pending_worker::{PendingWorker, WorkerCommand, start_worker};
48
49#[derive(Clone, Debug, Hash, PartialEq, Eq)]
50struct BatchKey {
51    catalog: String,
52    schema: String,
53    table_name: String,
54    skip_wal: bool,
55}
56
57// Match Prom batching: group by table name and keep WAL policies separate.
58fn batch_key_from_ctx(table_name: &str, ctx: &QueryContextRef) -> BatchKey {
59    BatchKey {
60        catalog: ctx.current_catalog().to_string(),
61        schema: ctx.current_schema().clone(),
62        table_name: table_name.to_string(),
63        skip_wal: ctx.skip_wal(),
64    }
65}
66
67type PendingWorkers = WorkerRegistry<BatchKey, WorkerCommand>;
68
69/// Admission and worker lookup for table-scoped bulk writes.
70pub struct TablePendingRowsBatcher {
71    workers: Arc<PendingWorkers>,
72    flush_policy: TimingFlushPolicy,
73    flush_limiter: FlushLimiter,
74    request_limiter: RequestLimiter,
75    worker_channel_capacity: usize,
76    pending_rows_batch_sync: bool,
77    worker_idle_timeout: Duration,
78    inserter: Arc<Inserter>,
79    flow_notifier: FlowNotifier,
80    shutdown: broadcast::Sender<()>,
81}
82
83impl TablePendingRowsBatcher {
84    /// Creates a shared timing batcher, rejecting disabled or unsupported limits.
85    pub fn try_new(
86        flush_interval: Duration,
87        max_batch_rows: usize,
88        max_concurrent_flushes: usize,
89        worker_channel_capacity: usize,
90        max_inflight_requests: usize,
91        flow_notification_queue_capacity: NonZeroUsize,
92        inserter: Arc<Inserter>,
93    ) -> Option<Arc<Self>> {
94        if worker_channel_capacity == 0
95            || worker_channel_capacity > Semaphore::MAX_PERMITS
96            || max_inflight_requests == 0
97            || max_inflight_requests > Semaphore::MAX_PERMITS
98        {
99            return None;
100        }
101        let flush_policy = TimingFlushPolicy::try_new(flush_interval, max_batch_rows)?;
102        let flush_limiter = FlushLimiter::try_new(max_concurrent_flushes)?;
103        let request_limiter = RequestLimiter::try_new(max_inflight_requests)?;
104        let flow_notifier = FlowNotifier::new(
105            inserter.table_flownode_set_cache().clone(),
106            inserter.node_manager().clone(),
107            flow_notification_queue_capacity,
108        )?;
109        let (shutdown, _) = broadcast::channel(1);
110        PENDING_WORKERS.set(0);
111        Some(Arc::new(Self {
112            workers: Arc::new(WorkerRegistry::new()),
113            flush_policy,
114            flush_limiter,
115            request_limiter,
116            worker_channel_capacity,
117            pending_rows_batch_sync: pending_rows_batch_sync_enabled(),
118            worker_idle_timeout: flush_interval.checked_mul(3).unwrap_or(flush_interval),
119            inserter,
120            flow_notifier,
121            shutdown,
122        }))
123    }
124
125    async fn worker(&self, key: &BatchKey) -> PendingWorker {
126        let (tx, receiver) = self
127            .workers
128            .get_or_create(key.clone(), self.worker_channel_capacity)
129            .await;
130        if let Some(rx) = receiver {
131            start_worker(
132                key.clone(),
133                tx.clone(),
134                self.workers.clone(),
135                rx,
136                self.shutdown.clone(),
137                self.flush_policy,
138                self.flush_limiter.clone(),
139                self.inserter.clone(),
140                self.flow_notifier.clone(),
141                self.worker_idle_timeout,
142            );
143            PENDING_WORKERS.set(self.workers.len().await as i64);
144        }
145        PendingWorker { tx }
146    }
147}
148
149impl Drop for TablePendingRowsBatcher {
150    fn drop(&mut self) {
151        let _ = self.shutdown.send(());
152    }
153}
154
155#[async_trait]
156impl PendingRowsBatcher for TablePendingRowsBatcher {
157    /// Acquires once per original request, before submitting any of its tables.
158    /// Retain a clone in every submission until its flush has completed.
159    async fn acquire(&self) -> Result<Arc<OwnedSemaphorePermit>> {
160        self.request_limiter.acquire().await.map_err(|_| {
161            UnexpectedSnafu {
162                violated: "batcher admission closed".to_string(),
163            }
164            .build()
165        })
166    }
167
168    /// Acknowledges according to the global batching policy. Cancellation does
169    /// not retract an admitted submission. Flush failures may be partial writes.
170    async fn submit(
171        &self,
172        table_info: TableInfoRef,
173        batch: RecordBatch,
174        ctx: QueryContextRef,
175        permit: Arc<OwnedSemaphorePermit>,
176    ) -> Result<Output> {
177        let total_rows = batch.num_rows();
178        let confirmed = ctx.confirmed_writes();
179        if total_rows == 0 {
180            return Ok(Output::new_with_affected_rows(0));
181        }
182        if table_info.catalog_name != ctx.current_catalog()
183            || table_info.schema_name != ctx.current_schema()
184        {
185            return UnexpectedSnafu {
186                violated: "batch table and request database differ".to_string(),
187            }
188            .fail();
189        }
190        let key = batch_key_from_ctx(&table_info.name, &ctx);
191        let (response_tx, response_rx) = oneshot::channel();
192        let progress = ctx.write_progress().cloned();
193        let pending = PendingBatch {
194            _completion: progress.as_ref().map(|progress| progress.start()),
195            table_info,
196            batch,
197            ctx,
198            response_tx,
199            _permit: permit,
200        };
201        let mut command = Some(WorkerCommand::Submit(pending));
202        for _ in 0..2 {
203            let worker = self.worker(&key).await;
204            let Some(pending) = command.take() else { break };
205            match worker.tx.send(pending).await {
206                Ok(()) => break,
207                Err(error) => {
208                    command = Some(error.0);
209                    if self.workers.remove_if_same(&key, &worker.tx).await {
210                        PENDING_WORKERS.set(self.workers.len().await as i64);
211                    }
212                }
213            }
214        }
215        if command.is_some() {
216            return UnexpectedSnafu {
217                violated: "batch worker channel closed".to_string(),
218            }
219            .fail();
220        }
221        let completion = async move {
222            response_rx.await.unwrap_or_else(|_| WriteProgress::failed(BoxedError::new(
223                UnexpectedSnafu { violated: "batch worker stopped before acknowledging its write; storage outcome unknown".to_string() }.build()
224            )))
225        };
226        let mut output = Output::new_with_affected_rows(total_rows);
227        if self.pending_rows_batch_sync && !confirmed {
228            completion.await.into_result().context(ExternalSnafu)?;
229        } else {
230            if let Some(progress) = progress {
231                progress.mark_accepted();
232            }
233            output
234                .meta
235                .write_completions
236                .push(WriteCompletion::with_progress(completion));
237        }
238        Ok(output)
239    }
240}
241
242#[cfg(test)]
243mod tests {
244    use std::collections::HashMap;
245    use std::num::NonZeroUsize;
246    use std::sync::atomic::{AtomicUsize, Ordering};
247    use std::sync::{Arc, Mutex};
248    use std::time::Duration;
249
250    use api::region::RegionResponse;
251    use api::v1::helper::{tag_column_schema, time_index_column_schema};
252    use api::v1::region::{RegionRequest, bulk_insert_request, region_request};
253    use api::v1::value::ValueData;
254    use api::v1::{ColumnDataType, Row, Rows, Value};
255    use arrow::array::{ArrayRef, Int32Array, TimestampMillisecondArray};
256    use arrow::datatypes::Schema as ArrowSchema;
257    use arrow::record_batch::RecordBatch;
258    use catalog::memory::MemoryCatalogManager;
259    use common_batcher::flush_limiter::FlushLimiter;
260    use common_batcher::flush_policy::timing::TimingFlushPolicy;
261    use common_batcher::request_limiter::RequestLimiter;
262    use common_batcher::worker_registry::WorkerRegistry;
263    use common_catalog::consts::default_engine;
264    use common_grpc::flight::FlightDecoder;
265    use common_meta::error::Result as MetaResult;
266    use common_meta::peer::Peer;
267    use common_meta::test_util::{MockDatanodeHandler, MockDatanodeManager};
268    use common_query::OutputData;
269    use common_query::request::QueryRequest;
270    use common_recordbatch::SendableRecordBatchStream;
271    use common_telemetry::info;
272    use datatypes::schema::{ColumnDefaultConstraint, SchemaBuilder};
273    use datatypes::value::Value as DtValue;
274    use datatypes::vectors::{Int32Vector, TimestampMillisecondVector, VectorRef};
275    use operator::batcher::PendingRowsBatcher;
276    use operator::insert::Inserter;
277    use operator::metrics::DIST_INGEST_ROW_COUNT;
278    use operator::req_convert::insert::rows_to_record_batch;
279    use operator::test_util::{
280        create_partition_rule_manager, new_test_table_info, prepare_mocked_backend,
281    };
282    use session::context::{Channel, QueryContext};
283    use store_api::storage::RegionId;
284    use table::dist_table::DistTable;
285    use table::metadata::TableInfoRef;
286    use table::requests::InsertRequest;
287    use tokio::sync::{broadcast, mpsc, oneshot};
288    use tokio::time::timeout;
289
290    use crate::batcher::table::flow_notifier::FlowNotifier;
291    use crate::batcher::table::pending_batch::PendingBatch;
292    use crate::batcher::table::pending_worker::{WorkerCommand, start_worker};
293    use crate::batcher::table::{PendingWorkers, TablePendingRowsBatcher, batch_key_from_ctx};
294    use crate::batcher::test_util::mock_table_flownode_cache;
295
296    #[derive(Clone)]
297    struct BulkHandler {
298        requests: Arc<Mutex<Vec<RecordBatch>>>,
299        report_missing_row: bool,
300        expected_skip_wal: bool,
301        expected_schema: String,
302    }
303
304    #[async_trait::async_trait]
305    impl MockDatanodeHandler for BulkHandler {
306        async fn handle(&self, peer: &Peer, request: RegionRequest) -> MetaResult<RegionResponse> {
307            assert_eq!(3, peer.id);
308            assert_eq!(self.expected_schema, request.header.unwrap().dbname);
309            let Some(region_request::Body::BulkInsert(request)) = request.body else {
310                panic!("expected bulk insert")
311            };
312            assert_eq!(RegionId::new(1, 3).as_u64(), request.region_id);
313            assert_eq!(self.expected_skip_wal, request.skip_wal);
314            assert!(request.partition_expr_version.is_some());
315            let Some(bulk_insert_request::Body::ArrowIpc(ipc)) = request.body else {
316                panic!("expected Arrow IPC")
317            };
318            let batch = FlightDecoder::try_from_schema_bytes(&ipc.schema)
319                .unwrap()
320                .try_decode_record_batch(&ipc.data_header, &ipc.payload)
321                .unwrap();
322            let rows = batch.num_rows();
323            self.requests.lock().unwrap().push(batch);
324            Ok(RegionResponse::new(
325                rows - usize::from(self.report_missing_row),
326            ))
327        }
328
329        async fn handle_query(
330            &self,
331            _: &Peer,
332            _: QueryRequest,
333        ) -> MetaResult<SendableRecordBatchStream> {
334            panic!("batching must not query the datanode")
335        }
336    }
337
338    #[tokio::test]
339    async fn test_acknowledgement_preserves_admission() {
340        use operator::error::UnexpectedSnafu;
341
342        use crate::batcher::pending_rows_batch_sync_enabled;
343        use crate::batcher::table::batch_key_from_ctx;
344        use crate::batcher::table::pending_batch::notify_batches;
345
346        for sync in [false, true] {
347            for fail in [false, true] {
348                let backend = prepare_mocked_backend().await;
349                let nodes = Arc::new(MockDatanodeManager::new(BulkHandler {
350                    requests: Arc::new(Mutex::new(Vec::new())),
351                    report_missing_row: false,
352                    expected_skip_wal: false,
353                    expected_schema: "public".to_string(),
354                }));
355                let inserter = Arc::new(Inserter::new(
356                    MemoryCatalogManager::new(),
357                    create_partition_rule_manager(backend).await,
358                    nodes,
359                    mock_table_flownode_cache(1, vec![]).await,
360                    true,
361                ));
362                let mut batcher = TablePendingRowsBatcher::try_new(
363                    Duration::from_secs(3600),
364                    1,
365                    1,
366                    1,
367                    1,
368                    NonZeroUsize::new(1).unwrap(),
369                    inserter,
370                )
371                .unwrap();
372                assert_eq!(
373                    batcher.pending_rows_batch_sync,
374                    pending_rows_batch_sync_enabled()
375                );
376                Arc::get_mut(&mut batcher).unwrap().pending_rows_batch_sync = sync;
377                let table = Arc::new(new_test_table_info(1, "ack", [0].into_iter()));
378                let progress = Arc::new(session::write_progress::WriteProgress::default());
379                let mut ctx = QueryContext::with("greptime", "public");
380                ctx.set_write_progress(progress.clone());
381                let ctx = Arc::new(ctx);
382                let key = batch_key_from_ctx(&table.name, &ctx);
383                // Hold the worker command to control completion independently of scheduling.
384                let (_, receiver) = batcher.workers.get_or_create(key, 1).await;
385                let mut receiver = receiver.unwrap();
386                let batch = RecordBatch::try_from_iter(vec![(
387                    "a",
388                    Arc::new(Int32Array::from(vec![1])) as ArrayRef,
389                )])
390                .unwrap();
391                let permit = batcher.acquire().await.unwrap();
392                let submitter = batcher.clone();
393                let submitted =
394                    tokio::spawn(async move { submitter.submit(table, batch, ctx, permit).await });
395                let WorkerCommand::Submit(pending) =
396                    timeout(Duration::from_secs(5), receiver.recv())
397                        .await
398                        .unwrap()
399                        .unwrap();
400                let mut submitted = Some(submitted);
401                let mut completion = None;
402                if sync {
403                    assert!(!submitted.as_ref().unwrap().is_finished());
404                } else {
405                    let mut output = timeout(Duration::from_secs(5), submitted.take().unwrap())
406                        .await
407                        .unwrap()
408                        .unwrap()
409                        .unwrap();
410                    assert_eq!(output.extract_rows_and_cost().0, 1);
411                    assert_eq!(output.meta.write_completions.len(), 1);
412                    let mut pending = Box::pin(output.meta.write_completions.pop().unwrap().wait());
413                    assert!(futures::poll!(pending.as_mut()).is_pending());
414                    completion = Some(pending);
415                }
416                assert_eq!(progress.accepted(), !sync);
417                assert_eq!(progress.written(), 0);
418                // Early acknowledgement must not release capacity before completion.
419                let acquire = batcher.acquire();
420                tokio::pin!(acquire);
421                assert!(futures::poll!(acquire.as_mut()).is_pending());
422                let result = if fail {
423                    Err(Arc::new(
424                        UnexpectedSnafu {
425                            violated: "flush failed".to_string(),
426                        }
427                        .build(),
428                    ))
429                } else {
430                    Ok(())
431                };
432                notify_batches(vec![pending], result);
433                if let Some(submitted) = submitted {
434                    let result = timeout(Duration::from_secs(5), submitted)
435                        .await
436                        .unwrap()
437                        .unwrap();
438                    assert_eq!(result.is_err(), fail);
439                    if !fail {
440                        let output = result.unwrap();
441                        assert_eq!(output.extract_rows_and_cost().0, 1);
442                        assert!(output.meta.write_completions.is_empty());
443                    }
444                }
445                if let Some(completion) = completion {
446                    let result = timeout(Duration::from_secs(5), completion).await.unwrap();
447                    assert_eq!(result.is_err(), fail);
448                }
449                timeout(Duration::from_secs(5), acquire)
450                    .await
451                    .unwrap()
452                    .unwrap();
453            }
454        }
455    }
456
457    fn rows(value: i32) -> Rows {
458        Rows {
459            schema: vec![
460                tag_column_schema("a", ColumnDataType::Int32),
461                time_index_column_schema("ts", ColumnDataType::TimestampMillisecond),
462            ],
463            rows: vec![Row {
464                values: vec![
465                    Value {
466                        value_data: Some(ValueData::I32Value(value)),
467                    },
468                    Value {
469                        value_data: Some(ValueData::TimestampMillisecondValue(
470                            1000 + i64::from(value),
471                        )),
472                    },
473                ],
474            }],
475        }
476    }
477
478    async fn run_bulk_case(
479        max_batch_rows: usize,
480        report_missing_row: bool,
481        shared_request: bool,
482    ) -> usize {
483        run_bulk_case_with_wal(max_batch_rows, report_missing_row, shared_request, false).await
484    }
485
486    async fn run_bulk_case_with_wal(
487        max_batch_rows: usize,
488        report_missing_row: bool,
489        shared_request: bool,
490        skip_wal: bool,
491    ) -> usize {
492        let backend = prepare_mocked_backend().await;
493        let partitions = create_partition_rule_manager(backend.clone()).await;
494        let captured = Arc::new(Mutex::new(Vec::new()));
495        // Isolate the database counter across concurrent test cases.
496        static NEXT_DATABASE: AtomicUsize = AtomicUsize::new(0);
497        let schema_name = format!(
498            "ingest_count_{}",
499            NEXT_DATABASE.fetch_add(1, Ordering::Relaxed)
500        );
501        let nodes = Arc::new(MockDatanodeManager::new(BulkHandler {
502            requests: captured.clone(),
503            report_missing_row,
504            expected_skip_wal: skip_wal,
505            expected_schema: schema_name.clone(),
506        }));
507        let inserter = Arc::new(Inserter::new(
508            MemoryCatalogManager::new(),
509            partitions.clone(),
510            nodes.clone(),
511            mock_table_flownode_cache(1, vec![]).await,
512            true,
513        ));
514        let mut table = new_test_table_info(1, "table_1", [1, 2, 3].into_iter());
515        table.schema_name = schema_name;
516        table.meta.engine = default_engine().to_string();
517        let mut columns = table.meta.schema.column_schemas().to_vec();
518        columns[2] = columns[2]
519            .clone()
520            .with_default_constraint(Some(ColumnDefaultConstraint::Value(DtValue::Int32(7))))
521            .unwrap();
522        table.meta.schema = Arc::new(
523            SchemaBuilder::try_from(columns)
524                .unwrap()
525                .version(123)
526                .build()
527                .unwrap(),
528        );
529        let table = Arc::new(table);
530        let first = rows_to_record_batch(&rows(1), &table).unwrap();
531        let second = rows_to_record_batch(&rows(2), &table).unwrap();
532        // Keep timer/concurrency/workload fixed; only the supplied row threshold
533        // determines whether these two submissions share a bulk request.
534        let batcher: Arc<dyn PendingRowsBatcher> = TablePendingRowsBatcher::try_new(
535            Duration::from_secs(3600),
536            max_batch_rows,
537            1,
538            4,
539            if shared_request { 1 } else { 4 },
540            NonZeroUsize::new(16).unwrap(),
541            inserter,
542        )
543        .unwrap();
544        let context = |channel| {
545            let mut ctx = QueryContext::with_channel("greptime", &table.schema_name, channel);
546            ctx.set_batching_enabled(true);
547            Arc::new(ctx)
548        };
549        let influx_ctx = context(Channel::Influx);
550        let opentsdb_ctx = context(Channel::Opentsdb);
551        let ingest_count =
552            DIST_INGEST_ROW_COUNT.with_label_values(&[influx_ctx.get_db_string().as_str()]);
553        assert_eq!(0, ingest_count.get());
554        influx_ctx.set_skip_wal(skip_wal);
555        opentsdb_ctx.set_skip_wal(skip_wal);
556        let (first_result, second_result) = timeout(Duration::from_secs(5), async {
557            if shared_request {
558                let permit = batcher.acquire().await.unwrap();
559                let (first, second) = tokio::join!(
560                    batcher.submit(table.clone(), first, influx_ctx, permit.clone()),
561                    batcher.submit(table.clone(), second, opentsdb_ctx, permit),
562                );
563                (
564                    first.map(|output| output.extract_rows_and_cost().0),
565                    second.map(|output| output.extract_rows_and_cost().0),
566                )
567            } else {
568                let inserter = Inserter::new(
569                    MemoryCatalogManager::new_with_table(DistTable::table(table.clone())),
570                    partitions,
571                    nodes,
572                    mock_table_flownode_cache(1, vec![]).await,
573                    true,
574                )
575                .with_pending_rows_batcher(Some(batcher));
576                let request = |value: i32| InsertRequest {
577                    catalog_name: table.catalog_name.clone(),
578                    schema_name: table.schema_name.clone(),
579                    table_name: table.name.clone(),
580                    columns_values: HashMap::from([
581                        (
582                            "a".to_string(),
583                            Arc::new(Int32Vector::from_slice([value])) as VectorRef,
584                        ),
585                        (
586                            "ts".to_string(),
587                            Arc::new(TimestampMillisecondVector::from_slice([
588                                1000 + i64::from(value)
589                            ])) as VectorRef,
590                        ),
591                    ]),
592                    skip_wal,
593                };
594                let (first, second) = tokio::join!(
595                    inserter.handle_table_insert(request(1), influx_ctx),
596                    inserter.handle_table_insert(request(2), opentsdb_ctx),
597                );
598                (
599                    first.map(|output| match output.data {
600                        OutputData::AffectedRows(rows) => rows,
601                        _ => panic!("expected affected rows"),
602                    }),
603                    second.map(|output| match output.data {
604                        OutputData::AffectedRows(rows) => rows,
605                        _ => panic!("expected affected rows"),
606                    }),
607                )
608            }
609        })
610        .await
611        .expect("the row threshold did not dispatch the combined bulk insert");
612        if report_missing_row {
613            assert!(first_result.is_err());
614            assert!(second_result.is_err());
615        } else {
616            assert_eq!(1, first_result.unwrap());
617            assert_eq!(1, second_result.unwrap());
618        }
619        assert_eq!(if report_missing_row { 0 } else { 2 }, ingest_count.get());
620        let requests = captured.lock().unwrap();
621        let mut actual_rows = Vec::new();
622        for batch in requests.iter() {
623            assert_eq!(
624                vec!["a", "ts", "b"],
625                batch
626                    .schema()
627                    .fields()
628                    .iter()
629                    .map(|field| field.name().as_str())
630                    .collect::<Vec<_>>()
631            );
632            let a = batch
633                .column(0)
634                .as_any()
635                .downcast_ref::<Int32Array>()
636                .unwrap();
637            let ts = batch
638                .column(1)
639                .as_any()
640                .downcast_ref::<TimestampMillisecondArray>()
641                .unwrap();
642            let b = batch
643                .column(2)
644                .as_any()
645                .downcast_ref::<Int32Array>()
646                .unwrap();
647            actual_rows.extend(
648                (0..batch.num_rows())
649                    .map(|index| (a.value(index), ts.value(index), b.value(index))),
650            );
651        }
652        actual_rows.sort_unstable();
653        assert_eq!(vec![(1, 1001, 7), (2, 1002, 7)], actual_rows);
654        requests.len()
655    }
656
657    #[tokio::test]
658    async fn test_inserter_counts_batched_rows_once() {
659        for report_missing_row in [false, true] {
660            assert_eq!(1, run_bulk_case(2, report_missing_row, false).await);
661        }
662    }
663
664    #[tokio::test]
665    async fn test_row_threshold_ablation_reduces_bulk_requests() {
666        for repetition in 0..2 {
667            let direct_count = run_bulk_case(1, false, false).await;
668            let combined_count = run_bulk_case(2, false, false).await;
669            assert_eq!(2, direct_count);
670            assert_eq!(1, combined_count);
671            info!(
672                "repetition={repetition} rows=2 threshold=1 bulk_requests={direct_count}; threshold=2 bulk_requests={combined_count}"
673            );
674        }
675    }
676
677    #[tokio::test]
678    async fn test_original_request_shares_one_admission_slot() {
679        assert_eq!(1, run_bulk_case(2, false, true).await);
680    }
681
682    #[tokio::test]
683    async fn test_bulk_preserves_skip_wal() {
684        for skip_wal in [false, true] {
685            assert_eq!(1, run_bulk_case_with_wal(2, false, false, skip_wal).await);
686        }
687    }
688
689    #[tokio::test]
690    async fn test_exemplar_coalesced_table_flush_retains_each_requests_progress() {
691        #[derive(Clone)]
692        struct PartialHandler;
693        #[async_trait::async_trait]
694        impl MockDatanodeHandler for PartialHandler {
695            async fn handle(&self, _: &Peer, request: RegionRequest) -> MetaResult<RegionResponse> {
696                let Some(region_request::Body::BulkInsert(request)) = request.body else {
697                    panic!("expected bulk")
698                };
699                let Some(bulk_insert_request::Body::ArrowIpc(ipc)) = request.body else {
700                    panic!("expected Arrow")
701                };
702                let batch = FlightDecoder::try_from_schema_bytes(&ipc.schema)
703                    .unwrap()
704                    .try_decode_record_batch(&ipc.data_header, &ipc.payload)
705                    .unwrap();
706                let values = batch
707                    .column(0)
708                    .as_any()
709                    .downcast_ref::<Int32Array>()
710                    .unwrap();
711                if values.values().contains(&20) {
712                    return common_meta::error::UnexpectedSnafu {
713                        err_msg: "lost acknowledgement",
714                    }
715                    .fail();
716                }
717                Ok(RegionResponse::new(batch.num_rows()))
718            }
719            async fn handle_query(
720                &self,
721                _: &Peer,
722                _: QueryRequest,
723            ) -> MetaResult<SendableRecordBatchStream> {
724                unreachable!()
725            }
726        }
727        for sync in [false, true] {
728            let partitions = create_partition_rule_manager(prepare_mocked_backend().await).await;
729            let inserter = Arc::new(Inserter::new(
730                MemoryCatalogManager::new(),
731                partitions,
732                Arc::new(MockDatanodeManager::new(PartialHandler)),
733                mock_table_flownode_cache(1, vec![]).await,
734                true,
735            ));
736            let mut batcher = TablePendingRowsBatcher::try_new(
737                Duration::from_secs(3600),
738                3,
739                1,
740                4,
741                4,
742                NonZeroUsize::new(16).unwrap(),
743                inserter,
744            )
745            .unwrap();
746            Arc::get_mut(&mut batcher).unwrap().pending_rows_batch_sync = sync;
747            let table = Arc::new(new_test_table_info(1, "table_1", [1, 2, 3].into_iter()));
748            let mut first = rows(1);
749            first.rows.extend(rows(20).rows);
750            let first = rows_to_record_batch(&first, &table).unwrap();
751            let second = rows_to_record_batch(&rows(2), &table).unwrap();
752            let counters =
753                [0, 1].map(|_| Arc::new(session::write_progress::WriteProgress::default()));
754            let contexts = counters.each_ref().map(|counter| {
755                let mut ctx = QueryContext::arc().with_confirmed_writes();
756                ctx.set_write_progress(counter.clone());
757                Arc::new(ctx)
758            });
759            let first_permit = batcher.acquire().await.unwrap();
760            let second_permit = batcher.acquire().await.unwrap();
761            let (first, second) = tokio::join!(
762                batcher.submit(table.clone(), first, contexts[0].clone(), first_permit),
763                batcher.submit(table, second, contexts[1].clone(), second_permit),
764            );
765            let first = timeout(Duration::from_secs(5), first.unwrap().write_progress())
766                .await
767                .unwrap();
768            let second = timeout(Duration::from_secs(5), second.unwrap().write_progress())
769                .await
770                .unwrap();
771            assert_eq!(first.affected_rows, 1);
772            assert!(first.error.is_some());
773            assert_eq!(second.affected_rows, 1);
774            assert!(second.error.is_none());
775            for counter in counters {
776                counter.wait().await;
777                assert_eq!(counter.written(), 1);
778            }
779        }
780    }
781
782    #[test]
783    fn test_batch_key_uses_name_and_wal_policy() {
784        let ctx = QueryContext::arc();
785        let table = new_test_table_info(1, "table_1", [1, 2, 3].into_iter());
786        let mut replacement = table.clone();
787        replacement.ident.table_id = 2;
788        replacement.ident.version += 1;
789        let key = batch_key_from_ctx(&table.name, &ctx);
790        assert_eq!(key, batch_key_from_ctx(&replacement.name, &ctx));
791        assert_ne!(key, batch_key_from_ctx("table_2", &ctx));
792        for (catalog, schema) in [("other", "public"), ("greptime", "other")] {
793            let other = Arc::new(QueryContext::with_channel(catalog, schema, Channel::Influx));
794            assert_ne!(key, batch_key_from_ctx(&table.name, &other));
795        }
796        ctx.set_skip_wal(true);
797        assert_ne!(key, batch_key_from_ctx(&table.name, &ctx));
798    }
799
800    const WORKER_TIMEOUT: Duration = Duration::from_secs(5);
801
802    #[derive(Clone)]
803    struct GatedHandler {
804        inner: BulkHandler,
805        entered: mpsc::UnboundedSender<oneshot::Sender<()>>,
806    }
807
808    #[async_trait::async_trait]
809    impl MockDatanodeHandler for GatedHandler {
810        async fn handle(&self, peer: &Peer, request: RegionRequest) -> MetaResult<RegionResponse> {
811            let response = self.inner.handle(peer, request).await?;
812            let (release_tx, release_rx) = oneshot::channel();
813            self.entered.send(release_tx).unwrap();
814            release_rx.await.unwrap();
815            Ok(response)
816        }
817
818        async fn handle_query(
819            &self,
820            _: &Peer,
821            _: QueryRequest,
822        ) -> MetaResult<SendableRecordBatchStream> {
823            panic!("worker must not query the datanode")
824        }
825    }
826
827    struct WorkerTest {
828        tx: mpsc::Sender<WorkerCommand>,
829        shutdown: broadcast::Sender<()>,
830        workers: Arc<PendingWorkers>,
831        table: TableInfoRef,
832        entered: mpsc::UnboundedReceiver<oneshot::Sender<()>>,
833    }
834
835    impl Drop for WorkerTest {
836        fn drop(&mut self) {
837            let _ = self.shutdown.send(());
838        }
839    }
840
841    impl WorkerTest {
842        async fn new(max_rows: usize, max_flushes: usize) -> Self {
843            let backend = prepare_mocked_backend().await;
844            let partitions = create_partition_rule_manager(backend.clone()).await;
845            let (entered_tx, entered) = mpsc::unbounded_channel();
846            let nodes = Arc::new(MockDatanodeManager::new(GatedHandler {
847                inner: BulkHandler {
848                    requests: Arc::new(Mutex::new(Vec::new())),
849                    report_missing_row: false,
850                    expected_skip_wal: false,
851                    expected_schema: "public".to_string(),
852                },
853                entered: entered_tx,
854            }));
855            let cache = mock_table_flownode_cache(1, vec![]).await;
856            let notifier =
857                FlowNotifier::new(cache.clone(), nodes.clone(), NonZeroUsize::new(16).unwrap())
858                    .unwrap();
859            let inserter = Arc::new(Inserter::new(
860                MemoryCatalogManager::new(),
861                partitions,
862                nodes,
863                cache,
864                true,
865            ));
866            let table = Arc::new(new_test_table_info(1, "table_1", [1, 2, 3].into_iter()));
867            let key = batch_key_from_ctx(&table.name, &QueryContext::arc());
868            let workers = Arc::new(WorkerRegistry::new());
869            let (tx, rx) = mpsc::channel(1);
870            workers.get_or_insert_with(key.clone(), || tx.clone()).await;
871            let (shutdown, _) = broadcast::channel(1);
872            start_worker(
873                key,
874                tx.clone(),
875                workers.clone(),
876                rx,
877                shutdown.clone(),
878                TimingFlushPolicy::try_new(Duration::from_secs(3600), max_rows).unwrap(),
879                FlushLimiter::try_new(max_flushes).unwrap(),
880                inserter,
881                notifier,
882                Duration::from_secs(10800),
883            );
884            Self {
885                tx,
886                shutdown,
887                workers,
888                table,
889                entered,
890            }
891        }
892
893        async fn submit(
894            &self,
895            count: usize,
896            changed_schema: bool,
897        ) -> oneshot::Receiver<common_query::WriteProgress> {
898            self.submit_with_table(count, changed_schema, self.table.clone())
899                .await
900        }
901
902        async fn submit_with_table(
903            &self,
904            count: usize,
905            changed_schema: bool,
906            table_info: TableInfoRef,
907        ) -> oneshot::Receiver<common_query::WriteProgress> {
908            let mut input = rows(1);
909            input.rows = vec![input.rows[0].clone(); count];
910            let mut batch = rows_to_record_batch(&input, &self.table).unwrap();
911            if changed_schema {
912                let schema = ArrowSchema::new_with_metadata(
913                    batch.schema().fields().clone(),
914                    [("test_version".to_string(), "2".to_string())]
915                        .into_iter()
916                        .collect(),
917                );
918                batch = RecordBatch::try_new(Arc::new(schema), batch.columns().to_vec()).unwrap();
919            }
920            let limiter = RequestLimiter::try_new(1).unwrap();
921            let (response_tx, response_rx) = oneshot::channel();
922            self.tx
923                .send(WorkerCommand::Submit(PendingBatch {
924                    _completion: None,
925                    table_info,
926                    batch,
927                    ctx: QueryContext::arc(),
928                    response_tx,
929                    _permit: limiter.acquire().await.unwrap(),
930                }))
931                .await
932                .unwrap();
933            // With capacity one, this proves the submitted command was dequeued.
934            drop(
935                timeout(WORKER_TIMEOUT, self.tx.reserve())
936                    .await
937                    .unwrap()
938                    .unwrap(),
939            );
940            response_rx
941        }
942
943        async fn entered(&mut self) -> oneshot::Sender<()> {
944            timeout(WORKER_TIMEOUT, self.entered.recv())
945                .await
946                .expect("flush did not reach datanode")
947                .unwrap()
948        }
949
950        async fn stopped(&self) {
951            timeout(WORKER_TIMEOUT, async {
952                while !self.workers.is_empty().await {
953                    tokio::task::yield_now().await;
954                }
955            })
956            .await
957            .expect("worker did not remove its registry entry");
958        }
959    }
960
961    #[tokio::test]
962    async fn test_same_schema_flushes_overlap() {
963        let mut worker = WorkerTest::new(1, 2).await;
964        let first = worker.submit(1, false).await;
965        let first_release = worker.entered().await;
966        let second = worker.submit(1, false).await;
967        let second_release = worker.entered().await;
968        // Both RPCs reached the datanode before either was allowed to complete.
969        first_release.send(()).unwrap();
970        second_release.send(()).unwrap();
971        timeout(WORKER_TIMEOUT, first)
972            .await
973            .unwrap()
974            .unwrap()
975            .into_result()
976            .unwrap();
977        timeout(WORKER_TIMEOUT, second)
978            .await
979            .unwrap()
980            .unwrap()
981            .into_result()
982            .unwrap();
983        worker.shutdown.send(()).unwrap();
984        worker.stopped().await;
985    }
986
987    #[tokio::test]
988    async fn test_schema_change_waits_for_prior_flush() {
989        let mut worker = WorkerTest::new(1, 2).await;
990        let first = worker.submit(1, false).await;
991        let first_release = worker.entered().await;
992        let second = worker.submit(1, true).await;
993        // A second flush permit is free, but the schema barrier must withhold the RPC.
994        assert!(
995            timeout(Duration::from_millis(100), worker.entered.recv())
996                .await
997                .is_err()
998        );
999        first_release.send(()).unwrap();
1000        let second_release = worker.entered().await;
1001        second_release.send(()).unwrap();
1002        timeout(WORKER_TIMEOUT, first)
1003            .await
1004            .unwrap()
1005            .unwrap()
1006            .into_result()
1007            .unwrap();
1008        timeout(WORKER_TIMEOUT, second)
1009            .await
1010            .unwrap()
1011            .unwrap()
1012            .into_result()
1013            .unwrap();
1014        worker.shutdown.send(()).unwrap();
1015        worker.stopped().await;
1016    }
1017
1018    #[tokio::test]
1019    async fn test_table_identity_change_flushes_pending_rows_before_switching() {
1020        for recreated in [false, true] {
1021            let mut worker = WorkerTest::new(2, 2).await;
1022            let first = worker.submit(1, false).await;
1023            let mut changed = (*worker.table).clone();
1024            if recreated {
1025                changed.ident.table_id = 2;
1026            } else {
1027                changed.ident.version += 1;
1028            }
1029            let mut second = worker.submit_with_table(2, false, Arc::new(changed)).await;
1030            // The first batch is below max_rows but must be flushed on identity change.
1031            let first_release = worker.entered().await;
1032            assert!(
1033                timeout(Duration::from_millis(100), worker.entered.recv())
1034                    .await
1035                    .is_err()
1036            );
1037            assert!(
1038                timeout(Duration::from_millis(100), &mut second)
1039                    .await
1040                    .is_err()
1041            );
1042            first_release.send(()).unwrap();
1043            timeout(WORKER_TIMEOUT, first)
1044                .await
1045                .unwrap()
1046                .unwrap()
1047                .into_result()
1048                .unwrap();
1049            if recreated {
1050                // The recreated table has no route in this fixture; its error belongs only to it.
1051                assert!(
1052                    timeout(WORKER_TIMEOUT, second)
1053                        .await
1054                        .unwrap()
1055                        .unwrap()
1056                        .into_result()
1057                        .is_err()
1058                );
1059            } else {
1060                let second_release = worker.entered().await;
1061                second_release.send(()).unwrap();
1062                timeout(WORKER_TIMEOUT, second)
1063                    .await
1064                    .unwrap()
1065                    .unwrap()
1066                    .into_result()
1067                    .unwrap();
1068            }
1069            worker.shutdown.send(()).unwrap();
1070            worker.stopped().await;
1071        }
1072    }
1073
1074    #[tokio::test]
1075    async fn test_shutdown_flushes_inline_without_joining_inflight() {
1076        let mut worker = WorkerTest::new(2, 1).await;
1077        let mut first = worker.submit(2, false).await;
1078        let first_release = worker.entered().await;
1079        let pending = worker.submit(1, false).await;
1080        worker.shutdown.send(()).unwrap();
1081        // The normal flush owns the sole permit. Shutdown must still start pending inline.
1082        let pending_release = worker.entered().await;
1083        pending_release.send(()).unwrap();
1084        timeout(WORKER_TIMEOUT, pending)
1085            .await
1086            .unwrap()
1087            .unwrap()
1088            .into_result()
1089            .unwrap();
1090        worker.stopped().await;
1091        assert!(matches!(
1092            first.try_recv(),
1093            Err(oneshot::error::TryRecvError::Empty)
1094        ));
1095        first_release.send(()).unwrap();
1096        timeout(WORKER_TIMEOUT, first)
1097            .await
1098            .unwrap()
1099            .unwrap()
1100            .into_result()
1101            .unwrap();
1102    }
1103}
servers/batcher/table.rs

servers/batcher/
table.rs
