servers/grpc/flight/
stream.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::collections::VecDeque;
16use std::pin::Pin;
17use std::task::{Context, Poll};
18use std::time::{Duration, Instant};
19
20use arrow_flight::FlightData;
21use common_error::ext::ErrorExt;
22use common_grpc::flight::{FlightEncoder, FlightMessage};
23use common_recordbatch::SendableRecordBatchStream;
24use common_telemetry::tracing::{Instrument, info_span};
25use common_telemetry::tracing_context::{FutureExt, TracingContext};
26use common_telemetry::{error, info, warn};
27use futures::channel::mpsc;
28use futures::channel::mpsc::Sender;
29use futures::{SinkExt, Stream, StreamExt};
30use pin_project::{pin_project, pinned_drop};
31use session::context::{
32    FLIGHT_METRICS_HEARTBEAT_INTERVAL, QueryContextRef,
33    SUPPORT_FLIGHT_METRICS_BEFORE_BATCH_EXTENSION_KEY,
34};
35use snafu::ResultExt;
36use tokio::task::JoinHandle;
37use tokio::time;
38
39use crate::error;
40use crate::grpc::FlightCompression;
41use crate::grpc::flight::TonicResult;
42
43/// Metrics collector for Flight stream with RAII logging pattern
44struct StreamMetrics {
45    send_schema_duration: Duration,
46    send_record_batch_duration: Duration,
47    send_metrics_duration: Duration,
48    fetch_content_duration: Duration,
49    record_batch_count: usize,
50    metrics_count: usize,
51    total_rows: usize,
52    total_bytes: usize,
53    should_log: bool,
54}
55
56impl StreamMetrics {
57    fn new(should_log: bool) -> Self {
58        Self {
59            send_schema_duration: Duration::ZERO,
60            send_record_batch_duration: Duration::ZERO,
61            send_metrics_duration: Duration::ZERO,
62            fetch_content_duration: Duration::ZERO,
63            record_batch_count: 0,
64            metrics_count: 0,
65            total_rows: 0,
66            total_bytes: 0,
67            should_log,
68        }
69    }
70}
71
72impl Drop for StreamMetrics {
73    fn drop(&mut self) {
74        if self.should_log {
75            info!(
76                "flight_data_stream finished: \
77                send_schema_duration={:?}, \
78                send_record_batch_duration={:?}, \
79                send_metrics_duration={:?}, \
80                fetch_content_duration={:?}, \
81                record_batch_count={}, \
82                metrics_count={}, \
83                total_rows={}, \
84                total_bytes={}",
85                self.send_schema_duration,
86                self.send_record_batch_duration,
87                self.send_metrics_duration,
88                self.fetch_content_duration,
89                self.record_batch_count,
90                self.metrics_count,
91                self.total_rows,
92                self.total_bytes
93            );
94        }
95    }
96}
97
98#[pin_project(PinnedDrop)]
99pub struct FlightRecordBatchStream {
100    #[pin]
101    rx: mpsc::Receiver<Result<FlightMessage, tonic::Status>>,
102    join_handle: JoinHandle<()>,
103    done: bool,
104    encoder: FlightEncoder,
105    buffer: VecDeque<FlightData>,
106}
107
108impl FlightRecordBatchStream {
109    async fn send_metrics(
110        tx: &mut Sender<TonicResult<FlightMessage>>,
111        metrics: &mut StreamMetrics,
112        metrics_str: String,
113    ) -> bool {
114        metrics.metrics_count += 1;
115        let start = Instant::now();
116        if let Err(e) = tx.send(Ok(FlightMessage::Metrics(metrics_str))).await {
117            warn!(e; "stop sending Flight data");
118            return false;
119        }
120        metrics.send_metrics_duration += start.elapsed();
121        true
122    }
123
124    async fn send_metrics_if_changed(
125        tx: &mut Sender<TonicResult<FlightMessage>>,
126        metrics: &mut StreamMetrics,
127        last_metrics_str: &mut Option<String>,
128        metrics_str: String,
129    ) -> bool {
130        if last_metrics_str.as_deref() == Some(metrics_str.as_str()) {
131            return true;
132        }
133
134        *last_metrics_str = Some(metrics_str.clone());
135        Self::send_metrics(tx, metrics, metrics_str).await
136    }
137
138    pub fn new(
139        recordbatches: SendableRecordBatchStream,
140        tracing_context: TracingContext,
141        compression: FlightCompression,
142        query_ctx: QueryContextRef,
143    ) -> Self {
144        let should_send_partial_metrics = query_ctx.explain_verbose();
145        let can_send_metrics_before_batch = query_ctx
146            .remote_query_id()
147            .zip(query_ctx.extension(SUPPORT_FLIGHT_METRICS_BEFORE_BATCH_EXTENSION_KEY))
148            .is_some_and(|(remote_query_id, capability)| capability == remote_query_id);
149        let (tx, rx) = mpsc::channel::<TonicResult<FlightMessage>>(1);
150        let join_handle = common_runtime::spawn_global(async move {
151            Self::flight_data_stream(
152                recordbatches,
153                tx,
154                should_send_partial_metrics,
155                can_send_metrics_before_batch,
156            )
157            .trace(tracing_context.attach(info_span!("flight_data_stream")))
158            .await
159        });
160        let encoder = if compression.arrow_compression() {
161            FlightEncoder::default()
162        } else {
163            FlightEncoder::with_compression_disabled()
164        };
165        Self {
166            rx,
167            join_handle,
168            done: false,
169            encoder,
170            buffer: VecDeque::new(),
171        }
172    }
173
174    async fn flight_data_stream(
175        mut recordbatches: SendableRecordBatchStream,
176        mut tx: Sender<TonicResult<FlightMessage>>,
177        should_send_partial_metrics: bool,
178        can_send_metrics_before_batch: bool,
179    ) {
180        let mut metrics = StreamMetrics::new(should_send_partial_metrics);
181        let mut last_metrics_str = None;
182
183        let schema = recordbatches.schema().arrow_schema().clone();
184        let start = Instant::now();
185        if let Err(e) = tx.send(Ok(FlightMessage::Schema(schema))).await {
186            warn!(e; "stop sending Flight data");
187            return;
188        }
189        metrics.send_schema_duration += start.elapsed();
190
191        loop {
192            let start = Instant::now();
193            let batch_or_err = if should_send_partial_metrics && can_send_metrics_before_batch {
194                match time::timeout(
195                    FLIGHT_METRICS_HEARTBEAT_INTERVAL,
196                    recordbatches.next().in_current_span(),
197                )
198                .await
199                {
200                    Ok(result) => result,
201                    Err(_) => {
202                        if let Some(metrics_str) = recordbatches
203                            .metrics()
204                            .and_then(|m| serde_json::to_string(&m).ok())
205                            && !Self::send_metrics_if_changed(
206                                &mut tx,
207                                &mut metrics,
208                                &mut last_metrics_str,
209                                metrics_str,
210                            )
211                            .await
212                        {
213                            return;
214                        }
215                        metrics.fetch_content_duration += start.elapsed();
216                        continue;
217                    }
218                }
219            } else {
220                recordbatches.next().in_current_span().await
221            };
222            metrics.fetch_content_duration += start.elapsed();
223            let Some(batch_or_err) = batch_or_err else {
224                break;
225            };
226            match batch_or_err {
227                Ok(recordbatch) => {
228                    metrics.total_rows += recordbatch.num_rows();
229                    metrics.record_batch_count += 1;
230                    metrics.total_bytes += recordbatch.df_record_batch().get_array_memory_size();
231
232                    let start = Instant::now();
233                    if let Err(e) = tx
234                        .send(Ok(FlightMessage::RecordBatch(
235                            recordbatch.into_df_record_batch(),
236                        )))
237                        .await
238                    {
239                        warn!(e; "stop sending Flight data");
240                        return;
241                    }
242                    metrics.send_record_batch_duration += start.elapsed();
243
244                    if should_send_partial_metrics
245                        && let Some(metrics_str) = recordbatches
246                            .metrics()
247                            .and_then(|m| serde_json::to_string(&m).ok())
248                        && {
249                            last_metrics_str = Some(metrics_str.clone());
250                            !Self::send_metrics(&mut tx, &mut metrics, metrics_str).await
251                        }
252                    {
253                        return;
254                    }
255                }
256                Err(e) => {
257                    if e.status_code().should_log_error() {
258                        error!("{e:?}");
259                    }
260
261                    let e = Err(e).context(error::CollectRecordbatchSnafu);
262                    if let Err(e) = tx.send(e.map_err(|x| x.into())).await {
263                        warn!(e; "stop sending Flight data");
264                    }
265                    return;
266                }
267            }
268        }
269        // make last package to pass metrics
270        if let Some(metrics_str) = recordbatches
271            .metrics()
272            .and_then(|m| serde_json::to_string(&m).ok())
273        {
274            let _ = Self::send_metrics(&mut tx, &mut metrics, metrics_str).await;
275        }
276    }
277}
278
279#[pinned_drop]
280impl PinnedDrop for FlightRecordBatchStream {
281    fn drop(self: Pin<&mut Self>) {
282        self.join_handle.abort();
283    }
284}
285
286impl Stream for FlightRecordBatchStream {
287    type Item = TonicResult<FlightData>;
288
289    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
290        let this = self.project();
291        if *this.done {
292            Poll::Ready(None)
293        } else {
294            if let Some(x) = this.buffer.pop_front() {
295                return Poll::Ready(Some(Ok(x)));
296            }
297            match this.rx.poll_next(cx) {
298                Poll::Ready(None) => {
299                    *this.done = true;
300                    Poll::Ready(None)
301                }
302                Poll::Ready(Some(result)) => match result {
303                    Ok(flight_message) => {
304                        let mut iter = this.encoder.encode(flight_message).into_iter();
305                        let Some(first) = iter.next() else {
306                            // Safety: `iter` on a type of `Vec1`, which is guaranteed to have
307                            // at least one element.
308                            unreachable!()
309                        };
310                        this.buffer.extend(iter);
311                        Poll::Ready(Some(Ok(first)))
312                    }
313                    Err(e) => {
314                        *this.done = true;
315                        Poll::Ready(Some(Err(e)))
316                    }
317                },
318                Poll::Pending => Poll::Pending,
319            }
320        }
321    }
322}
323
324#[cfg(test)]
325mod test {
326    use std::pin::Pin;
327    use std::sync::Arc;
328    use std::task::{Context, Poll};
329    use std::time::Duration;
330
331    use common_grpc::flight::{FlightDecoder, FlightMessage};
332    use common_recordbatch::adapter::RecordBatchMetrics;
333    use common_recordbatch::{OrderOption, RecordBatch, RecordBatchStream, RecordBatches};
334    use datatypes::prelude::*;
335    use datatypes::schema::{ColumnSchema, Schema, SchemaRef};
336    use datatypes::vectors::Int32Vector;
337    use futures::StreamExt;
338    use session::context::{QueryContext, SUPPORT_FLIGHT_METRICS_BEFORE_BATCH_EXTENSION_KEY};
339
340    use super::*;
341
342    struct PendingMetricsStream {
343        schema: SchemaRef,
344        metrics: RecordBatchMetrics,
345    }
346
347    struct MetricsThenBatchStream {
348        schema: SchemaRef,
349        metrics: RecordBatchMetrics,
350        rx: tokio::sync::mpsc::UnboundedReceiver<common_recordbatch::error::Result<RecordBatch>>,
351    }
352
353    fn query_context_with_matching_capability() -> Arc<QueryContext> {
354        let query_ctx = QueryContext::arc();
355        let remote_query_id = query_ctx
356            .remote_query_id()
357            .expect("query context must have remote query id")
358            .to_string();
359        let mut query_ctx = (*query_ctx).clone();
360        query_ctx.set_extension(
361            SUPPORT_FLIGHT_METRICS_BEFORE_BATCH_EXTENSION_KEY,
362            remote_query_id,
363        );
364        Arc::new(query_ctx)
365    }
366
367    fn query_context_with_capability(capability: &str) -> Arc<QueryContext> {
368        let mut query_ctx = (*QueryContext::arc()).clone();
369        query_ctx.set_extension(
370            SUPPORT_FLIGHT_METRICS_BEFORE_BATCH_EXTENSION_KEY,
371            capability,
372        );
373        Arc::new(query_ctx)
374    }
375
376    impl RecordBatchStream for PendingMetricsStream {
377        fn schema(&self) -> SchemaRef {
378            self.schema.clone()
379        }
380
381        fn output_ordering(&self) -> Option<&[OrderOption]> {
382            None
383        }
384
385        fn metrics(&self) -> Option<RecordBatchMetrics> {
386            Some(self.metrics.clone())
387        }
388    }
389
390    impl Stream for PendingMetricsStream {
391        type Item = common_recordbatch::error::Result<RecordBatch>;
392
393        fn poll_next(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
394            Poll::Pending
395        }
396    }
397
398    impl RecordBatchStream for MetricsThenBatchStream {
399        fn schema(&self) -> SchemaRef {
400            self.schema.clone()
401        }
402
403        fn output_ordering(&self) -> Option<&[OrderOption]> {
404            None
405        }
406
407        fn metrics(&self) -> Option<RecordBatchMetrics> {
408            Some(self.metrics.clone())
409        }
410    }
411
412    impl Stream for MetricsThenBatchStream {
413        type Item = common_recordbatch::error::Result<RecordBatch>;
414
415        fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
416            self.rx.poll_recv(cx)
417        }
418    }
419
420    #[tokio::test]
421    async fn test_flight_record_batch_stream() {
422        let schema = Arc::new(Schema::new(vec![ColumnSchema::new(
423            "a",
424            ConcreteDataType::int32_datatype(),
425            false,
426        )]));
427
428        let v: VectorRef = Arc::new(Int32Vector::from_slice([1, 2]));
429        let recordbatch = RecordBatch::new(schema.clone(), vec![v]).unwrap();
430
431        let recordbatches = RecordBatches::try_new(schema.clone(), vec![recordbatch.clone()])
432            .unwrap()
433            .as_stream();
434        let mut stream = FlightRecordBatchStream::new(
435            recordbatches,
436            TracingContext::default(),
437            FlightCompression::default(),
438            QueryContext::arc(),
439        );
440
441        let mut raw_data = Vec::with_capacity(2);
442        raw_data.push(stream.next().await.unwrap().unwrap());
443        raw_data.push(stream.next().await.unwrap().unwrap());
444        assert!(stream.next().await.is_none());
445        assert!(stream.done);
446
447        let decoder = &mut FlightDecoder::default();
448        let mut flight_messages = raw_data
449            .into_iter()
450            .map(|x| decoder.try_decode(&x).unwrap().unwrap())
451            .collect::<Vec<FlightMessage>>();
452        assert_eq!(flight_messages.len(), 2);
453
454        match flight_messages.remove(0) {
455            FlightMessage::Schema(actual_schema) => {
456                assert_eq!(&actual_schema, schema.arrow_schema());
457            }
458            _ => unreachable!(),
459        }
460
461        match flight_messages.remove(0) {
462            FlightMessage::RecordBatch(actual_recordbatch) => {
463                assert_eq!(&actual_recordbatch, recordbatch.df_record_batch());
464            }
465            _ => unreachable!(),
466        }
467    }
468
469    #[tokio::test]
470    async fn test_flight_record_batch_stream_emits_metrics_while_pending() {
471        let schema = Arc::new(Schema::new(vec![ColumnSchema::new(
472            "a",
473            ConcreteDataType::int32_datatype(),
474            false,
475        )]));
476        let metrics = RecordBatchMetrics {
477            elapsed_compute: 42,
478            ..Default::default()
479        };
480        let recordbatches = Box::pin(PendingMetricsStream {
481            schema: schema.clone(),
482            metrics,
483        });
484        let query_ctx = query_context_with_matching_capability();
485        query_ctx.set_explain_verbose(true);
486        let mut stream = FlightRecordBatchStream::new(
487            recordbatches,
488            TracingContext::default(),
489            FlightCompression::default(),
490            query_ctx,
491        );
492
493        let decoder = &mut FlightDecoder::default();
494        let schema_data = stream.next().await.unwrap().unwrap();
495        match decoder.try_decode(&schema_data).unwrap().unwrap() {
496            FlightMessage::Schema(actual_schema) => {
497                assert_eq!(&actual_schema, schema.arrow_schema());
498            }
499            _ => unreachable!(),
500        }
501
502        let metrics_data = tokio::time::timeout(Duration::from_secs(2), stream.next())
503            .await
504            .unwrap()
505            .unwrap()
506            .unwrap();
507        match decoder.try_decode(&metrics_data).unwrap().unwrap() {
508            FlightMessage::Metrics(metrics) => {
509                let metrics: RecordBatchMetrics = serde_json::from_str(&metrics).unwrap();
510                assert_eq!(metrics.elapsed_compute, 42);
511            }
512            other => panic!("expected metrics message, got {other:?}"),
513        }
514    }
515
516    #[tokio::test]
517    async fn test_flight_record_batch_stream_continues_after_pending_metrics() {
518        let schema = Arc::new(Schema::new(vec![ColumnSchema::new(
519            "a",
520            ConcreteDataType::int32_datatype(),
521            false,
522        )]));
523        let metrics = RecordBatchMetrics {
524            elapsed_compute: 42,
525            ..Default::default()
526        };
527        let recordbatch = RecordBatch::new(
528            schema.clone(),
529            vec![Arc::new(Int32Vector::from_slice([1])) as VectorRef],
530        )
531        .unwrap();
532        let expected_recordbatch = recordbatch.df_record_batch().clone();
533        let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
534        let recordbatches = Box::pin(MetricsThenBatchStream {
535            schema: schema.clone(),
536            metrics,
537            rx,
538        });
539        let query_ctx = query_context_with_matching_capability();
540        query_ctx.set_explain_verbose(true);
541        let mut stream = FlightRecordBatchStream::new(
542            recordbatches,
543            TracingContext::default(),
544            FlightCompression::default(),
545            query_ctx,
546        );
547
548        let decoder = &mut FlightDecoder::default();
549        let schema_data = stream.next().await.unwrap().unwrap();
550        assert!(matches!(
551            decoder.try_decode(&schema_data).unwrap().unwrap(),
552            FlightMessage::Schema(_)
553        ));
554
555        let metrics_data = tokio::time::timeout(Duration::from_secs(2), stream.next())
556            .await
557            .unwrap()
558            .unwrap()
559            .unwrap();
560        assert!(matches!(
561            decoder.try_decode(&metrics_data).unwrap().unwrap(),
562            FlightMessage::Metrics(_)
563        ));
564
565        tx.send(Ok(recordbatch)).unwrap();
566        let batch_data = tokio::time::timeout(Duration::from_secs(2), stream.next())
567            .await
568            .unwrap()
569            .unwrap()
570            .unwrap();
571        match decoder.try_decode(&batch_data).unwrap().unwrap() {
572            FlightMessage::RecordBatch(actual_recordbatch) => {
573                assert_eq!(&actual_recordbatch, &expected_recordbatch);
574            }
575            other => panic!("expected record batch after pending metrics, got {other:?}"),
576        }
577    }
578
579    #[tokio::test]
580    async fn test_flight_record_batch_stream_requires_capability_for_pre_batch_metrics() {
581        let schema = Arc::new(Schema::new(vec![ColumnSchema::new(
582            "a",
583            ConcreteDataType::int32_datatype(),
584            false,
585        )]));
586        let recordbatches = Box::pin(PendingMetricsStream {
587            schema: schema.clone(),
588            metrics: RecordBatchMetrics {
589                elapsed_compute: 42,
590                ..Default::default()
591            },
592        });
593        let query_ctx = QueryContext::arc();
594        query_ctx.set_explain_verbose(true);
595        let mut stream = FlightRecordBatchStream::new(
596            recordbatches,
597            TracingContext::default(),
598            FlightCompression::default(),
599            query_ctx,
600        );
601
602        let decoder = &mut FlightDecoder::default();
603        let schema_data = stream.next().await.unwrap().unwrap();
604        assert!(matches!(
605            decoder.try_decode(&schema_data).unwrap().unwrap(),
606            FlightMessage::Schema(_)
607        ));
608        assert!(
609            tokio::time::timeout(
610                FLIGHT_METRICS_HEARTBEAT_INTERVAL + Duration::from_millis(200),
611                stream.next()
612            )
613            .await
614            .is_err(),
615            "pre-batch Metrics must be gated by client capability"
616        );
617    }
618
619    #[tokio::test]
620    async fn test_flight_record_batch_stream_rejects_spoofed_capability_for_pre_batch_metrics() {
621        let schema = Arc::new(Schema::new(vec![ColumnSchema::new(
622            "a",
623            ConcreteDataType::int32_datatype(),
624            false,
625        )]));
626        let recordbatches = Box::pin(PendingMetricsStream {
627            schema: schema.clone(),
628            metrics: RecordBatchMetrics {
629                elapsed_compute: 42,
630                ..Default::default()
631            },
632        });
633        let query_ctx = query_context_with_capability("true");
634        query_ctx.set_explain_verbose(true);
635        let mut stream = FlightRecordBatchStream::new(
636            recordbatches,
637            TracingContext::default(),
638            FlightCompression::default(),
639            query_ctx,
640        );
641
642        let decoder = &mut FlightDecoder::default();
643        let schema_data = stream.next().await.unwrap().unwrap();
644        assert!(matches!(
645            decoder.try_decode(&schema_data).unwrap().unwrap(),
646            FlightMessage::Schema(_)
647        ));
648        assert!(
649            tokio::time::timeout(
650                FLIGHT_METRICS_HEARTBEAT_INTERVAL + Duration::from_millis(200),
651                stream.next()
652            )
653            .await
654            .is_err(),
655            "pre-batch Metrics must reject a spoofed capability"
656        );
657    }
658
659    #[tokio::test]
660    async fn test_flight_record_batch_stream_requires_explain_verbose_for_pre_batch_metrics() {
661        let schema = Arc::new(Schema::new(vec![ColumnSchema::new(
662            "a",
663            ConcreteDataType::int32_datatype(),
664            false,
665        )]));
666        let recordbatches = Box::pin(PendingMetricsStream {
667            schema: schema.clone(),
668            metrics: RecordBatchMetrics {
669                elapsed_compute: 42,
670                ..Default::default()
671            },
672        });
673        let query_ctx = query_context_with_matching_capability();
674        let mut stream = FlightRecordBatchStream::new(
675            recordbatches,
676            TracingContext::default(),
677            FlightCompression::default(),
678            query_ctx,
679        );
680
681        let decoder = &mut FlightDecoder::default();
682        let schema_data = stream.next().await.unwrap().unwrap();
683        assert!(matches!(
684            decoder.try_decode(&schema_data).unwrap().unwrap(),
685            FlightMessage::Schema(_)
686        ));
687        assert!(
688            tokio::time::timeout(
689                FLIGHT_METRICS_HEARTBEAT_INTERVAL + Duration::from_millis(200),
690                stream.next()
691            )
692            .await
693            .is_err(),
694            "pre-batch Metrics must be gated by explain verbose even when capability is set"
695        );
696    }
697}
servers/grpc/flight/stream.rs

servers/grpc/flight/
stream.rs
