servers/batcher/table/
flow_notifier.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
15use std::num::NonZeroUsize;
16
17use arrow::record_batch::RecordBatch;
18use common_meta::cache::TableFlownodeSetCacheRef;
19use common_meta::node_manager::NodeManagerRef;
20use common_telemetry::error;
21use lazy_static::lazy_static;
22use operator::req_convert::insert::extract_timestamps;
23use prometheus::{IntCounterVec, register_int_counter_vec};
24use table::metadata::TableInfoRef;
25
26use crate::batcher::flow_notifier::{
27    FlowNotifier as NotificationQueue, start_flow_notification_worker,
28};
29use crate::batcher::flow_sender::FlowNotification;
30
31lazy_static! {
32    static ref FLOW_NOTIFICATION_DROPPED: IntCounterVec = register_int_counter_vec!(
33        "greptime_table_batcher_flow_notification_dropped_total",
34        "Ordinary batch flow notifications dropped before delivery",
35        &["reason"]
36    )
37    .unwrap();
38}
39
40/// Enqueues best-effort flow notifications independently of write completion.
41#[derive(Clone)]
42pub(in crate::batcher::table) struct FlowNotifier {
43    notifier: NotificationQueue,
44}
45
46impl FlowNotifier {
47    /// Starts bounded delivery on the shared runtime, or rejects unsupported capacity.
48    pub fn new(
49        cache: TableFlownodeSetCacheRef,
50        node_manager: NodeManagerRef,
51        capacity: NonZeroUsize,
52    ) -> Option<Self> {
53        let (notifier, receiver) =
54            NotificationQueue::try_new(capacity.get(), FLOW_NOTIFICATION_DROPPED.clone())?;
55        start_flow_notification_worker(receiver, cache, node_manager);
56        Some(Self { notifier })
57    }
58
59    /// Enqueues one notification for batches from the same table.
60    /// Never waits for queue capacity or delivery and never changes the write result.
61    pub fn notify<'a>(
62        &self,
63        table_info: TableInfoRef,
64        batches: impl IntoIterator<Item = &'a RecordBatch>,
65    ) {
66        let Some(timestamp_column) = table_info.meta.schema.timestamp_column() else {
67            return;
68        };
69        let mut timestamps = Vec::new();
70        for batch in batches {
71            match extract_timestamps(batch, &timestamp_column.name) {
72                Ok(batch_timestamps) => timestamps.extend(batch_timestamps),
73                Err(error) => {
74                    error!(error; "Failed to extract flow notification timestamps, table_id: {}", table_info.table_id());
75                }
76            }
77        }
78        if timestamps.is_empty() {
79            return;
80        }
81        self.enqueue(FlowNotification {
82            table_id: table_info.table_id(),
83            timestamps,
84        });
85    }
86
87    fn enqueue(&self, notification: FlowNotification) {
88        self.notifier.try_notify(notification);
89    }
90}
91
92#[cfg(test)]
93mod tests {
94    use std::sync::Arc;
95
96    use api::helper::ColumnDataTypeWrapper;
97    use api::region::RegionResponse;
98    use api::v1::helper::{tag_column_schema, time_index_column_schema};
99    use api::v1::region::{RegionRequest, bulk_insert_request, region_request};
100    use api::v1::value::ValueData;
101    use api::v1::{ColumnDataType, Row, Rows, Value};
102    use arrow::array::{
103        ArrayRef, Int32Array, TimestampMicrosecondArray, TimestampMillisecondArray,
104        TimestampNanosecondArray, TimestampSecondArray,
105    };
106    use arrow::record_batch::RecordBatch;
107    use catalog::memory::MemoryCatalogManager;
108    use common_grpc::flight::FlightDecoder;
109    use common_meta::error::Result as MetaResult;
110    use common_meta::peer::Peer;
111    use common_meta::test_util::{MockDatanodeHandler, MockDatanodeManager};
112    use common_query::request::QueryRequest;
113    use common_recordbatch::SendableRecordBatchStream;
114    use datatypes::data_type::ConcreteDataType;
115    use datatypes::schema::{ColumnSchema, Schema};
116    use operator::insert::Inserter;
117    use operator::req_convert::insert::rows_to_record_batch;
118    use operator::test_util::{
119        create_partition_rule_manager, new_test_table_info, prepare_mocked_backend,
120    };
121    use session::context::QueryContext;
122    use tokio::sync::{Semaphore, oneshot};
123
124    use crate::batcher::flow_notifier::FlowNotifier as NotificationQueue;
125    use crate::batcher::table::batch::{Batch, flush_batch};
126    use crate::batcher::table::flow_notifier::{
127        FLOW_NOTIFICATION_DROPPED, FlowNotification, FlowNotifier,
128    };
129    use crate::batcher::table::metrics;
130    use crate::batcher::table::pending_batch::PendingBatch;
131    use crate::batcher::test_util::mock_table_flownode_cache;
132
133    #[tokio::test]
134    async fn test_flush_coalesces_notifications_for_successful_submissions() {
135        #[derive(Clone)]
136        struct PartialHandler;
137
138        #[async_trait::async_trait]
139        impl MockDatanodeHandler for PartialHandler {
140            async fn handle(&self, _: &Peer, request: RegionRequest) -> MetaResult<RegionResponse> {
141                let Some(region_request::Body::BulkInsert(request)) = request.body else {
142                    panic!("expected bulk insert")
143                };
144                let Some(bulk_insert_request::Body::ArrowIpc(ipc)) = request.body else {
145                    panic!("expected Arrow IPC")
146                };
147                let batch = FlightDecoder::try_from_schema_bytes(&ipc.schema)
148                    .unwrap()
149                    .try_decode_record_batch(&ipc.data_header, &ipc.payload)
150                    .unwrap();
151                let values = batch
152                    .column(0)
153                    .as_any()
154                    .downcast_ref::<Int32Array>()
155                    .unwrap();
156                if values.values().contains(&20) {
157                    return common_meta::error::UnexpectedSnafu {
158                        err_msg: "lost acknowledgement",
159                    }
160                    .fail();
161                }
162                Ok(RegionResponse::new(batch.num_rows()))
163            }
164
165            async fn handle_query(
166                &self,
167                _: &Peer,
168                _: QueryRequest,
169            ) -> MetaResult<SendableRecordBatchStream> {
170                unreachable!()
171            }
172        }
173
174        let partitions = create_partition_rule_manager(prepare_mocked_backend().await).await;
175        let inserter = Arc::new(Inserter::new(
176            MemoryCatalogManager::new(),
177            partitions,
178            Arc::new(MockDatanodeManager::new(PartialHandler)),
179            mock_table_flownode_cache(1, vec![]).await,
180            true,
181        ));
182        let table = Arc::new(new_test_table_info(1, "table_1", [1, 2, 3].into_iter()));
183        let cases = [
184            (
185                vec![vec![1], vec![2]],
186                vec![(1, false), (1, false)],
187                vec![60_000, 120_000],
188                true,
189            ),
190            (
191                vec![vec![1, 20], vec![2], vec![3]],
192                vec![(1, true), (1, false), (1, false)],
193                vec![120_000, 180_000],
194                false,
195            ),
196            (
197                vec![vec![20], vec![20]],
198                vec![(0, true), (0, true)],
199                vec![],
200                false,
201            ),
202            (
203                vec![vec![1]; 100],
204                vec![(1, false); 100],
205                vec![60_000; 100],
206                true,
207            ),
208        ];
209        let ctx = Arc::new(QueryContext::arc().with_confirmed_writes());
210        let ingested = operator::metrics::DIST_INGEST_ROW_COUNT
211            .with_label_values(&[ctx.get_db_string().as_str()]);
212        for (inputs, expected_progress, expected_timestamps, successful_flush) in cases {
213            let permits = Arc::new(Semaphore::new(inputs.len()));
214            let mut submissions = Vec::new();
215            let mut completions = Vec::new();
216            let mut total_rows = 0;
217            for values in inputs {
218                let rows = Rows {
219                    schema: vec![
220                        tag_column_schema("a", ColumnDataType::Int32),
221                        time_index_column_schema("ts", ColumnDataType::TimestampMillisecond),
222                    ],
223                    rows: values
224                        .into_iter()
225                        .map(|value| Row {
226                            values: vec![
227                                Value {
228                                    value_data: Some(ValueData::I32Value(value)),
229                                },
230                                Value {
231                                    value_data: Some(ValueData::TimestampMillisecondValue(
232                                        i64::from(value) * 60_000,
233                                    )),
234                                },
235                            ],
236                        })
237                        .collect(),
238                };
239                total_rows += rows.rows.len();
240                let (response_tx, response_rx) = oneshot::channel();
241                submissions.push(PendingBatch {
242                    _completion: None,
243                    table_info: table.clone(),
244                    batch: rows_to_record_batch(&rows, &table).unwrap(),
245                    ctx: ctx.clone(),
246                    response_tx,
247                    _permit: Arc::new(permits.clone().acquire_owned().await.unwrap()),
248                });
249                completions.push(response_rx);
250            }
251            // Keep a one-slot queue undrained so per-submission notifications lose timestamps.
252            let (notifier, mut notifications) =
253                NotificationQueue::try_new(1, FLOW_NOTIFICATION_DROPPED.clone()).unwrap();
254            let successes = metrics::FLUSH_TOTAL.get();
255            let failures = metrics::FLUSH_FAILURES.get();
256            let samples = metrics::FLUSH_ROWS.get_sample_count();
257            let sum = metrics::FLUSH_ROWS.get_sample_sum();
258            let dropped = metrics::FLUSH_DROPPED_ROWS.get();
259            let written = ingested.get();
260            flush_batch(
261                Batch {
262                    submissions,
263                    total_rows,
264                },
265                inserter.clone(),
266                FlowNotifier { notifier },
267            )
268            .await;
269
270            let expected_written: usize = expected_progress.iter().map(|(rows, _)| rows).sum();
271            assert_eq!(
272                metrics::FLUSH_TOTAL.get(),
273                successes + u64::from(successful_flush)
274            );
275            assert_eq!(
276                metrics::FLUSH_FAILURES.get(),
277                failures + u64::from(!successful_flush)
278            );
279            assert_eq!(
280                metrics::FLUSH_ROWS.get_sample_count(),
281                samples + u64::from(successful_flush)
282            );
283            assert_eq!(
284                metrics::FLUSH_ROWS.get_sample_sum(),
285                sum + if successful_flush {
286                    total_rows as f64
287                } else {
288                    0.0
289                }
290            );
291            assert_eq!(
292                metrics::FLUSH_DROPPED_ROWS.get(),
293                dropped + (total_rows - expected_written) as u64
294            );
295            assert_eq!(ingested.get(), written + expected_written as u64);
296            for (completion, (affected_rows, failed)) in
297                completions.into_iter().zip(expected_progress)
298            {
299                let progress = completion.await.unwrap();
300                assert_eq!(progress.affected_rows, affected_rows);
301                assert_eq!(progress.error.is_some(), failed);
302            }
303            if !expected_timestamps.is_empty() {
304                let notification = notifications.try_recv().unwrap();
305                assert_eq!(notification.table_id, 1);
306                assert_eq!(notification.timestamps, expected_timestamps);
307            }
308            assert!(notifications.try_recv().is_err());
309        }
310    }
311
312    #[test]
313    fn test_full_and_closed_queues_do_not_block() {
314        let (notifier, receiver) =
315            NotificationQueue::try_new(1, FLOW_NOTIFICATION_DROPPED.clone()).unwrap();
316        let notifier = FlowNotifier { notifier };
317        let full = FLOW_NOTIFICATION_DROPPED.with_label_values(&["full"]);
318        let closed = FLOW_NOTIFICATION_DROPPED.with_label_values(&["closed"]);
319        let before_full = full.get();
320        let before_closed = closed.get();
321        for _ in 0..2 {
322            notifier.enqueue(FlowNotification {
323                table_id: 1,
324                timestamps: vec![42],
325            });
326        }
327        assert!(full.get() > before_full);
328        drop(receiver);
329        notifier.enqueue(FlowNotification {
330            table_id: 1,
331            timestamps: vec![42],
332        });
333        assert!(closed.get() > before_closed);
334    }
335    #[test]
336    fn test_notification_preserves_native_timestamp_units() {
337        let cases: Vec<(ColumnDataType, ArrayRef)> = vec![
338            (
339                ColumnDataType::TimestampSecond,
340                Arc::new(TimestampSecondArray::from(vec![Some(42), None])),
341            ),
342            (
343                ColumnDataType::TimestampMillisecond,
344                Arc::new(TimestampMillisecondArray::from(vec![Some(42), None])),
345            ),
346            (
347                ColumnDataType::TimestampMicrosecond,
348                Arc::new(TimestampMicrosecondArray::from(vec![Some(42), None])),
349            ),
350            (
351                ColumnDataType::TimestampNanosecond,
352                Arc::new(TimestampNanosecondArray::from(vec![Some(42), None])),
353            ),
354        ];
355        for (datatype, array) in cases {
356            let mut table = new_test_table_info(1, "test", [0].into_iter());
357            let column = ColumnSchema::new(
358                "ts",
359                ConcreteDataType::from(ColumnDataTypeWrapper::new(datatype, None)),
360                true,
361            )
362            .with_time_index(true);
363            table.meta.schema = Arc::new(Schema::new(vec![column]));
364            let batch = RecordBatch::try_new(table.meta.schema.arrow_schema().clone(), vec![array])
365                .unwrap();
366            let (notifier, mut receiver) =
367                NotificationQueue::try_new(1, FLOW_NOTIFICATION_DROPPED.clone()).unwrap();
368            FlowNotifier { notifier }.notify(Arc::new(table), [&batch, &batch]);
369            let notification = receiver.try_recv().unwrap();
370            assert_eq!(notification.table_id, 1);
371            assert_eq!(notification.timestamps, vec![42, 42]);
372            assert!(receiver.try_recv().is_err());
373        }
374    }
375}
servers/batcher/table/flow_notifier.rs

servers/batcher/table/
flow_notifier.rs
