operator/write_dispatcher/
progress.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::{HashMap, HashSet};
16use std::sync::Arc;
17
18use api::v1::region::{RegionRequest, region_request::Body};
19use backon::{BackoffBuilder, ExponentialBuilder};
20use common_error::ext::{BoxedError, ErrorExt, RetryHint};
21use common_meta::error::{Error, UnexpectedSnafu};
22use common_meta::peer::PeerRef;
23use common_query::WriteProgress;
24use futures::future;
25use prometheus::Histogram;
26use tokio_util::sync::CancellationToken;
27
28use crate::write_dispatcher::{
29    RoutedWrite, WriteAttempt, WriteAttribution, WriteDispatcher, get_peer, write_region_id,
30};
31
32/// A region payload and its original-input indices, in payload row order.
33/// Ownership stays local; it is never added to the schema or sent to storage.
34pub struct TrackedWrite {
35    pub request: RegionRequest,
36    pub owners: Vec<usize>,
37}
38
39/// A wire request retaining one row-owner vector per region entry.
40struct TrackedBatch {
41    request: RegionRequest,
42    owners: Vec<Vec<usize>>,
43}
44
45impl WriteDispatcher {
46    /// Collects acknowledgements per input across shared flushes and repartition retries.
47    /// Records each datanode send attempt, including retries, in the optional histogram.
48    pub async fn dispatch_tracked(
49        &self,
50        requests: Vec<TrackedWrite>,
51        inputs: usize,
52        histogram: Option<&Histogram>,
53    ) -> Vec<WriteProgress> {
54        let (writes, progress) = self.prepare_tracked(requests, inputs).await;
55        self.execute_tracked(writes, progress, histogram, None)
56            .await
57    }
58
59    /// Resolves routes before tracked writes are detached from the caller.
60    pub(crate) async fn prepare_tracked(
61        &self,
62        requests: Vec<TrackedWrite>,
63        inputs: usize,
64    ) -> (Vec<RoutedWrite>, Vec<WriteProgress>) {
65        let mut progress = vec![WriteProgress::default(); inputs];
66        if requests
67            .iter()
68            .any(|request| request.owners.iter().any(|owner| *owner >= inputs))
69        {
70            let error = UnexpectedSnafu {
71                err_msg: "Write owner exceeds original input count",
72            }
73            .build();
74            return (
75                Vec::new(),
76                vec![WriteProgress::failed(BoxedError::new(error)); inputs],
77            );
78        }
79        let regions = match requests
80            .iter()
81            .map(|write| write_region_id(&write.request))
82            .collect::<common_meta::error::Result<Vec<_>>>()
83        {
84            Ok(regions) => regions,
85            Err(error) => {
86                return (
87                    Vec::new(),
88                    vec![WriteProgress::failed(BoxedError::new(error)); inputs],
89                );
90            }
91        };
92        let peers = match self.resolve_peers(&regions).await {
93            Ok(peers) => peers,
94            Err(error) => {
95                return (
96                    Vec::new(),
97                    vec![WriteProgress::failed(BoxedError::new(error)); inputs],
98                );
99            }
100        };
101        let mut pending = Vec::with_capacity(requests.len());
102        for (write, region) in requests.into_iter().zip(regions) {
103            match get_peer(&peers, region) {
104                Ok(peer) => pending.push(RoutedWrite {
105                    peer,
106                    request: Arc::new(write.request),
107                    attribution: WriteAttribution::default(),
108                    owners: Some(write.owners),
109                }),
110                Err(error) => record_failure(&mut progress, &write.owners, error),
111            }
112        }
113        (pending, progress)
114    }
115
116    /// Drains in-flight writes while cancellation stops subsequent retries.
117    pub(crate) async fn execute_tracked(
118        &self,
119        mut pending: Vec<RoutedWrite>,
120        mut progress: Vec<WriteProgress>,
121        histogram: Option<&Histogram>,
122        cancellation: Option<&CancellationToken>,
123    ) -> Vec<WriteProgress> {
124        let mut backoff = self.retry.as_ref().map(|retry| {
125            ExponentialBuilder::default()
126                .with_min_delay(retry.options.initial_backoff)
127                .with_max_delay(retry.options.max_backoff)
128                .with_factor(retry.options.backoff_multiplier)
129                .with_max_times(retry.options.max_retries)
130                .with_jitter()
131                .build()
132        });
133        let mut failures = Vec::new();
134        loop {
135            failures.clear();
136            // Partial responses can be attributed directly when all rows belong to one input.
137            let attempts = if progress.len() == 1
138                && self.retry.is_some()
139                && pending.len() > 1
140                && pending
141                    .iter()
142                    .any(|write| matches!(write.request.body, Some(Body::Inserts(_))))
143            {
144                self.send_tracked_batches(pending, histogram).await
145            } else {
146                future::join_all(
147                    pending
148                        .into_iter()
149                        .map(|write| self.send_tracked_write(write, histogram)),
150                )
151                .await
152            };
153            pending = Vec::new();
154            let mut refresh = Vec::new();
155            let mut terminal = false;
156            for (owners, attempt) in attempts {
157                match attempt {
158                    WriteAttempt::Success(response) if response.affected_rows == owners.len() => {
159                        for run in owners.chunk_by(|a, b| a == b) {
160                            progress[run[0]].affected_rows += run.len();
161                        }
162                    }
163                    WriteAttempt::Success(response) => {
164                        terminal = true;
165                        record_failure(&mut progress, &owners, UnexpectedSnafu {
166                            err_msg: format!("Write acknowledged {} rows, expected {}; cannot attribute this acknowledgement", response.affected_rows, owners.len()),
167                        }.build());
168                    }
169                    WriteAttempt::Failed {
170                        error,
171                        request: Some(write),
172                    } => {
173                        if error.retry_hint() == RetryHint::RetryAfterRouteRefresh {
174                            refresh.push(write);
175                        } else {
176                            pending.push(write);
177                        }
178                        failures.push((owners, error));
179                    }
180                    WriteAttempt::Failed {
181                        error,
182                        request: None,
183                    } => {
184                        terminal = true;
185                        record_failure(&mut progress, &owners, error);
186                    }
187                }
188            }
189            if failures.is_empty() {
190                break;
191            }
192            let delay = (!terminal)
193                .then(|| backoff.as_mut().and_then(Iterator::next))
194                .flatten();
195            let (Some(delay), Some(retry)) = (delay, &self.retry) else {
196                break;
197            };
198            if let Some(cancellation) = cancellation {
199                tokio::select! {
200                    biased;
201                    _ = cancellation.cancelled() => break,
202                    _ = tokio::time::sleep(delay) => {}
203                }
204            } else {
205                tokio::time::sleep(delay).await;
206            }
207            if cancellation.is_some_and(CancellationToken::is_cancelled) {
208                break;
209            }
210            if !refresh.is_empty() {
211                let owners = refresh
212                    .iter()
213                    .flat_map(|write| write.owners.iter().flatten().copied())
214                    .collect::<Vec<_>>();
215                match retry.reroute_failed_writes(refresh).await {
216                    Ok(writes) => pending.extend(writes),
217                    Err(error) => record_failure(&mut progress, &owners, error),
218                }
219            }
220            if pending.is_empty() || cancellation.is_some_and(CancellationToken::is_cancelled) {
221                break;
222            }
223        }
224        for (owners, error) in failures {
225            record_failure(&mut progress, &owners, error);
226        }
227        progress
228    }
229
230    async fn send_tracked_write(
231        &self,
232        mut write: RoutedWrite,
233        histogram: Option<&Histogram>,
234    ) -> (Vec<usize>, WriteAttempt) {
235        let owners = write.owners.take().unwrap_or_default();
236        let mut attempt = if let Some(retry) = &self.retry {
237            retry.send(write, histogram).await
238        } else {
239            let _timer = histogram.map(Histogram::start_timer);
240            let datanode = self.node_manager.datanode(&write.peer).await;
241            match datanode.handle(write.request).await {
242                Ok(response) => WriteAttempt::Success(response),
243                Err(error) => WriteAttempt::Failed {
244                    error,
245                    request: None,
246                },
247            }
248        };
249        if let WriteAttempt::Failed {
250            request: Some(write),
251            ..
252        } = &mut attempt
253        {
254            write.owners = Some(owners.clone());
255        }
256        (owners, attempt)
257    }
258
259    async fn send_tracked_batches(
260        &self,
261        writes: Vec<RoutedWrite>,
262        histogram: Option<&Histogram>,
263    ) -> Vec<(Vec<usize>, WriteAttempt)> {
264        let mut batches: HashMap<PeerRef, Vec<TrackedBatch>> = HashMap::new();
265        let mut singles = Vec::new();
266        let mut attempts = Vec::new();
267        for write in writes {
268            if !matches!(write.request.body, Some(Body::Inserts(_))) {
269                singles.push(write);
270                continue;
271            }
272            let owners = write.owners.unwrap_or_default();
273            let mut request = match Arc::try_unwrap(write.request) {
274                Ok(request) => request,
275                Err(_) => {
276                    attempts.push((
277                        owners,
278                        WriteAttempt::Failed {
279                            error: UnexpectedSnafu {
280                                err_msg: "Region request is still shared before batching",
281                            }
282                            .build(),
283                            request: None,
284                        },
285                    ));
286                    continue;
287                }
288            };
289            let batches = batches.entry(write.peer).or_default();
290            if let Some(batch) = batches
291                .iter_mut()
292                .find(|batch| batch.request.header == request.header)
293                && let (Some(Body::Inserts(existing)), Some(Body::Inserts(next))) =
294                    (&mut batch.request.body, &mut request.body)
295            {
296                existing.requests.append(&mut next.requests);
297                batch.owners.push(owners);
298            } else {
299                batches.push(TrackedBatch {
300                    request,
301                    owners: vec![owners],
302                });
303            }
304        }
305        let batches = future::join_all(batches.into_iter().flat_map(|(peer, batches)| {
306            batches.into_iter().map(move |batch| {
307                self.send_batch(
308                    Arc::clone(&peer),
309                    batch.request,
310                    batch.owners,
311                    histogram,
312                    None,
313                )
314            })
315        }));
316        let singles = future::join_all(
317            singles
318                .into_iter()
319                .map(|write| self.send_tracked_write(write, histogram)),
320        );
321        let (batches, singles) = future::join(batches, singles).await;
322        attempts.extend(batches.into_iter().flatten());
323        attempts.extend(singles);
324        attempts
325    }
326}
327
328fn record_failure(progress: &mut [WriteProgress], owners: &[usize], error: Error) {
329    let error = Arc::new(BoxedError::new(error));
330    for owner in owners.iter().copied().collect::<HashSet<_>>() {
331        progress[owner].error.get_or_insert_with(|| error.clone());
332    }
333}
operator/write_dispatcher/progress.rs

operator/write_dispatcher/
progress.rs
