servers/batcher/table/
pending_batch.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::sync::Arc;
16
17use arrow::record_batch::RecordBatch;
18use common_error::ext::BoxedError;
19use common_query::WriteProgress;
20use operator::error::Error;
21use session::context::QueryContextRef;
22use snafu::IntoError;
23use table::metadata::TableInfoRef;
24use tokio::sync::{OwnedSemaphorePermit, oneshot};
25
26/// One complete input and its completion/admission ownership.
27pub(in crate::batcher::table) struct PendingBatch {
28    pub table_info: TableInfoRef,
29    pub batch: RecordBatch,
30    pub ctx: QueryContextRef,
31    pub response_tx: oneshot::Sender<WriteProgress>,
32    pub _permit: Arc<OwnedSemaphorePermit>,
33    pub _completion: Option<session::write_progress::PendingWrite>,
34}
35
36pub(in crate::batcher::table) fn notify_batches(
37    batches: Vec<PendingBatch>,
38    result: Result<(), Arc<Error>>,
39) {
40    for batch in batches {
41        let progress = match &result {
42            Ok(()) => WriteProgress {
43                affected_rows: batch.batch.num_rows(),
44                error: None,
45            },
46            Err(error) => WriteProgress::failed(BoxedError::new(
47                operator::error::BatchFlushSnafu.into_error(error.clone()),
48            )),
49        };
50        let _ = batch.response_tx.send(progress);
51    }
52}
53
54#[cfg(test)]
55mod tests {
56    use arrow::array::Int32Array;
57    use arrow::datatypes::{DataType, Field, Schema};
58    use operator::error::UnexpectedSnafu;
59    use operator::test_util::new_test_table_info;
60    use session::context::QueryContext;
61    use tokio::sync::Semaphore;
62
63    use crate::batcher::table::pending_batch::*;
64
65    #[tokio::test]
66    async fn test_completion_fans_out_and_releases_admission() {
67        for fail in [false, true] {
68            let semaphore = Arc::new(Semaphore::new(3));
69            let mut batches = Vec::new();
70            let mut receivers = Vec::new();
71            for value in 0..3 {
72                let (response_tx, response_rx) = oneshot::channel();
73                let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)]));
74                batches.push(PendingBatch {
75                    _completion: None,
76                    table_info: Arc::new(new_test_table_info(1, "t", [0].into_iter())),
77                    batch: RecordBatch::try_new(
78                        schema,
79                        vec![Arc::new(Int32Array::from(vec![value]))],
80                    )
81                    .unwrap(),
82                    ctx: QueryContext::arc(),
83                    response_tx,
84                    _permit: Arc::new(semaphore.clone().acquire_owned().await.unwrap()),
85                });
86                receivers.push(response_rx);
87            }
88            assert_eq!(0, semaphore.available_permits());
89            // A cancelled waiter must not prevent delivery to the others.
90            drop(receivers.pop());
91            let error = Arc::new(
92                UnexpectedSnafu {
93                    violated: "test flush failure".to_string(),
94                }
95                .build(),
96            );
97            notify_batches(batches, if fail { Err(error.clone()) } else { Ok(()) });
98            assert_eq!(3, semaphore.available_permits());
99            for receiver in receivers {
100                let progress = receiver.await.unwrap();
101                assert_eq!(progress.error.is_some(), fail);
102                assert_eq!(progress.affected_rows, if fail { 0 } else { 1 });
103                if let Some(error) = progress.error {
104                    assert!(error.to_string().contains("test flush failure"));
105                }
106            }
107        }
108    }
109}
servers/batcher/table/pending_batch.rs

servers/batcher/table/
pending_batch.rs
