frontend/
stream_wrapper.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::pin::Pin;
16use std::task::{Context, Poll};
17
18use catalog::process_manager::Ticket;
19use common_recordbatch::adapter::RecordBatchMetrics;
20use common_recordbatch::{OrderOption, RecordBatch, RecordBatchStream, SendableRecordBatchStream};
21use datatypes::schema::SchemaRef;
22use futures::Stream;
23
24pub struct CancellableStreamWrapper {
25    inner: SendableRecordBatchStream,
26    ticket: Ticket,
27    cancel_on_drop: bool,
28    finished: bool,
29}
30
31impl Unpin for CancellableStreamWrapper {}
32
33impl CancellableStreamWrapper {
34    pub fn new(stream: SendableRecordBatchStream, ticket: Ticket) -> Self {
35        Self {
36            inner: stream,
37            ticket,
38            cancel_on_drop: false,
39            finished: false,
40        }
41    }
42
43    pub fn new_cancel_on_drop(stream: SendableRecordBatchStream, ticket: Ticket) -> Self {
44        Self {
45            inner: stream,
46            ticket,
47            cancel_on_drop: true,
48            finished: false,
49        }
50    }
51}
52
53impl Stream for CancellableStreamWrapper {
54    type Item = common_recordbatch::error::Result<RecordBatch>;
55
56    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
57        let this = &mut *self;
58        if this.ticket.cancellation_handle.is_cancelled() {
59            return Poll::Ready(Some(common_recordbatch::error::StreamCancelledSnafu.fail()));
60        }
61
62        if let Poll::Ready(res) = Pin::new(&mut this.inner).poll_next(cx) {
63            if res.is_none() {
64                this.finished = true;
65            }
66            return Poll::Ready(res);
67        }
68
69        // on pending, register cancellation waker.
70        this.ticket.cancellation_handle.waker().register(cx.waker());
71        // check if canceled again.
72        if this.ticket.cancellation_handle.is_cancelled() {
73            return Poll::Ready(Some(common_recordbatch::error::StreamCancelledSnafu.fail()));
74        }
75        Poll::Pending
76    }
77}
78
79impl Drop for CancellableStreamWrapper {
80    fn drop(&mut self) {
81        if self.cancel_on_drop && !self.finished {
82            self.ticket.cancellation_handle.cancel();
83        }
84    }
85}
86
87impl RecordBatchStream for CancellableStreamWrapper {
88    fn schema(&self) -> SchemaRef {
89        self.inner.schema()
90    }
91
92    fn output_ordering(&self) -> Option<&[OrderOption]> {
93        self.inner.output_ordering()
94    }
95
96    fn metrics(&self) -> Option<RecordBatchMetrics> {
97        self.inner.metrics()
98    }
99}
100
101#[cfg(test)]
102mod tests {
103    use std::pin::Pin;
104    use std::sync::Arc;
105    use std::task::{Context, Poll};
106    use std::time::Duration;
107
108    use catalog::process_manager::ProcessManager;
109    use common_recordbatch::adapter::RecordBatchMetrics;
110    use common_recordbatch::{OrderOption, RecordBatch, RecordBatchStream};
111    use datatypes::data_type::ConcreteDataType;
112    use datatypes::prelude::VectorRef;
113    use datatypes::schema::{ColumnSchema, Schema, SchemaRef};
114    use datatypes::vectors::Int32Vector;
115    use futures::{Stream, StreamExt};
116    use tokio::time::{sleep, timeout};
117
118    use super::CancellableStreamWrapper;
119
120    // Mock stream for testing
121    struct MockRecordBatchStream {
122        schema: SchemaRef,
123        batches: Vec<common_recordbatch::error::Result<RecordBatch>>,
124        current: usize,
125        delay: Option<Duration>,
126    }
127
128    impl MockRecordBatchStream {
129        fn new(batches: Vec<common_recordbatch::error::Result<RecordBatch>>) -> Self {
130            let schema = Arc::new(Schema::new(vec![ColumnSchema::new(
131                "test_col",
132                ConcreteDataType::int32_datatype(),
133                false,
134            )]));
135
136            Self {
137                schema,
138                batches,
139                current: 0,
140                delay: None,
141            }
142        }
143
144        fn with_delay(mut self, delay: Duration) -> Self {
145            self.delay = Some(delay);
146            self
147        }
148    }
149
150    impl Stream for MockRecordBatchStream {
151        type Item = common_recordbatch::error::Result<RecordBatch>;
152
153        fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
154            if let Some(delay) = self.delay {
155                // Simulate async delay
156                let waker = cx.waker().clone();
157                let delay_clone = delay;
158                tokio::spawn(async move {
159                    sleep(delay_clone).await;
160                    waker.wake();
161                });
162                self.delay = None; // Only delay once
163                return Poll::Pending;
164            }
165
166            if self.current >= self.batches.len() {
167                return Poll::Ready(None);
168            }
169
170            let batch = self.batches[self.current].as_ref().unwrap().clone();
171            self.current += 1;
172            Poll::Ready(Some(Ok(batch)))
173        }
174    }
175
176    impl RecordBatchStream for MockRecordBatchStream {
177        fn schema(&self) -> SchemaRef {
178            self.schema.clone()
179        }
180
181        fn output_ordering(&self) -> Option<&[OrderOption]> {
182            None
183        }
184
185        fn metrics(&self) -> Option<RecordBatchMetrics> {
186            None
187        }
188    }
189
190    fn create_test_batch() -> RecordBatch {
191        let schema = Arc::new(Schema::new(vec![ColumnSchema::new(
192            "test_col",
193            ConcreteDataType::int32_datatype(),
194            false,
195        )]));
196        RecordBatch::new(
197            schema,
198            vec![Arc::new(Int32Vector::from_values(0..3)) as VectorRef],
199        )
200        .unwrap()
201    }
202
203    #[tokio::test]
204    async fn test_stream_completes_normally() {
205        let batch = create_test_batch();
206        let mock_stream = MockRecordBatchStream::new(vec![Ok(batch.clone())]);
207        let process_manager = Arc::new(ProcessManager::new("".to_string(), None));
208        let ticket = process_manager.register_query(
209            "catalog".to_string(),
210            vec![],
211            "query".to_string(),
212            "client".to_string(),
213            None,
214            None,
215        );
216
217        let mut cancellable_stream = CancellableStreamWrapper::new(Box::pin(mock_stream), ticket);
218
219        let result = cancellable_stream.next().await;
220        assert!(result.is_some());
221        assert!(result.unwrap().is_ok());
222
223        let end_result = cancellable_stream.next().await;
224        assert!(end_result.is_none());
225    }
226
227    #[tokio::test]
228    async fn test_stream_cancelled_before_start() {
229        let batch = create_test_batch();
230        let mock_stream = MockRecordBatchStream::new(vec![Ok(batch)]);
231        let process_manager = Arc::new(ProcessManager::new("".to_string(), None));
232        let ticket = process_manager.register_query(
233            "catalog".to_string(),
234            vec![],
235            "query".to_string(),
236            "client".to_string(),
237            None,
238            None,
239        );
240
241        // Cancel before creating the wrapper
242        ticket.cancellation_handle.cancel();
243
244        let mut cancellable_stream = CancellableStreamWrapper::new(Box::pin(mock_stream), ticket);
245
246        let result = cancellable_stream.next().await;
247        assert!(result.is_some());
248        assert!(result.unwrap().is_err());
249    }
250
251    #[tokio::test]
252    async fn test_stream_cancelled_during_execution() {
253        let batch = create_test_batch();
254        let mock_stream =
255            MockRecordBatchStream::new(vec![Ok(batch)]).with_delay(Duration::from_millis(100));
256        let process_manager = Arc::new(ProcessManager::new("".to_string(), None));
257        let ticket = process_manager.register_query(
258            "catalog".to_string(),
259            vec![],
260            "query".to_string(),
261            "client".to_string(),
262            None,
263            None,
264        );
265        let cancellation_handle = ticket.cancellation_handle.clone();
266
267        let mut cancellable_stream = CancellableStreamWrapper::new(Box::pin(mock_stream), ticket);
268
269        // Cancel after a short delay
270        tokio::spawn(async move {
271            sleep(Duration::from_millis(50)).await;
272            cancellation_handle.cancel();
273        });
274
275        let result = cancellable_stream.next().await;
276        assert!(result.is_some());
277        assert!(result.unwrap().is_err());
278    }
279
280    #[tokio::test]
281    async fn test_stream_completes_before_cancellation() {
282        let batch = create_test_batch();
283        let mock_stream = MockRecordBatchStream::new(vec![Ok(batch.clone())]);
284        let process_manager = Arc::new(ProcessManager::new("".to_string(), None));
285        let ticket = process_manager.register_query(
286            "catalog".to_string(),
287            vec![],
288            "query".to_string(),
289            "client".to_string(),
290            None,
291            None,
292        );
293        let cancellation_handle = ticket.cancellation_handle.clone();
294
295        let mut cancellable_stream = CancellableStreamWrapper::new(Box::pin(mock_stream), ticket);
296
297        // Try to cancel after the stream should have completed
298        tokio::spawn(async move {
299            sleep(Duration::from_millis(100)).await;
300            cancellation_handle.cancel();
301        });
302
303        let result = cancellable_stream.next().await;
304        assert!(result.is_some());
305        assert!(result.unwrap().is_ok());
306    }
307
308    #[tokio::test]
309    async fn test_multiple_batches() {
310        let batch1 = create_test_batch();
311        let batch2 = create_test_batch();
312        let mock_stream = MockRecordBatchStream::new(vec![Ok(batch1), Ok(batch2)]);
313        let process_manager = Arc::new(ProcessManager::new("".to_string(), None));
314        let ticket = process_manager.register_query(
315            "catalog".to_string(),
316            vec![],
317            "query".to_string(),
318            "client".to_string(),
319            None,
320            None,
321        );
322
323        let mut cancellable_stream = CancellableStreamWrapper::new(Box::pin(mock_stream), ticket);
324
325        // First batch
326        let result1 = cancellable_stream.next().await;
327        assert!(result1.is_some());
328        assert!(result1.unwrap().is_ok());
329
330        // Second batch
331        let result2 = cancellable_stream.next().await;
332        assert!(result2.is_some());
333        assert!(result2.unwrap().is_ok());
334
335        // End of stream
336        let end_result = cancellable_stream.next().await;
337        assert!(end_result.is_none());
338    }
339
340    #[tokio::test]
341    async fn test_record_batch_stream_methods() {
342        let batch = create_test_batch();
343        let mock_stream = MockRecordBatchStream::new(vec![Ok(batch)]);
344        let process_manager = Arc::new(ProcessManager::new("".to_string(), None));
345        let ticket = process_manager.register_query(
346            "catalog".to_string(),
347            vec![],
348            "query".to_string(),
349            "client".to_string(),
350            None,
351            None,
352        );
353
354        let cancellable_stream = CancellableStreamWrapper::new(Box::pin(mock_stream), ticket);
355
356        // Test schema method
357        let schema = cancellable_stream.schema();
358        assert_eq!(schema.column_schemas().len(), 1);
359        assert_eq!(schema.column_schemas()[0].name, "test_col");
360
361        // Test output_ordering method
362        assert!(cancellable_stream.output_ordering().is_none());
363
364        // Test metrics method
365        assert!(cancellable_stream.metrics().is_none());
366    }
367
368    #[tokio::test]
369    async fn test_cancellation_during_pending_poll() {
370        let batch = create_test_batch();
371        let mock_stream =
372            MockRecordBatchStream::new(vec![Ok(batch)]).with_delay(Duration::from_millis(200));
373        let process_manager = Arc::new(ProcessManager::new("".to_string(), None));
374        let ticket = process_manager.register_query(
375            "catalog".to_string(),
376            vec![],
377            "query".to_string(),
378            "client".to_string(),
379            None,
380            None,
381        );
382        let cancellation_handle = ticket.cancellation_handle.clone();
383
384        let mut cancellable_stream = CancellableStreamWrapper::new(Box::pin(mock_stream), ticket);
385
386        // Cancel while the stream is pending
387        tokio::spawn(async move {
388            sleep(Duration::from_millis(50)).await;
389            cancellation_handle.cancel();
390        });
391
392        let result = timeout(Duration::from_millis(300), cancellable_stream.next()).await;
393        assert!(result.is_ok());
394        let stream_result = result.unwrap();
395        assert!(stream_result.is_some());
396        assert!(stream_result.unwrap().is_err());
397    }
398
399    #[tokio::test]
400    async fn test_cancel_on_drop_cancels_unfinished_stream() {
401        let batch = create_test_batch();
402        let mock_stream = MockRecordBatchStream::new(vec![Ok(batch)]);
403        let process_manager = Arc::new(ProcessManager::new("".to_string(), None));
404        let ticket = process_manager.register_query(
405            "catalog".to_string(),
406            vec![],
407            "query".to_string(),
408            "client".to_string(),
409            None,
410            None,
411        );
412        let cancellation_handle = ticket.cancellation_handle.clone();
413
414        let cancellable_stream =
415            CancellableStreamWrapper::new_cancel_on_drop(Box::pin(mock_stream), ticket);
416        drop(cancellable_stream);
417
418        assert!(cancellation_handle.is_cancelled());
419    }
420
421    #[tokio::test]
422    async fn test_cancel_on_drop_does_not_cancel_finished_stream() {
423        let batch = create_test_batch();
424        let mock_stream = MockRecordBatchStream::new(vec![Ok(batch)]);
425        let process_manager = Arc::new(ProcessManager::new("".to_string(), None));
426        let ticket = process_manager.register_query(
427            "catalog".to_string(),
428            vec![],
429            "query".to_string(),
430            "client".to_string(),
431            None,
432            None,
433        );
434        let cancellation_handle = ticket.cancellation_handle.clone();
435
436        let mut cancellable_stream =
437            CancellableStreamWrapper::new_cancel_on_drop(Box::pin(mock_stream), ticket);
438        assert!(cancellable_stream.next().await.unwrap().is_ok());
439        assert!(cancellable_stream.next().await.is_none());
440        drop(cancellable_stream);
441
442        assert!(!cancellation_handle.is_cancelled());
443    }
444}
frontend/stream_wrapper.rs

frontend/
stream_wrapper.rs
