mito2/read/
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::sync::Arc;
18use std::task::{Context, Poll};
19use std::time::Instant;
20
21use async_stream::try_stream;
22use common_error::ext::BoxedError;
23use common_recordbatch::error::ExternalSnafu;
24use common_recordbatch::{DfRecordBatch, RecordBatch};
25use futures::stream::BoxStream;
26use futures::{Stream, StreamExt};
27use snafu::ResultExt;
28use tokio::sync::{Semaphore, mpsc};
29use tokio::task::AbortHandle;
30
31use crate::cache::CacheStrategy;
32use crate::error::{JoinSnafu, Result, UnexpectedSnafu};
33use crate::read::BoxedRecordBatchStream;
34use crate::read::flat_projection::FlatProjectionMapper;
35use crate::read::scan_util::PartitionMetrics;
36use crate::read::series_scan::SeriesBatch;
37
38/// All kinds of [`Batch`]es to produce in scanner.
39pub enum ScanBatch {
40    Series(SeriesBatch),
41    RecordBatch(DfRecordBatch),
42}
43
44pub type ScanBatchStream = BoxStream<'static, Result<ScanBatch>>;
45
46/// A stream that takes [`ScanBatch`]es and produces (converts them to) [`RecordBatch`]es.
47pub(crate) struct ConvertBatchStream {
48    inner: ScanBatchStream,
49    projection_mapper: Arc<FlatProjectionMapper>,
50    cache_strategy: CacheStrategy,
51    partition_metrics: PartitionMetrics,
52    pending: VecDeque<RecordBatch>,
53}
54
55impl ConvertBatchStream {
56    pub(crate) fn new(
57        inner: ScanBatchStream,
58        projection_mapper: Arc<FlatProjectionMapper>,
59        cache_strategy: CacheStrategy,
60        partition_metrics: PartitionMetrics,
61    ) -> Self {
62        Self {
63            inner,
64            projection_mapper,
65            cache_strategy,
66            partition_metrics,
67            pending: VecDeque::new(),
68        }
69    }
70
71    fn convert(&mut self, batch: ScanBatch) -> common_recordbatch::error::Result<RecordBatch> {
72        match batch {
73            ScanBatch::Series(series) => {
74                debug_assert!(
75                    self.pending.is_empty(),
76                    "ConvertBatchStream should not convert a new SeriesBatch when pending batches exist"
77                );
78
79                let SeriesBatch::Flat(flat_batch) = series;
80                // Safety: Only flat format returns this batch.
81                for batch in flat_batch.batches {
82                    self.pending.push_back(
83                        self.projection_mapper
84                            .convert(&batch, &self.cache_strategy)?,
85                    );
86                }
87
88                let output_schema = self.projection_mapper.output_schema();
89                Ok(self
90                    .pending
91                    .pop_front()
92                    .unwrap_or_else(|| RecordBatch::new_empty(output_schema)))
93            }
94            ScanBatch::RecordBatch(df_record_batch) => {
95                // Safety: Only flat format returns this batch.
96                self.projection_mapper
97                    .convert(&df_record_batch, &self.cache_strategy)
98            }
99        }
100    }
101}
102
103impl Stream for ConvertBatchStream {
104    type Item = common_recordbatch::error::Result<RecordBatch>;
105
106    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
107        if let Some(batch) = self.pending.pop_front() {
108            return Poll::Ready(Some(Ok(batch)));
109        }
110
111        let batch = futures::ready!(self.inner.poll_next_unpin(cx));
112        let Some(batch) = batch else {
113            return Poll::Ready(None);
114        };
115
116        let record_batch = match batch {
117            Ok(batch) => {
118                let start = Instant::now();
119                let record_batch = self.convert(batch);
120                self.partition_metrics
121                    .inc_convert_batch_cost(start.elapsed());
122                record_batch
123            }
124            Err(e) => Err(BoxedError::new(e)).context(ExternalSnafu),
125        };
126        Poll::Ready(Some(record_batch))
127    }
128}
129
130/// Polls `stream` on its own compaction runtime task, buffering up to
131/// `capacity` batches ahead of the consumer. While polling, the task holds a
132/// permit of `semaphore` if one is given.
133///
134/// The channel also closes if the task panics, so the returned stream checks
135/// how the task ended before reporting the end of the input. Dropping the
136/// returned stream aborts the task.
137pub(crate) fn spawn_compact_stream(
138    mut stream: BoxedRecordBatchStream,
139    capacity: usize,
140    semaphore: Option<Arc<Semaphore>>,
141) -> BoxedRecordBatchStream {
142    let (sender, mut receiver) = mpsc::channel(capacity);
143    let task = common_runtime::spawn_compact(async move {
144        loop {
145            // Releases the permit before sending, so a full channel never holds it.
146            let permit = match &semaphore {
147                Some(semaphore) => match semaphore.acquire().await {
148                    Ok(permit) => Some(permit),
149                    Err(e) => {
150                        let reason = format!("Failed to acquire a stream permit: {e}");
151                        let _ = sender.send(UnexpectedSnafu { reason }.fail()).await;
152                        break;
153                    }
154                },
155                None => None,
156            };
157            let next = stream.next().await;
158            drop(permit);
159            let Some(result) = next else {
160                break;
161            };
162            // Stops after forwarding the first error or once the consumer is gone.
163            let failed = result.is_err();
164            if sender.send(result).await.is_err() || failed {
165                break;
166            }
167        }
168    });
169    // Own the guard before the lazy stream can be dropped without being polled.
170    let abort = AbortOnDrop(task.abort_handle());
171    Box::pin(try_stream! {
172        let _abort = abort;
173        while let Some(batch) = receiver.recv().await {
174            yield batch?;
175        }
176        task.await.context(JoinSnafu)?;
177    })
178}
179
180/// Aborts a task when its consumer is dropped before the end.
181struct AbortOnDrop(AbortHandle);
182
183impl Drop for AbortOnDrop {
184    fn drop(&mut self) {
185        self.0.abort();
186    }
187}
188
189#[cfg(test)]
190mod tests {
191    use std::sync::Arc;
192    use std::task::Poll;
193    use std::time::Duration;
194
195    use datatypes::arrow::array::{ArrayRef, AsArray, Int64Array};
196    use datatypes::arrow::datatypes::{DataType, Field, Int64Type, Schema};
197    use datatypes::arrow::record_batch::RecordBatch;
198    use tokio::sync::oneshot;
199    use tokio::time::timeout;
200
201    use super::*;
202    use crate::error::{Error, InvalidRecordBatchSnafu};
203    use crate::read::flat_merge::FlatMergeReader;
204
205    /// Keeps the producer pending and exposes when its source is polled and dropped.
206    struct PendingStream {
207        polled: Option<oneshot::Sender<()>>,
208        dropped: Option<oneshot::Sender<()>>,
209    }
210
211    impl PendingStream {
212        fn new() -> (Self, oneshot::Receiver<()>, oneshot::Receiver<()>) {
213            let (polled_tx, polled_rx) = oneshot::channel();
214            let (dropped_tx, dropped_rx) = oneshot::channel();
215            (
216                Self {
217                    polled: Some(polled_tx),
218                    dropped: Some(dropped_tx),
219                },
220                polled_rx,
221                dropped_rx,
222            )
223        }
224    }
225
226    impl Stream for PendingStream {
227        type Item = Result<RecordBatch>;
228
229        fn poll_next(mut self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<Option<Self::Item>> {
230            if let Some(polled) = self.polled.take() {
231                let _ = polled.send(());
232            }
233            Poll::Pending
234        }
235    }
236
237    impl Drop for PendingStream {
238        fn drop(&mut self) {
239            if let Some(dropped) = self.dropped.take() {
240                let _ = dropped.send(());
241            }
242        }
243    }
244
245    async fn wait_for_signal(signal: oneshot::Receiver<()>) {
246        // Notifications establish ordering; the timeout only bounds a broken test.
247        timeout(Duration::from_secs(10), signal)
248            .await
249            .expect("timed out waiting for the source")
250            .expect("source dropped without sending the signal");
251    }
252
253    fn batch(value: i64) -> RecordBatch {
254        let schema = Arc::new(Schema::new(vec![Field::new("v", DataType::Int64, false)]));
255        let column = Arc::new(Int64Array::from(vec![value])) as ArrayRef;
256        RecordBatch::try_new(schema, vec![column]).unwrap()
257    }
258
259    fn values_stream(values: std::ops::Range<i64>) -> BoxedRecordBatchStream {
260        Box::pin(futures::stream::iter(values.map(|value| Ok(batch(value)))))
261    }
262
263    /// Drains `stream`, returning the values read before it ended or failed.
264    async fn drain(mut stream: BoxedRecordBatchStream) -> (Vec<i64>, Option<Error>) {
265        let mut values = Vec::new();
266        while let Some(result) = stream.next().await {
267            match result {
268                Ok(batch) => values.push(batch.column(0).as_primitive::<Int64Type>().value(0)),
269                Err(e) => return (values, Some(e)),
270            }
271        }
272        (values, None)
273    }
274
275    #[tokio::test]
276    async fn test_spawn_compact_stream_forwards_batches_in_order() {
277        let (values, error) = drain(spawn_compact_stream(values_stream(0..5), 2, None)).await;
278
279        assert_eq!(vec![0, 1, 2, 3, 4], values);
280        assert!(error.is_none());
281    }
282
283    #[tokio::test]
284    async fn test_spawn_compact_stream_stops_at_first_error() {
285        let stream = futures::stream::iter(vec![
286            Ok(batch(0)),
287            InvalidRecordBatchSnafu { reason: "boom" }.fail(),
288            Ok(batch(2)),
289        ]);
290
291        let (values, error) = drain(spawn_compact_stream(Box::pin(stream), 2, None)).await;
292
293        assert_eq!(vec![0], values);
294        assert!(
295            matches!(error, Some(Error::InvalidRecordBatch { .. })),
296            "{error:?}"
297        );
298    }
299
300    #[tokio::test]
301    async fn test_spawn_compact_stream_reports_panic_instead_of_end() {
302        let mut polled = false;
303        let stream = futures::stream::poll_fn(move |_| {
304            if std::mem::replace(&mut polled, true) {
305                panic!("source panicked");
306            }
307            Poll::Ready(Some(Ok(batch(0))))
308        });
309
310        let (values, error) = drain(spawn_compact_stream(Box::pin(stream), 2, None)).await;
311
312        // A truncated input must never look like a complete one.
313        assert_eq!(vec![0], values);
314        assert!(matches!(error, Some(Error::Join { .. })), "{error:?}");
315    }
316
317    #[rstest::rstest]
318    #[case::before_first_poll(false)]
319    #[case::after_first_poll(true)]
320    #[tokio::test]
321    async fn test_dropped_consumer_cancels_pending_source(#[case] poll_consumer: bool) {
322        let (source, polled, dropped) = PendingStream::new();
323        let semaphore = Arc::new(Semaphore::new(1));
324        let mut consumer = spawn_compact_stream(Box::pin(source), 1, Some(semaphore.clone()));
325        wait_for_signal(polled).await;
326        assert_eq!(0, semaphore.available_permits());
327        if poll_consumer {
328            assert!(futures::poll!(consumer.next()).is_pending());
329        }
330
331        // A guard created inside the lazy stream misses the before-first-poll case.
332        drop(consumer);
333        wait_for_signal(dropped).await;
334        let _permit = timeout(Duration::from_secs(10), semaphore.acquire())
335            .await
336            .expect("producer retained its permit")
337            .unwrap();
338    }
339
340    #[tokio::test]
341    async fn test_merge_initialization_error_cancels_unpolled_source() {
342        let (source, polled, dropped) = PendingStream::new();
343        let consumer = spawn_compact_stream(Box::pin(source), 1, None);
344        wait_for_signal(polled).await;
345        let error_stream = Box::pin(futures::stream::iter(vec![
346            InvalidRecordBatchSnafu { reason: "boom" }.fail(),
347        ]));
348        let sources = vec![spawn_compact_stream(error_stream, 1, None), consumer];
349
350        let result = FlatMergeReader::new(batch(0).schema(), sources, 1024, None).await;
351
352        assert!(matches!(result, Err(Error::InvalidRecordBatch { .. })));
353        wait_for_signal(dropped).await;
354    }
355
356    #[tokio::test]
357    async fn test_cancelled_merge_initialization_drops_all_sources() {
358        let (first, first_polled, first_dropped) = PendingStream::new();
359        let (second, second_polled, second_dropped) = PendingStream::new();
360        let sources = vec![
361            spawn_compact_stream(Box::pin(first), 1, None),
362            spawn_compact_stream(Box::pin(second), 1, None),
363        ];
364        wait_for_signal(first_polled).await;
365        wait_for_signal(second_polled).await;
366        let mut initialization =
367            Box::pin(FlatMergeReader::new(batch(0).schema(), sources, 1024, None));
368        assert!(futures::poll!(&mut initialization).is_pending());
369
370        // Only the first consumer has been polled by the sequential initialization.
371        drop(initialization);
372
373        wait_for_signal(first_dropped).await;
374        wait_for_signal(second_dropped).await;
375    }
376
377    #[tokio::test]
378    async fn test_spawn_compact_streams_share_semaphore() {
379        let semaphore = Arc::new(Semaphore::new(1));
380        let streams: Vec<_> = (0..3)
381            .map(|i| {
382                spawn_compact_stream(
383                    values_stream(i * 10..i * 10 + 3),
384                    1,
385                    Some(semaphore.clone()),
386                )
387            })
388            .collect();
389
390        for (i, stream) in streams.into_iter().enumerate() {
391            let start = i as i64 * 10;
392            let (values, error) = drain(stream).await;
393            assert_eq!((start..start + 3).collect::<Vec<_>>(), values);
394            assert!(error.is_none());
395        }
396    }
397}
mito2/read/stream.rs

mito2/read/
stream.rs
