Skip to main content

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