operator/
write_dispatcher.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
15mod options;
16#[cfg(test)]
17mod tests;
18
19use std::collections::{HashMap, HashSet, hash_map::Entry};
20use std::sync::Arc;
21
22use api::region::RegionResponse;
23use api::v1::region::{DeleteRequests, InsertRequests, RegionRequest, region_request::Body};
24use backon::{BackoffBuilder, ExponentialBuilder};
25use catalog::CatalogManagerRef;
26use common_error::ext::{BoxedError, ErrorExt, RetryHint};
27use common_meta::cache_invalidator::{CacheInvalidatorRef, Context};
28use common_meta::error::{Error, ExternalSnafu, Result, TableInfoNotFoundSnafu, UnexpectedSnafu};
29use common_meta::instruction::CacheIdent;
30use common_meta::node_manager::NodeManagerRef;
31use common_meta::peer::{Peer, PeerRef};
32use common_meta::rpc::router;
33use futures::future;
34use partition::manager::{PartitionRuleManager, PartitionRuleManagerRef};
35use prometheus::Histogram;
36use snafu::{OptionExt, ResultExt, ensure};
37use store_api::storage::RegionId;
38use table::table_name::TableName;
39use tokio_util::sync::CancellationToken;
40
41pub use self::options::WriteRetryOptions;
42use crate::error::{self, FindRegionLeaderSnafu, RequestRegionSnafu};
43use crate::req_convert::TablePartitioner;
44
45/// Dispatches already partitioned writes, retaining only failed payloads for retries.
46pub struct WriteDispatcher {
47    node_manager: NodeManagerRef,
48    partition_manager: PartitionRuleManagerRef,
49    retry: Option<WriteRetry>,
50}
51
52/// Shared reference to a [`WriteDispatcher`].
53pub type WriteDispatcherRef = Arc<WriteDispatcher>;
54
55/// The result of one write RPC attempt.
56pub(crate) enum WriteAttempt {
57    Success(RegionResponse),
58    Failed {
59        error: Error,
60        /// Retained when retries are enabled and the original client error allows retry.
61        /// The remaining retry budget is checked separately.
62        request: Option<RoutedWrite>,
63    },
64}
65
66impl WriteDispatcher {
67    /// Creates an owned-request dispatcher without retries, including for standalone.
68    pub fn new(node_manager: NodeManagerRef, partition_manager: PartitionRuleManagerRef) -> Self {
69        Self {
70            node_manager,
71            partition_manager,
72            retry: None,
73        }
74    }
75
76    /// Enables shared remote sends and retries using frontend configuration and caches.
77    pub fn with_retry(
78        mut self,
79        options: WriteRetryOptions,
80        catalog_manager: CatalogManagerRef,
81        cache_invalidator: CacheInvalidatorRef,
82    ) -> Self {
83        if options.max_retries > 0 {
84            self.retry = Some(WriteRetry {
85                node_manager: self.node_manager.clone(),
86                partition_manager: self.partition_manager.clone(),
87                options,
88                catalog_manager,
89                cache_invalidator,
90            });
91        } else {
92            self.retry = None;
93        }
94        self
95    }
96
97    /// Writes planned requests without repeating admission, default values, or Flow side effects.
98    pub async fn dispatch(&self, requests: Vec<RegionRequest>) -> Result<RegionResponse> {
99        self.dispatch_inner(requests, None).await
100    }
101
102    /// Records each datanode send attempt, including retries, in the supplied histogram.
103    pub async fn dispatch_with_metrics(
104        &self,
105        requests: Vec<RegionRequest>,
106        histogram: Histogram,
107    ) -> Result<RegionResponse> {
108        self.dispatch_inner(requests, Some(histogram)).await
109    }
110
111    /// Resolves routes before writes are detached from the caller.
112    pub(crate) async fn prepare(
113        &self,
114        requests: Vec<RegionRequest>,
115    ) -> error::Result<Vec<(PeerRef, RegionRequest)>> {
116        let requests = requests
117            .into_iter()
118            .map(|mut request| {
119                let header = request.header.take();
120                (header, split_request(request))
121            })
122            .collect::<Vec<_>>();
123        let regions = requests
124            .iter()
125            .flat_map(|(_, requests)| requests)
126            .map(write_region_id)
127            .collect::<Result<Vec<_>>>()
128            .context(RequestRegionSnafu)?;
129        let peers = self
130            .resolve_peers(&regions)
131            .await
132            .context(FindRegionLeaderSnafu)?;
133        let mut writes = Vec::new();
134        // Preserve each original envelope and peer batching when retries are disabled.
135        for (header, requests) in requests {
136            let mut by_peer: HashMap<PeerRef, RegionRequest> = HashMap::new();
137            for request in requests {
138                let peer = get_peer(
139                    &peers,
140                    write_region_id(&request).context(RequestRegionSnafu)?,
141                )
142                .context(RequestRegionSnafu)?;
143                match by_peer.entry(peer) {
144                    Entry::Vacant(entry) => {
145                        entry.insert(request);
146                    }
147                    Entry::Occupied(mut entry) => match (&mut entry.get_mut().body, request.body) {
148                        (Some(Body::Inserts(existing)), Some(Body::Inserts(next))) => {
149                            existing.requests.extend(next.requests)
150                        }
151                        (Some(Body::Deletes(existing)), Some(Body::Deletes(next))) => {
152                            existing.requests.extend(next.requests)
153                        }
154                        _ => {
155                            return UnexpectedSnafu {
156                                err_msg: "Cannot combine different write kinds",
157                            }
158                            .fail()
159                            .context(RequestRegionSnafu);
160                        }
161                    },
162                }
163            }
164            for (peer, mut request) in by_peer {
165                request.header = header.clone();
166                writes.push((peer, request));
167            }
168        }
169        Ok(writes)
170    }
171
172    /// Sends requests using routes already resolved by the caller.
173    /// Cancellation stops future retries while allowing in-flight requests to complete.
174    pub(crate) async fn execute(
175        &self,
176        writes: Vec<(PeerRef, RegionRequest)>,
177        histogram: Option<Histogram>,
178        cancellation: Option<&CancellationToken>,
179    ) -> Result<RegionResponse> {
180        let histogram = histogram.as_ref();
181        let writes = if self.retry.is_some() {
182            writes
183                .into_iter()
184                .flat_map(|(peer, request)| {
185                    split_request(request)
186                        .into_iter()
187                        .map(move |request| (Arc::clone(&peer), request))
188                })
189                .collect()
190        } else {
191            writes
192        };
193        let attempts = future::join_all(
194            writes
195                .into_iter()
196                .map(|(peer, request)| self.send(peer, request, histogram)),
197        )
198        .await;
199        self.finish(attempts, histogram, cancellation).await
200    }
201
202    /// Sends the first attempt immediately so bulk encoding can overlap other sends.
203    pub(crate) async fn send(
204        &self,
205        peer: PeerRef,
206        request: RegionRequest,
207        histogram: Option<&Histogram>,
208    ) -> WriteAttempt {
209        if let Some(retry) = &self.retry {
210            return retry
211                .send(
212                    RoutedWrite {
213                        peer,
214                        request: Arc::new(request),
215                    },
216                    histogram,
217                )
218                .await;
219        }
220        let _timer = histogram.map(Histogram::start_timer);
221        let datanode = self.node_manager.datanode(&peer).await;
222        match datanode.handle(request).await {
223            Ok(response) => WriteAttempt::Success(response),
224            Err(error) => WriteAttempt::Failed {
225                error,
226                request: None,
227            },
228        }
229    }
230
231    /// Collects first attempts and retries their failures together, refreshing routes once per round.
232    pub(crate) async fn finish(
233        &self,
234        attempts: Vec<WriteAttempt>,
235        histogram: Option<&Histogram>,
236        cancellation: Option<&CancellationToken>,
237    ) -> Result<RegionResponse> {
238        if let Some(retry) = &self.retry {
239            if retry.should_retry(&attempts, 0) {
240                return retry.retry(attempts, histogram, cancellation).await;
241            }
242            let mut response = RegionResponse::new(0);
243            let failures = merge_attempts(attempts, &mut response)?;
244            return match failures.into_iter().next() {
245                Some((_, error)) => Err(error),
246                None => Ok(response),
247            };
248        }
249        let mut response = RegionResponse::new(0);
250        for attempt in attempts {
251            match attempt {
252                WriteAttempt::Success(result) => merge_response(&mut response, result),
253                WriteAttempt::Failed { error, .. } => return Err(error),
254            }
255        }
256        Ok(response)
257    }
258
259    /// Resolves only the requested peers, loading each table route once.
260    pub(crate) async fn resolve_peers(
261        &self,
262        regions: &[RegionId],
263    ) -> partition::error::Result<HashMap<RegionId, PeerRef>> {
264        resolve_peers(&self.partition_manager, regions).await
265    }
266
267    async fn dispatch_inner(
268        &self,
269        requests: Vec<RegionRequest>,
270        histogram: Option<Histogram>,
271    ) -> Result<RegionResponse> {
272        let writes = self
273            .prepare(requests)
274            .await
275            .map_err(BoxedError::new)
276            .context(ExternalSnafu)?;
277        self.execute(writes, histogram, None).await
278    }
279}
280
281/// Sends shared requests and coordinates retries with refreshed routes.
282struct WriteRetry {
283    node_manager: NodeManagerRef,
284    partition_manager: PartitionRuleManagerRef,
285    options: WriteRetryOptions,
286    catalog_manager: CatalogManagerRef,
287    cache_invalidator: CacheInvalidatorRef,
288}
289
290/// A single-region write request paired with its destination peer.
291pub(crate) struct RoutedWrite {
292    peer: PeerRef,
293    request: Arc<RegionRequest>,
294}
295
296impl WriteRetry {
297    async fn send(&self, write: RoutedWrite, histogram: Option<&Histogram>) -> WriteAttempt {
298        let _timer = histogram.map(Histogram::start_timer);
299        let requester = self.node_manager.datanode(&write.peer).await;
300        match requester.handle(Arc::clone(&write.request)).await {
301            Ok(response) => WriteAttempt::Success(response),
302            Err(error) => {
303                let request = error.retry_hint().is_retryable().then_some(write);
304                WriteAttempt::Failed { error, request }
305            }
306        }
307    }
308
309    fn should_retry(&self, attempts: &[WriteAttempt], retries: usize) -> bool {
310        if retries >= self.options.max_retries {
311            return false;
312        }
313        let mut retryable = false;
314        for attempt in attempts {
315            match attempt {
316                WriteAttempt::Success(_) => {}
317                WriteAttempt::Failed {
318                    request: Some(_), ..
319                } => retryable = true,
320                WriteAttempt::Failed { request: None, .. } => return false,
321            }
322        }
323        retryable
324    }
325
326    async fn retry(
327        &self,
328        mut attempts: Vec<WriteAttempt>,
329        histogram: Option<&Histogram>,
330        cancellation: Option<&CancellationToken>,
331    ) -> Result<RegionResponse> {
332        let mut response = RegionResponse::new(0);
333        let mut backoff = ExponentialBuilder::default()
334            .with_min_delay(self.options.initial_backoff)
335            .with_max_delay(self.options.max_backoff)
336            .with_factor(self.options.backoff_multiplier)
337            .with_max_times(self.options.max_retries)
338            .with_jitter()
339            .build();
340        let mut retries = 0;
341        while self.should_retry(&attempts, retries) {
342            let Some(delay) = backoff.next() else {
343                break;
344            };
345            if let Some(cancellation) = cancellation {
346                tokio::select! {
347                    biased;
348                    _ = cancellation.cancelled() => break,
349                    _ = tokio::time::sleep(delay) => {}
350                }
351            } else {
352                tokio::time::sleep(delay).await;
353            }
354            if cancellation.is_some_and(CancellationToken::is_cancelled) {
355                break;
356            }
357            let failures = merge_attempts(attempts, &mut response)?;
358            let mut first_error = None;
359            let mut refresh = Vec::new();
360            let mut pending = Vec::new();
361            for (write, error) in failures {
362                if error.retry_hint() == RetryHint::RetryAfterRouteRefresh {
363                    refresh.push(write);
364                } else {
365                    pending.push(write);
366                }
367                first_error.get_or_insert(error);
368            }
369            if !refresh.is_empty() {
370                pending.extend(self.reroute_failed_writes(refresh).await?);
371            }
372            if let Some(error) =
373                first_error.filter(|_| cancellation.is_some_and(CancellationToken::is_cancelled))
374            {
375                return Err(error);
376            }
377            attempts =
378                future::join_all(pending.into_iter().map(|write| self.send(write, histogram)))
379                    .await;
380            retries += 1;
381        }
382        match merge_attempts(attempts, &mut response)?.into_iter().next() {
383            Some((_, error)) => Err(error),
384            None => Ok(response),
385        }
386    }
387
388    async fn reroute_failed_writes(&self, writes: Vec<RoutedWrite>) -> Result<Vec<RoutedWrite>> {
389        let regions = writes
390            .iter()
391            .map(|write| write_region_id(&write.request))
392            .collect::<Result<Vec<_>>>()?;
393        self.invalidate_routes(&regions).await?;
394        let mut table_names = HashMap::new();
395        for table_id in regions
396            .iter()
397            .map(|region| region.table_id())
398            .collect::<HashSet<_>>()
399        {
400            // Raw metadata is only used to recover the name of the failed request's table.
401            let table_info = self
402                .catalog_manager
403                .table_info_by_id(table_id)
404                .await
405                .map_err(BoxedError::new)
406                .context(ExternalSnafu)?
407                .with_context(|| TableInfoNotFoundSnafu {
408                    table: table_id.to_string(),
409                })?;
410            table_names.insert(
411                table_id,
412                TableName::new(
413                    &table_info.catalog_name,
414                    &table_info.schema_name,
415                    &table_info.name,
416                ),
417            );
418        }
419        // TableId invalidation refreshes metadata and routes, but TableCache is keyed by name.
420        // Use the refreshed names in case a table was renamed while the write was pending.
421        let caches = table_names
422            .values()
423            .cloned()
424            .map(CacheIdent::TableName)
425            .collect::<Vec<_>>();
426        self.cache_invalidator
427            .invalidate(&Context::default(), &caches)
428            .await?;
429        let mut partitioners = HashMap::new();
430        for (table_id, name) in table_names {
431            let table = self
432                .catalog_manager
433                .table(
434                    &name.catalog_name,
435                    &name.schema_name,
436                    &name.table_name,
437                    None,
438                )
439                .await
440                .map_err(BoxedError::new)
441                .context(ExternalSnafu)?
442                .with_context(|| TableInfoNotFoundSnafu {
443                    table: name.to_string(),
444                })?;
445            let table_info = table.table_info();
446            ensure!(
447                table_info.table_id() == table_id,
448                UnexpectedSnafu {
449                    err_msg: format!(
450                        "Table {name} changed ID from {table_id} to {} while retrying writes",
451                        table_info.table_id()
452                    ),
453                }
454            );
455            let partitioner = TablePartitioner::try_new(&self.partition_manager, &table_info)
456                .await
457                .map_err(BoxedError::new)
458                .context(ExternalSnafu)?;
459            partitioners.insert(table_id, (table_info, partitioner));
460        }
461        let mut rebuilt_requests = Vec::new();
462        for write in writes {
463            let request = Arc::try_unwrap(write.request).map_err(|_| {
464                UnexpectedSnafu {
465                    err_msg: "Write request is still shared after its RPC completed",
466                }
467                .build()
468            })?;
469            let table_id = write_region_id(&request)?.table_id();
470            let (table_info, partitioner) =
471                partitioners.get(&table_id).context(UnexpectedSnafu {
472                    err_msg: "Refreshed write has no partition snapshot",
473                })?;
474            let requests = partitioner
475                .rebuild_write_requests(request, table_info)
476                .map_err(BoxedError::new)
477                .context(ExternalSnafu)?;
478            rebuilt_requests.extend(requests);
479        }
480        self.prepare(rebuilt_requests).await
481    }
482
483    async fn prepare(&self, requests: Vec<RegionRequest>) -> Result<Vec<RoutedWrite>> {
484        let regions = requests
485            .iter()
486            .map(write_region_id)
487            .collect::<Result<Vec<_>>>()?;
488        let peers = resolve_peers(&self.partition_manager, &regions)
489            .await
490            .map_err(BoxedError::new)
491            .context(ExternalSnafu)?;
492        requests
493            .into_iter()
494            .zip(regions)
495            .map(|(request, region)| {
496                Ok(RoutedWrite {
497                    peer: get_peer(&peers, region)?,
498                    request: Arc::new(request),
499                })
500            })
501            .collect()
502    }
503
504    async fn invalidate_routes(&self, regions: &[RegionId]) -> Result<()> {
505        let mut table_ids = regions
506            .iter()
507            .map(|region| region.table_id())
508            .collect::<HashSet<_>>();
509        let mut physical_ids = HashSet::new();
510        for table_id in &table_ids {
511            let (physical_id, _) = self
512                .partition_manager
513                .find_physical_table_route_with_id(*table_id)
514                .await
515                .map_err(BoxedError::new)
516                .context(ExternalSnafu)?;
517            physical_ids.insert(physical_id);
518        }
519        table_ids.extend(physical_ids);
520        let caches = table_ids
521            .into_iter()
522            .map(CacheIdent::TableId)
523            .collect::<Vec<_>>();
524        self.cache_invalidator
525            .invalidate(&Context::default(), &caches)
526            .await
527    }
528}
529
530async fn resolve_peers(
531    partition_manager: &PartitionRuleManager,
532    regions: &[RegionId],
533) -> partition::error::Result<HashMap<RegionId, PeerRef>> {
534    let table_ids = regions
535        .iter()
536        .map(|region| region.table_id())
537        .collect::<HashSet<_>>();
538    let mut physical_tables = HashMap::new();
539    let mut routes = HashMap::new();
540    for table_id in table_ids {
541        let (physical_id, route) = partition_manager
542            .find_physical_table_route_with_id(table_id)
543            .await?;
544        physical_tables.insert(table_id, physical_id);
545        routes.entry(physical_id).or_insert(route);
546    }
547    let mut shared_peers: HashMap<&Peer, PeerRef> = HashMap::new();
548    let mut peers = HashMap::with_capacity(regions.len());
549    for region in regions {
550        if peers.contains_key(region) {
551            continue;
552        }
553        let peer = physical_tables
554            .get(&region.table_id())
555            .and_then(|physical_id| routes.get(physical_id))
556            .and_then(|route| {
557                router::find_region_leader_ref(&route.region_routes, region.region_number())
558            })
559            .context(partition::error::FindLeaderSnafu {
560                region_id: *region,
561                table_id: region.table_id(),
562            })?;
563        let peer = shared_peers
564            .entry(peer)
565            .or_insert_with(|| Arc::new(peer.clone()));
566        peers.insert(*region, Arc::clone(peer));
567    }
568    Ok(peers)
569}
570
571fn get_peer(peers: &HashMap<RegionId, PeerRef>, region: RegionId) -> Result<PeerRef> {
572    peers.get(&region).cloned().context(UnexpectedSnafu {
573        err_msg: "Planned write has no resolved peer",
574    })
575}
576
577fn merge_attempts(
578    attempts: Vec<WriteAttempt>,
579    response: &mut RegionResponse,
580) -> Result<Vec<(RoutedWrite, Error)>> {
581    let mut failures = Vec::new();
582    let mut terminal_error = None;
583    for attempt in attempts {
584        match attempt {
585            WriteAttempt::Success(result) => merge_response(response, result),
586            WriteAttempt::Failed {
587                error,
588                request: Some(write),
589            } => failures.push((write, error)),
590            WriteAttempt::Failed {
591                error,
592                request: None,
593            } => {
594                terminal_error.get_or_insert(error);
595            }
596        }
597    }
598    // A terminal error takes precedence over retained retryable failures.
599    match terminal_error {
600        Some(error) => Err(error),
601        None => Ok(failures),
602    }
603}
604
605fn merge_response(response: &mut RegionResponse, result: RegionResponse) {
606    response.affected_rows += result.affected_rows;
607    response.extensions.extend(result.extensions);
608    response.metadata = result.metadata;
609}
610
611fn split_request(request: RegionRequest) -> Vec<RegionRequest> {
612    let RegionRequest { header, body } = request;
613    let bodies: Vec<Body> = match body {
614        Some(Body::Inserts(inserts)) => inserts
615            .requests
616            .into_iter()
617            .map(|request| {
618                Body::Inserts(InsertRequests {
619                    requests: vec![request],
620                })
621            })
622            .collect(),
623        Some(Body::Deletes(deletes)) => deletes
624            .requests
625            .into_iter()
626            .map(|request| {
627                Body::Deletes(DeleteRequests {
628                    requests: vec![request],
629                })
630            })
631            .collect(),
632        body => return vec![RegionRequest { header, body }],
633    };
634    bodies
635        .into_iter()
636        .map(|body| RegionRequest {
637            header: header.clone(),
638            body: Some(body),
639        })
640        .collect()
641}
642
643fn write_region_id(request: &RegionRequest) -> Result<RegionId> {
644    let region_id = match &request.body {
645        Some(Body::Inserts(inserts)) => match inserts.requests.as_slice() {
646            [request] => Some(request.region_id),
647            _ => None,
648        },
649        Some(Body::Deletes(deletes)) => match deletes.requests.as_slice() {
650            [request] => Some(request.region_id),
651            _ => None,
652        },
653        Some(Body::BulkInsert(bulk)) => Some(bulk.region_id),
654        _ => None,
655    };
656    region_id.map(RegionId::from_u64).ok_or_else(|| {
657        UnexpectedSnafu {
658            err_msg: "Expected a single-region write request",
659        }
660        .build()
661    })
662}
operator/write_dispatcher.rs

operator/
write_dispatcher.rs
