operator/statement/export_logical_tables/
writers.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_datasource::packed_writer::PackedWriterRef;
19use futures::stream::FuturesUnordered;
20use futures::{FutureExt, StreamExt};
21use object_store::ObjectStore;
22use snafu::{OptionExt, ResultExt, ensure};
23use tokio::sync::{OwnedSemaphorePermit, Semaphore, mpsc};
24use tokio::task::JoinHandle;
25use tokio_util::sync::CancellationToken;
26
27use crate::error::{self, Result};
28use crate::statement::export_logical_tables::{
29    ActiveWriter, LogicalTableExportLimits, LogicalTableProjection, check_cancelled,
30    map_writer_error,
31};
32
33const PAYLOAD_BYTES: usize = 64 * 1024 * 1024;
34
35/// Request-wide downstream payload and writer admission, separate from query memory.
36pub(crate) struct ExportWriteBudget {
37    writers: Arc<Semaphore>,
38    bytes: Arc<Semaphore>,
39    max_writers: usize,
40}
41
42impl ExportWriteBudget {
43    pub(crate) fn new(parallelism: usize) -> Arc<Self> {
44        let max_writers = parallelism.max(1);
45        Arc::new(Self {
46            writers: Arc::new(Semaphore::new(max_writers)),
47            bytes: Arc::new(Semaphore::new(PAYLOAD_BYTES)),
48            max_writers,
49        })
50    }
51
52    #[cfg(test)]
53    pub(crate) fn available(&self) -> (usize, usize) {
54        (
55            self.writers.available_permits(),
56            self.bytes.available_permits(),
57        )
58    }
59
60    pub(crate) async fn writer(&self, token: &CancellationToken) -> Result<OwnedSemaphorePermit> {
61        tokio::select! {
62            biased;
63            _ = token.cancelled() => error::LogicalTableExportCancelledSnafu.fail(),
64            permit = self.writers.clone().acquire_owned() => permit.map_err(|_| error::UnexpectedSnafu { violated: "writer budget closed" }.build()),
65        }
66    }
67
68    pub(crate) async fn reserve(
69        &self,
70        size: usize,
71        token: &CancellationToken,
72    ) -> Result<OwnedSemaphorePermit> {
73        ensure!(
74            size <= PAYLOAD_BYTES,
75            error::LogicalTableExportResourceSnafu {
76                reason: "batch backing buffers exceed request payload budget"
77            }
78        );
79        tokio::select! {
80            biased;
81            _ = token.cancelled() => error::LogicalTableExportCancelledSnafu.fail(),
82            permit = self.bytes.clone().acquire_many_owned(size as u32) => permit.map_err(|_| error::UnexpectedSnafu { violated: "payload budget closed" }.build()),
83        }
84    }
85
86    /// Accounts for retained input allocations, conversion capacity growth and
87    /// array metadata. Pass zero backing when the caller detaches oversized slices.
88    pub(crate) fn conversion_budget(
89        backing: usize,
90        columns: usize,
91        requested: usize,
92    ) -> Result<(usize, usize)> {
93        let overhead = columns.saturating_mul(1024);
94        let available = PAYLOAD_BYTES
95            .saturating_sub(backing)
96            .saturating_sub(overhead)
97            / 4;
98        let conversion = requested.min(available);
99        ensure!(
100            conversion > 0,
101            error::LogicalTableExportResourceSnafu {
102                reason: "batch backing buffers exceed request payload budget"
103            }
104        );
105        Ok((conversion, backing.saturating_add(overhead)))
106    }
107}
108
109pub(crate) struct Payload {
110    pub(crate) batch: RecordBatch,
111    pub(crate) permit: OwnedSemaphorePermit,
112}
113
114pub(crate) struct TableWriters {
115    pub(crate) packed: Option<PackedWriterRef>,
116    current: Option<(u32, mpsc::Sender<Payload>)>,
117    tasks: FuturesUnordered<JoinHandle<Result<()>>>,
118    budget: Arc<ExportWriteBudget>,
119}
120
121impl TableWriters {
122    pub(crate) fn new(budget: Arc<ExportWriteBudget>) -> Self {
123        Self {
124            current: None,
125            packed: None,
126            tasks: FuturesUnordered::new(),
127            budget,
128        }
129    }
130
131    #[cfg(test)]
132    pub(crate) fn pending_tasks(&self) -> usize {
133        self.tasks.len()
134    }
135
136    pub(crate) fn table_id(&self) -> Option<u32> {
137        self.current.as_ref().map(|(id, _)| *id)
138    }
139
140    pub(crate) fn close_input(&mut self) {
141        self.current = None;
142    }
143
144    pub(crate) async fn open(
145        &mut self,
146        id: u32,
147        table: &LogicalTableProjection,
148        store: &ObjectStore,
149        limits: LogicalTableExportLimits,
150        token: &CancellationToken,
151    ) -> Result<()> {
152        // EOF must precede acquiring the next slot, including when P is one.
153        self.close_input();
154        let permit = self.budget.writer(token).await?;
155        self.reap_for_admission(token).await?;
156        let writer = if let Some(packed) = &self.packed {
157            ActiveWriter::open_packed(table, id, store, limits, packed.clone())?
158        } else {
159            ActiveWriter::open(table, store, limits).await?
160        };
161        let (sender, receiver) = mpsc::channel(2);
162        self.current = Some((id, sender));
163        let token = token.clone();
164        self.tasks.push(common_runtime::spawn_global(async move {
165            let _permit = permit;
166            let guard = token.clone().drop_guard();
167            let result = run_writer(writer, receiver, &token).await;
168            if result.is_ok() {
169                guard.disarm();
170            }
171            result
172        }));
173        Ok(())
174    }
175
176    async fn reap_for_admission(&mut self, token: &CancellationToken) -> Result<()> {
177        while let Some(Some(result)) = self.tasks.next().now_or_never() {
178            result.context(error::JoinTaskSnafu)??;
179        }
180        // A worker can release its permit before its JoinHandle becomes ready.
181        while self.tasks.len() > self.budget.max_writers {
182            let result = tokio::select! {
183                biased;
184                _ = token.cancelled() => return error::LogicalTableExportCancelledSnafu.fail(),
185                result = self.tasks.next() => result.context(error::UnexpectedSnafu { violated: "writer task queue unexpectedly empty" })?,
186            };
187            result.context(error::JoinTaskSnafu)??;
188        }
189        Ok(())
190    }
191
192    pub(crate) async fn send(&self, payload: Payload, token: &CancellationToken) -> Result<()> {
193        let (_, sender) = self.current.as_ref().context(error::UnexpectedSnafu {
194            violated: "missing logical writer",
195        })?;
196        tokio::select! {
197            biased;
198            _ = token.cancelled() => error::LogicalTableExportCancelledSnafu.fail(),
199            result = sender.send(payload) => result.map_err(|_| error::LogicalTableExportCancelledSnafu.build()),
200        }
201    }
202
203    pub(crate) async fn drain(
204        &mut self,
205        result: Result<()>,
206        token: &CancellationToken,
207    ) -> Result<()> {
208        self.close_input();
209        let mut first_error = result.err();
210        if first_error.is_some() {
211            token.cancel();
212        }
213        while let Some(result) = self.tasks.next().await {
214            if let Err(err) = result.context(error::JoinTaskSnafu).and_then(|r| r) {
215                retain_error(&mut first_error, err);
216                token.cancel();
217            }
218        }
219        match first_error {
220            Some(err) => Err(err),
221            None => Ok(()),
222        }
223    }
224}
225
226pub(crate) fn retain_error(first: &mut Option<error::Error>, error: error::Error) {
227    if first.as_ref().is_none_or(|err| {
228        matches!(
229            err,
230            error::Error::LogicalTableExportCancelled { .. }
231                | error::Error::DatabaseExportCancelled { .. }
232        )
233    }) {
234        *first = Some(error);
235    }
236}
237
238async fn run_writer(
239    mut writer: ActiveWriter,
240    mut receiver: mpsc::Receiver<Payload>,
241    token: &CancellationToken,
242) -> Result<()> {
243    let result = async {
244        loop {
245            let payload = tokio::select! {
246                biased;
247                _ = token.cancelled() => return error::LogicalTableExportCancelledSnafu.fail(),
248                payload = receiver.recv() => payload,
249            };
250            let Some(Payload { batch, permit }) = payload else {
251                break;
252            };
253            let result = writer.writer.write(batch, Some(token)).await;
254            drop(permit);
255            result.map_err(|error| map_writer_error(error, &writer.path))?;
256        }
257        writer
258            .writer
259            .finish(Some(token))
260            .await
261            .map_err(|error| map_writer_error(error, &writer.path))?;
262        check_cancelled(token)
263    }
264    .await;
265    if result.is_err() {
266        token.cancel();
267        receiver.close();
268        drop(receiver);
269        if let Err(error) = writer.writer.abort().await {
270            common_telemetry::warn!(error; "Failed to abort Metric export file");
271        }
272    }
273    result
274}
275
276#[cfg(test)]
277mod tests {
278    use tokio::sync::oneshot;
279
280    use super::*;
281
282    #[tokio::test]
283    async fn released_permit_does_not_hide_pending_join_handles() {
284        for parallelism in [1, 4] {
285            let budget = ExportWriteBudget::new(parallelism);
286            let token = CancellationToken::new();
287            let mut writers = TableWriters::new(budget.clone());
288            let mut resume = Vec::new();
289            for _ in 0..=parallelism {
290                let permit = budget.writer(&token).await.unwrap();
291                let (released_tx, released_rx) = oneshot::channel();
292                let (resume_tx, resume_rx) = oneshot::channel();
293                writers.tasks.push(common_runtime::spawn_global(async move {
294                    drop(permit);
295                    released_tx.send(()).unwrap();
296                    resume_rx.await.unwrap();
297                    Ok(())
298                }));
299                released_rx.await.unwrap();
300                resume.push(resume_tx);
301            }
302            assert_eq!(budget.available().0, parallelism);
303
304            let mut reap = Box::pin(writers.reap_for_admission(&token));
305            assert!(reap.as_mut().now_or_never().is_none());
306            resume.pop().unwrap().send(()).unwrap();
307            reap.await.unwrap();
308            assert_eq!(writers.pending_tasks(), parallelism);
309
310            for tx in resume {
311                tx.send(()).unwrap();
312            }
313            writers.drain(Ok(()), &token).await.unwrap();
314            assert_eq!(budget.available().0, parallelism);
315        }
316    }
317}
operator/statement/export_logical_tables/writers.rs

operator/statement/export_logical_tables/
writers.rs
