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};
35use operator::batcher::PendingRowsBatcher;
36use operator::error::{BatchFlushSnafu, 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        if total_rows == 0 {
179            return Ok(Output::new_with_affected_rows(0));
180        }
181        if table_info.catalog_name != ctx.current_catalog()
182            || table_info.schema_name != ctx.current_schema()
183        {
184            return UnexpectedSnafu {
185                violated: "batch table and request database differ".to_string(),
186            }
187            .fail();
188        }
189        let key = batch_key_from_ctx(&table_info.name, &ctx);
190        let (response_tx, response_rx) = oneshot::channel();
191        let pending = PendingBatch {
192            table_info,
193            batch,
194            ctx,
195            response_tx,
196            _permit: permit,
197        };
198        let mut command = Some(WorkerCommand::Submit(pending));
199        for _ in 0..2 {
200            let worker = self.worker(&key).await;
201            let Some(pending) = command.take() else { break };
202            match worker.tx.send(pending).await {
203                Ok(()) => break,
204                Err(error) => {
205                    command = Some(error.0);
206                    if self.workers.remove_if_same(&key, &worker.tx).await {
207                        PENDING_WORKERS.set(self.workers.len().await as i64);
208                    }
209                }
210            }
211        }
212        if command.is_some() {
213            return UnexpectedSnafu {
214                violated: "batch worker channel closed".to_string(),
215            }
216            .fail();
217        }
218        let completion = async move {
219            response_rx
220                .await
221                .map_err(|_| {
222                    UnexpectedSnafu {
223                        violated: "batch worker stopped before reporting its write result"
224                            .to_string(),
225                    }
226                    .build()
227                })?
228                .context(BatchFlushSnafu)
229        };
230        let mut output = Output::new_with_affected_rows(total_rows);
231        if self.pending_rows_batch_sync {
232            completion.await?;
233        } else {
234            output
235                .meta
236                .write_completions
237                .push(WriteCompletion::new(async move {
238                    completion.await.map_err(BoxedError::new)
239                }));
240        }
241        Ok(output)
242    }
243}
244
245#[cfg(test)]
246mod tests {
247    use std::collections::HashMap;
248    use std::num::NonZeroUsize;
249    use std::sync::atomic::{AtomicUsize, Ordering};
250    use std::sync::{Arc, Mutex};
251    use std::time::Duration;
252
253    use api::region::RegionResponse;
254    use api::v1::helper::{tag_column_schema, time_index_column_schema};
255    use api::v1::region::{RegionRequest, bulk_insert_request, region_request};
256    use api::v1::value::ValueData;
257    use api::v1::{ColumnDataType, Row, Rows, Value};
258    use arrow::array::{ArrayRef, Int32Array, TimestampMillisecondArray};
259    use arrow::datatypes::Schema as ArrowSchema;
260    use arrow::record_batch::RecordBatch;
261    use catalog::memory::MemoryCatalogManager;
262    use common_batcher::flush_limiter::FlushLimiter;
263    use common_batcher::flush_policy::timing::TimingFlushPolicy;
264    use common_batcher::request_limiter::RequestLimiter;
265    use common_batcher::worker_registry::WorkerRegistry;
266    use common_catalog::consts::default_engine;
267    use common_grpc::flight::FlightDecoder;
268    use common_meta::error::Result as MetaResult;
269    use common_meta::peer::Peer;
270    use common_meta::test_util::{MockDatanodeHandler, MockDatanodeManager};
271    use common_query::OutputData;
272    use common_query::request::QueryRequest;
273    use common_recordbatch::SendableRecordBatchStream;
274    use common_telemetry::info;
275    use datatypes::schema::{ColumnDefaultConstraint, SchemaBuilder};
276    use datatypes::value::Value as DtValue;
277    use datatypes::vectors::{Int32Vector, TimestampMillisecondVector, VectorRef};
278    use operator::batcher::PendingRowsBatcher;
279    use operator::error::Error;
280    use operator::insert::Inserter;
281    use operator::metrics::DIST_INGEST_ROW_COUNT;
282    use operator::req_convert::insert::rows_to_record_batch;
283    use operator::test_util::{
284        create_partition_rule_manager, new_test_table_info, prepare_mocked_backend,
285    };
286    use session::context::{Channel, QueryContext};
287    use store_api::storage::RegionId;
288    use table::dist_table::DistTable;
289    use table::metadata::TableInfoRef;
290    use table::requests::InsertRequest;
291    use tokio::sync::{broadcast, mpsc, oneshot};
292    use tokio::time::timeout;
293
294    use crate::batcher::table::flow_notifier::FlowNotifier;
295    use crate::batcher::table::pending_batch::PendingBatch;
296    use crate::batcher::table::pending_worker::{WorkerCommand, start_worker};
297    use crate::batcher::table::{PendingWorkers, TablePendingRowsBatcher, batch_key_from_ctx};
298    use crate::batcher::test_util::mock_table_flownode_cache;
299
300    #[derive(Clone)]
301    struct BulkHandler {
302        requests: Arc<Mutex<Vec<RecordBatch>>>,
303        report_missing_row: bool,
304        expected_skip_wal: bool,
305        expected_schema: String,
306    }
307
308    #[async_trait::async_trait]
309    impl MockDatanodeHandler for BulkHandler {
310        async fn handle(&self, peer: &Peer, request: RegionRequest) -> MetaResult<RegionResponse> {
311            assert_eq!(3, peer.id);
312            assert_eq!(self.expected_schema, request.header.unwrap().dbname);
313            let Some(region_request::Body::BulkInsert(request)) = request.body else {
314                panic!("expected bulk insert")
315            };
316            assert_eq!(RegionId::new(1, 3).as_u64(), request.region_id);
317            assert_eq!(self.expected_skip_wal, request.skip_wal);
318            assert!(request.partition_expr_version.is_some());
319            let Some(bulk_insert_request::Body::ArrowIpc(ipc)) = request.body else {
320                panic!("expected Arrow IPC")
321            };
322            let batch = FlightDecoder::try_from_schema_bytes(&ipc.schema)
323                .unwrap()
324                .try_decode_record_batch(&ipc.data_header, &ipc.payload)
325                .unwrap();
326            let rows = batch.num_rows();
327            self.requests.lock().unwrap().push(batch);
328            Ok(RegionResponse::new(
329                rows - usize::from(self.report_missing_row),
330            ))
331        }
332
333        async fn handle_query(
334            &self,
335            _: &Peer,
336            _: QueryRequest,
337        ) -> MetaResult<SendableRecordBatchStream> {
338            panic!("batching must not query the datanode")
339        }
340    }
341
342    #[tokio::test]
343    async fn test_acknowledgement_preserves_admission() {
344        use operator::error::UnexpectedSnafu;
345
346        use crate::batcher::pending_rows_batch_sync_enabled;
347        use crate::batcher::table::batch_key_from_ctx;
348        use crate::batcher::table::pending_batch::notify_batches;
349
350        for sync in [false, true] {
351            for fail in [false, true] {
352                let backend = prepare_mocked_backend().await;
353                let nodes = Arc::new(MockDatanodeManager::new(BulkHandler {
354                    requests: Arc::new(Mutex::new(Vec::new())),
355                    report_missing_row: false,
356                    expected_skip_wal: false,
357                    expected_schema: "public".to_string(),
358                }));
359                let inserter = Arc::new(Inserter::new(
360                    MemoryCatalogManager::new(),
361                    create_partition_rule_manager(backend).await,
362                    nodes,
363                    mock_table_flownode_cache(1, vec![]).await,
364                    true,
365                ));
366                let mut batcher = TablePendingRowsBatcher::try_new(
367                    Duration::from_secs(3600),
368                    1,
369                    1,
370                    1,
371                    1,
372                    NonZeroUsize::new(1).unwrap(),
373                    inserter,
374                )
375                .unwrap();
376                assert_eq!(
377                    batcher.pending_rows_batch_sync,
378                    pending_rows_batch_sync_enabled()
379                );
380                Arc::get_mut(&mut batcher).unwrap().pending_rows_batch_sync = sync;
381                let table = Arc::new(new_test_table_info(1, "ack", [0].into_iter()));
382                let ctx = QueryContext::arc();
383                let key = batch_key_from_ctx(&table.name, &ctx);
384                // Hold the worker command to control completion independently of scheduling.
385                let (_, receiver) = batcher.workers.get_or_create(key, 1).await;
386                let mut receiver = receiver.unwrap();
387                let batch = RecordBatch::try_from_iter(vec![(
388                    "a",
389                    Arc::new(Int32Array::from(vec![1])) as ArrayRef,
390                )])
391                .unwrap();
392                let permit = batcher.acquire().await.unwrap();
393                let submitter = batcher.clone();
394                let submitted =
395                    tokio::spawn(async move { submitter.submit(table, batch, ctx, permit).await });
396                let WorkerCommand::Submit(pending) =
397                    timeout(Duration::from_secs(5), receiver.recv())
398                        .await
399                        .unwrap()
400                        .unwrap();
401                let mut submitted = Some(submitted);
402                let mut completion = None;
403                if sync {
404                    assert!(!submitted.as_ref().unwrap().is_finished());
405                } else {
406                    let mut output = timeout(Duration::from_secs(5), submitted.take().unwrap())
407                        .await
408                        .unwrap()
409                        .unwrap()
410                        .unwrap();
411                    assert_eq!(output.extract_rows_and_cost().0, 1);
412                    assert_eq!(output.meta.write_completions.len(), 1);
413                    let mut pending = Box::pin(output.meta.write_completions.pop().unwrap().wait());
414                    assert!(futures::poll!(pending.as_mut()).is_pending());
415                    completion = Some(pending);
416                }
417                // Early acknowledgement must not release capacity before completion.
418                let acquire = batcher.acquire();
419                tokio::pin!(acquire);
420                assert!(futures::poll!(acquire.as_mut()).is_pending());
421                let result = if fail {
422                    Err(Arc::new(
423                        UnexpectedSnafu {
424                            violated: "flush failed".to_string(),
425                        }
426                        .build(),
427                    ))
428                } else {
429                    Ok(())
430                };
431                notify_batches(vec![pending], result);
432                if let Some(submitted) = submitted {
433                    let result = timeout(Duration::from_secs(5), submitted)
434                        .await
435                        .unwrap()
436                        .unwrap();
437                    assert_eq!(result.is_err(), fail);
438                    if !fail {
439                        let output = result.unwrap();
440                        assert_eq!(output.extract_rows_and_cost().0, 1);
441                        assert!(output.meta.write_completions.is_empty());
442                    }
443                }
444                if let Some(completion) = completion {
445                    let result = timeout(Duration::from_secs(5), completion).await.unwrap();
446                    assert_eq!(result.is_err(), fail);
447                }
448                timeout(Duration::from_secs(5), acquire)
449                    .await
450                    .unwrap()
451                    .unwrap();
452            }
453        }
454    }
455
456    fn rows(value: i32) -> Rows {
457        Rows {
458            schema: vec![
459                tag_column_schema("a", ColumnDataType::Int32),
460                time_index_column_schema("ts", ColumnDataType::TimestampMillisecond),
461            ],
462            rows: vec![Row {
463                values: vec![
464                    Value {
465                        value_data: Some(ValueData::I32Value(value)),
466                    },
467                    Value {
468                        value_data: Some(ValueData::TimestampMillisecondValue(
469                            1000 + i64::from(value),
470                        )),
471                    },
472                ],
473            }],
474        }
475    }
476
477    async fn run_bulk_case(
478        max_batch_rows: usize,
479        report_missing_row: bool,
480        shared_request: bool,
481    ) -> usize {
482        run_bulk_case_with_wal(max_batch_rows, report_missing_row, shared_request, false).await
483    }
484
485    async fn run_bulk_case_with_wal(
486        max_batch_rows: usize,
487        report_missing_row: bool,
488        shared_request: bool,
489        skip_wal: bool,
490    ) -> usize {
491        let backend = prepare_mocked_backend().await;
492        let partitions = create_partition_rule_manager(backend.clone()).await;
493        let captured = Arc::new(Mutex::new(Vec::new()));
494        // Isolate the database counter across concurrent test cases.
495        static NEXT_DATABASE: AtomicUsize = AtomicUsize::new(0);
496        let schema_name = format!(
497            "ingest_count_{}",
498            NEXT_DATABASE.fetch_add(1, Ordering::Relaxed)
499        );
500        let nodes = Arc::new(MockDatanodeManager::new(BulkHandler {
501            requests: captured.clone(),
502            report_missing_row,
503            expected_skip_wal: skip_wal,
504            expected_schema: schema_name.clone(),
505        }));
506        let inserter = Arc::new(Inserter::new(
507            MemoryCatalogManager::new(),
508            partitions.clone(),
509            nodes.clone(),
510            mock_table_flownode_cache(1, vec![]).await,
511            true,
512        ));
513        let mut table = new_test_table_info(1, "table_1", [1, 2, 3].into_iter());
514        table.schema_name = schema_name;
515        table.meta.engine = default_engine().to_string();
516        let mut columns = table.meta.schema.column_schemas().to_vec();
517        columns[2] = columns[2]
518            .clone()
519            .with_default_constraint(Some(ColumnDefaultConstraint::Value(DtValue::Int32(7))))
520            .unwrap();
521        table.meta.schema = Arc::new(
522            SchemaBuilder::try_from(columns)
523                .unwrap()
524                .version(123)
525                .build()
526                .unwrap(),
527        );
528        let table = Arc::new(table);
529        let first = rows_to_record_batch(&rows(1), &table).unwrap();
530        let second = rows_to_record_batch(&rows(2), &table).unwrap();
531        // Keep timer/concurrency/workload fixed; only the supplied row threshold
532        // determines whether these two submissions share a bulk request.
533        let batcher: Arc<dyn PendingRowsBatcher> = TablePendingRowsBatcher::try_new(
534            Duration::from_secs(3600),
535            max_batch_rows,
536            1,
537            4,
538            if shared_request { 1 } else { 4 },
539            NonZeroUsize::new(16).unwrap(),
540            inserter,
541        )
542        .unwrap();
543        let context = |channel| {
544            let mut ctx = QueryContext::with_channel("greptime", &table.schema_name, channel);
545            ctx.set_batching_enabled(true);
546            Arc::new(ctx)
547        };
548        let influx_ctx = context(Channel::Influx);
549        let opentsdb_ctx = context(Channel::Opentsdb);
550        let ingest_count =
551            DIST_INGEST_ROW_COUNT.with_label_values(&[influx_ctx.get_db_string().as_str()]);
552        assert_eq!(0, ingest_count.get());
553        influx_ctx.set_skip_wal(skip_wal);
554        opentsdb_ctx.set_skip_wal(skip_wal);
555        let (first_result, second_result) = timeout(Duration::from_secs(5), async {
556            if shared_request {
557                let permit = batcher.acquire().await.unwrap();
558                let (first, second) = tokio::join!(
559                    batcher.submit(table.clone(), first, influx_ctx, permit.clone()),
560                    batcher.submit(table.clone(), second, opentsdb_ctx, permit),
561                );
562                (
563                    first.map(|output| output.extract_rows_and_cost().0),
564                    second.map(|output| output.extract_rows_and_cost().0),
565                )
566            } else {
567                let inserter = Inserter::new(
568                    MemoryCatalogManager::new_with_table(DistTable::table(table.clone())),
569                    partitions,
570                    nodes,
571                    mock_table_flownode_cache(1, vec![]).await,
572                    true,
573                )
574                .with_pending_rows_batcher(Some(batcher));
575                let request = |value: i32| InsertRequest {
576                    catalog_name: table.catalog_name.clone(),
577                    schema_name: table.schema_name.clone(),
578                    table_name: table.name.clone(),
579                    columns_values: HashMap::from([
580                        (
581                            "a".to_string(),
582                            Arc::new(Int32Vector::from_slice([value])) as VectorRef,
583                        ),
584                        (
585                            "ts".to_string(),
586                            Arc::new(TimestampMillisecondVector::from_slice([
587                                1000 + i64::from(value)
588                            ])) as VectorRef,
589                        ),
590                    ]),
591                    skip_wal,
592                };
593                let (first, second) = tokio::join!(
594                    inserter.handle_table_insert(request(1), influx_ctx),
595                    inserter.handle_table_insert(request(2), opentsdb_ctx),
596                );
597                (
598                    first.map(|output| match output.data {
599                        OutputData::AffectedRows(rows) => rows,
600                        _ => panic!("expected affected rows"),
601                    }),
602                    second.map(|output| match output.data {
603                        OutputData::AffectedRows(rows) => rows,
604                        _ => panic!("expected affected rows"),
605                    }),
606                )
607            }
608        })
609        .await
610        .expect("the row threshold did not dispatch the combined bulk insert");
611        if report_missing_row {
612            assert!(first_result.is_err());
613            assert!(second_result.is_err());
614        } else {
615            assert_eq!(1, first_result.unwrap());
616            assert_eq!(1, second_result.unwrap());
617        }
618        assert_eq!(if report_missing_row { 0 } else { 2 }, ingest_count.get());
619        let requests = captured.lock().unwrap();
620        let mut actual_rows = Vec::new();
621        for batch in requests.iter() {
622            assert_eq!(
623                vec!["a", "ts", "b"],
624                batch
625                    .schema()
626                    .fields()
627                    .iter()
628                    .map(|field| field.name().as_str())
629                    .collect::<Vec<_>>()
630            );
631            let a = batch
632                .column(0)
633                .as_any()
634                .downcast_ref::<Int32Array>()
635                .unwrap();
636            let ts = batch
637                .column(1)
638                .as_any()
639                .downcast_ref::<TimestampMillisecondArray>()
640                .unwrap();
641            let b = batch
642                .column(2)
643                .as_any()
644                .downcast_ref::<Int32Array>()
645                .unwrap();
646            actual_rows.extend(
647                (0..batch.num_rows())
648                    .map(|index| (a.value(index), ts.value(index), b.value(index))),
649            );
650        }
651        actual_rows.sort_unstable();
652        assert_eq!(vec![(1, 1001, 7), (2, 1002, 7)], actual_rows);
653        requests.len()
654    }
655
656    #[tokio::test]
657    async fn test_inserter_counts_batched_rows_once() {
658        for report_missing_row in [false, true] {
659            assert_eq!(1, run_bulk_case(2, report_missing_row, false).await);
660        }
661    }
662
663    #[tokio::test]
664    async fn test_row_threshold_ablation_reduces_bulk_requests() {
665        for repetition in 0..2 {
666            let direct_count = run_bulk_case(1, false, false).await;
667            let combined_count = run_bulk_case(2, false, false).await;
668            assert_eq!(2, direct_count);
669            assert_eq!(1, combined_count);
670            info!(
671                "repetition={repetition} rows=2 threshold=1 bulk_requests={direct_count}; threshold=2 bulk_requests={combined_count}"
672            );
673        }
674    }
675
676    #[tokio::test]
677    async fn test_original_request_shares_one_admission_slot() {
678        assert_eq!(1, run_bulk_case(2, false, true).await);
679    }
680
681    #[tokio::test]
682    async fn test_bulk_preserves_skip_wal() {
683        for skip_wal in [false, true] {
684            assert_eq!(1, run_bulk_case_with_wal(2, false, false, skip_wal).await);
685        }
686    }
687
688    #[test]
689    fn test_batch_key_uses_name_and_wal_policy() {
690        let ctx = QueryContext::arc();
691        let table = new_test_table_info(1, "table_1", [1, 2, 3].into_iter());
692        let mut replacement = table.clone();
693        replacement.ident.table_id = 2;
694        replacement.ident.version += 1;
695        let key = batch_key_from_ctx(&table.name, &ctx);
696        assert_eq!(key, batch_key_from_ctx(&replacement.name, &ctx));
697        assert_ne!(key, batch_key_from_ctx("table_2", &ctx));
698        for (catalog, schema) in [("other", "public"), ("greptime", "other")] {
699            let other = Arc::new(QueryContext::with_channel(catalog, schema, Channel::Influx));
700            assert_ne!(key, batch_key_from_ctx(&table.name, &other));
701        }
702        ctx.set_skip_wal(true);
703        assert_ne!(key, batch_key_from_ctx(&table.name, &ctx));
704    }
705
706    const WORKER_TIMEOUT: Duration = Duration::from_secs(5);
707
708    #[derive(Clone)]
709    struct GatedHandler {
710        inner: BulkHandler,
711        entered: mpsc::UnboundedSender<oneshot::Sender<()>>,
712    }
713
714    #[async_trait::async_trait]
715    impl MockDatanodeHandler for GatedHandler {
716        async fn handle(&self, peer: &Peer, request: RegionRequest) -> MetaResult<RegionResponse> {
717            let response = self.inner.handle(peer, request).await?;
718            let (release_tx, release_rx) = oneshot::channel();
719            self.entered.send(release_tx).unwrap();
720            release_rx.await.unwrap();
721            Ok(response)
722        }
723
724        async fn handle_query(
725            &self,
726            _: &Peer,
727            _: QueryRequest,
728        ) -> MetaResult<SendableRecordBatchStream> {
729            panic!("worker must not query the datanode")
730        }
731    }
732
733    struct WorkerTest {
734        tx: mpsc::Sender<WorkerCommand>,
735        shutdown: broadcast::Sender<()>,
736        workers: Arc<PendingWorkers>,
737        table: TableInfoRef,
738        entered: mpsc::UnboundedReceiver<oneshot::Sender<()>>,
739    }
740
741    impl Drop for WorkerTest {
742        fn drop(&mut self) {
743            let _ = self.shutdown.send(());
744        }
745    }
746
747    impl WorkerTest {
748        async fn new(max_rows: usize, max_flushes: usize) -> Self {
749            let backend = prepare_mocked_backend().await;
750            let partitions = create_partition_rule_manager(backend.clone()).await;
751            let (entered_tx, entered) = mpsc::unbounded_channel();
752            let nodes = Arc::new(MockDatanodeManager::new(GatedHandler {
753                inner: BulkHandler {
754                    requests: Arc::new(Mutex::new(Vec::new())),
755                    report_missing_row: false,
756                    expected_skip_wal: false,
757                    expected_schema: "public".to_string(),
758                },
759                entered: entered_tx,
760            }));
761            let cache = mock_table_flownode_cache(1, vec![]).await;
762            let notifier =
763                FlowNotifier::new(cache.clone(), nodes.clone(), NonZeroUsize::new(16).unwrap())
764                    .unwrap();
765            let inserter = Arc::new(Inserter::new(
766                MemoryCatalogManager::new(),
767                partitions,
768                nodes,
769                cache,
770                true,
771            ));
772            let table = Arc::new(new_test_table_info(1, "table_1", [1, 2, 3].into_iter()));
773            let key = batch_key_from_ctx(&table.name, &QueryContext::arc());
774            let workers = Arc::new(WorkerRegistry::new());
775            let (tx, rx) = mpsc::channel(1);
776            workers.get_or_insert_with(key.clone(), || tx.clone()).await;
777            let (shutdown, _) = broadcast::channel(1);
778            start_worker(
779                key,
780                tx.clone(),
781                workers.clone(),
782                rx,
783                shutdown.clone(),
784                TimingFlushPolicy::try_new(Duration::from_secs(3600), max_rows).unwrap(),
785                FlushLimiter::try_new(max_flushes).unwrap(),
786                inserter,
787                notifier,
788                Duration::from_secs(10800),
789            );
790            Self {
791                tx,
792                shutdown,
793                workers,
794                table,
795                entered,
796            }
797        }
798
799        async fn submit(
800            &self,
801            count: usize,
802            changed_schema: bool,
803        ) -> oneshot::Receiver<Result<(), Arc<Error>>> {
804            self.submit_with_table(count, changed_schema, self.table.clone())
805                .await
806        }
807
808        async fn submit_with_table(
809            &self,
810            count: usize,
811            changed_schema: bool,
812            table_info: TableInfoRef,
813        ) -> oneshot::Receiver<Result<(), Arc<Error>>> {
814            let mut input = rows(1);
815            input.rows = vec![input.rows[0].clone(); count];
816            let mut batch = rows_to_record_batch(&input, &self.table).unwrap();
817            if changed_schema {
818                let schema = ArrowSchema::new_with_metadata(
819                    batch.schema().fields().clone(),
820                    [("test_version".to_string(), "2".to_string())]
821                        .into_iter()
822                        .collect(),
823                );
824                batch = RecordBatch::try_new(Arc::new(schema), batch.columns().to_vec()).unwrap();
825            }
826            let limiter = RequestLimiter::try_new(1).unwrap();
827            let (response_tx, response_rx) = oneshot::channel();
828            self.tx
829                .send(WorkerCommand::Submit(PendingBatch {
830                    table_info,
831                    batch,
832                    ctx: QueryContext::arc(),
833                    response_tx,
834                    _permit: limiter.acquire().await.unwrap(),
835                }))
836                .await
837                .unwrap();
838            // With capacity one, this proves the submitted command was dequeued.
839            drop(
840                timeout(WORKER_TIMEOUT, self.tx.reserve())
841                    .await
842                    .unwrap()
843                    .unwrap(),
844            );
845            response_rx
846        }
847
848        async fn entered(&mut self) -> oneshot::Sender<()> {
849            timeout(WORKER_TIMEOUT, self.entered.recv())
850                .await
851                .expect("flush did not reach datanode")
852                .unwrap()
853        }
854
855        async fn stopped(&self) {
856            timeout(WORKER_TIMEOUT, async {
857                while !self.workers.is_empty().await {
858                    tokio::task::yield_now().await;
859                }
860            })
861            .await
862            .expect("worker did not remove its registry entry");
863        }
864    }
865
866    #[tokio::test]
867    async fn test_same_schema_flushes_overlap() {
868        let mut worker = WorkerTest::new(1, 2).await;
869        let first = worker.submit(1, false).await;
870        let first_release = worker.entered().await;
871        let second = worker.submit(1, false).await;
872        let second_release = worker.entered().await;
873        // Both RPCs reached the datanode before either was allowed to complete.
874        first_release.send(()).unwrap();
875        second_release.send(()).unwrap();
876        timeout(WORKER_TIMEOUT, first)
877            .await
878            .unwrap()
879            .unwrap()
880            .unwrap();
881        timeout(WORKER_TIMEOUT, second)
882            .await
883            .unwrap()
884            .unwrap()
885            .unwrap();
886        worker.shutdown.send(()).unwrap();
887        worker.stopped().await;
888    }
889
890    #[tokio::test]
891    async fn test_schema_change_waits_for_prior_flush() {
892        let mut worker = WorkerTest::new(1, 2).await;
893        let first = worker.submit(1, false).await;
894        let first_release = worker.entered().await;
895        let second = worker.submit(1, true).await;
896        // A second flush permit is free, but the schema barrier must withhold the RPC.
897        assert!(
898            timeout(Duration::from_millis(100), worker.entered.recv())
899                .await
900                .is_err()
901        );
902        first_release.send(()).unwrap();
903        let second_release = worker.entered().await;
904        second_release.send(()).unwrap();
905        timeout(WORKER_TIMEOUT, first)
906            .await
907            .unwrap()
908            .unwrap()
909            .unwrap();
910        timeout(WORKER_TIMEOUT, second)
911            .await
912            .unwrap()
913            .unwrap()
914            .unwrap();
915        worker.shutdown.send(()).unwrap();
916        worker.stopped().await;
917    }
918
919    #[tokio::test]
920    async fn test_table_identity_change_flushes_pending_rows_before_switching() {
921        for recreated in [false, true] {
922            let mut worker = WorkerTest::new(2, 2).await;
923            let first = worker.submit(1, false).await;
924            let mut changed = (*worker.table).clone();
925            if recreated {
926                changed.ident.table_id = 2;
927            } else {
928                changed.ident.version += 1;
929            }
930            let mut second = worker.submit_with_table(2, false, Arc::new(changed)).await;
931            // The first batch is below max_rows but must be flushed on identity change.
932            let first_release = worker.entered().await;
933            assert!(
934                timeout(Duration::from_millis(100), worker.entered.recv())
935                    .await
936                    .is_err()
937            );
938            assert!(
939                timeout(Duration::from_millis(100), &mut second)
940                    .await
941                    .is_err()
942            );
943            first_release.send(()).unwrap();
944            timeout(WORKER_TIMEOUT, first)
945                .await
946                .unwrap()
947                .unwrap()
948                .unwrap();
949            if recreated {
950                // The recreated table has no route in this fixture; its error belongs only to it.
951                assert!(
952                    timeout(WORKER_TIMEOUT, second)
953                        .await
954                        .unwrap()
955                        .unwrap()
956                        .is_err()
957                );
958            } else {
959                let second_release = worker.entered().await;
960                second_release.send(()).unwrap();
961                timeout(WORKER_TIMEOUT, second)
962                    .await
963                    .unwrap()
964                    .unwrap()
965                    .unwrap();
966            }
967            worker.shutdown.send(()).unwrap();
968            worker.stopped().await;
969        }
970    }
971
972    #[tokio::test]
973    async fn test_shutdown_flushes_inline_without_joining_inflight() {
974        let mut worker = WorkerTest::new(2, 1).await;
975        let mut first = worker.submit(2, false).await;
976        let first_release = worker.entered().await;
977        let pending = worker.submit(1, false).await;
978        worker.shutdown.send(()).unwrap();
979        // The normal flush owns the sole permit. Shutdown must still start pending inline.
980        let pending_release = worker.entered().await;
981        pending_release.send(()).unwrap();
982        timeout(WORKER_TIMEOUT, pending)
983            .await
984            .unwrap()
985            .unwrap()
986            .unwrap();
987        worker.stopped().await;
988        assert!(matches!(
989            first.try_recv(),
990            Err(oneshot::error::TryRecvError::Empty)
991        ));
992        first_release.send(()).unwrap();
993        timeout(WORKER_TIMEOUT, first)
994            .await
995            .unwrap()
996            .unwrap()
997            .unwrap();
998    }
999}
servers/batcher/table.rs

servers/batcher/
table.rs
