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        mut request: RegionRequest,
207        histogram: Option<&Histogram>,
208    ) -> WriteAttempt {
209        if let Some(retry) = &self.retry {
210            if matches!(&request.body, Some(Body::Inserts(_) | Body::Deletes(_))) {
211                api::region::set_accept_partial_result(request.header.get_or_insert_default());
212            }
213            return retry
214                .send(
215                    RoutedWrite {
216                        peer,
217                        request: Arc::new(request),
218                    },
219                    histogram,
220                )
221                .await;
222        }
223        let _timer = histogram.map(Histogram::start_timer);
224        let datanode = self.node_manager.datanode(&peer).await;
225        match datanode.handle(request).await {
226            Ok(response) => WriteAttempt::Success(response),
227            Err(error) => WriteAttempt::Failed {
228                error,
229                request: None,
230            },
231        }
232    }
233
234    /// Collects first attempts and retries their failures together, refreshing routes once per round.
235    pub(crate) async fn finish(
236        &self,
237        attempts: Vec<WriteAttempt>,
238        histogram: Option<&Histogram>,
239        cancellation: Option<&CancellationToken>,
240    ) -> Result<RegionResponse> {
241        if let Some(retry) = &self.retry {
242            if retry.should_retry(&attempts, 0) {
243                return retry.retry(attempts, histogram, cancellation).await;
244            }
245            let mut response = RegionResponse::new(0);
246            let failures = merge_attempts(attempts, &mut response)?;
247            return match failures.into_iter().next() {
248                Some((_, error)) => Err(error),
249                None => Ok(response),
250            };
251        }
252        let mut response = RegionResponse::new(0);
253        for attempt in attempts {
254            match attempt {
255                WriteAttempt::Success(result) => merge_response(&mut response, result),
256                WriteAttempt::Failed { error, .. } => return Err(error),
257            }
258        }
259        Ok(response)
260    }
261
262    /// Resolves only the requested peers, loading each table route once.
263    pub(crate) async fn resolve_peers(
264        &self,
265        regions: &[RegionId],
266    ) -> partition::error::Result<HashMap<RegionId, PeerRef>> {
267        resolve_peers(&self.partition_manager, regions).await
268    }
269
270    async fn dispatch_inner(
271        &self,
272        requests: Vec<RegionRequest>,
273        histogram: Option<Histogram>,
274    ) -> Result<RegionResponse> {
275        let writes = self
276            .prepare(requests)
277            .await
278            .map_err(BoxedError::new)
279            .context(ExternalSnafu)?;
280        self.execute(writes, histogram, None).await
281    }
282}
283
284/// Sends shared requests and coordinates retries with refreshed routes.
285struct WriteRetry {
286    node_manager: NodeManagerRef,
287    partition_manager: PartitionRuleManagerRef,
288    options: WriteRetryOptions,
289    catalog_manager: CatalogManagerRef,
290    cache_invalidator: CacheInvalidatorRef,
291}
292
293/// A single-region write request paired with its destination peer.
294pub(crate) struct RoutedWrite {
295    peer: PeerRef,
296    request: Arc<RegionRequest>,
297}
298
299impl WriteRetry {
300    async fn send(&self, write: RoutedWrite, histogram: Option<&Histogram>) -> WriteAttempt {
301        let _timer = histogram.map(Histogram::start_timer);
302        let requester = self.node_manager.datanode(&write.peer).await;
303        match requester.handle(Arc::clone(&write.request)).await {
304            Ok(response) => WriteAttempt::Success(response),
305            Err(error) => {
306                let request = error.retry_hint().is_retryable().then_some(write);
307                WriteAttempt::Failed { error, request }
308            }
309        }
310    }
311
312    fn should_retry(&self, attempts: &[WriteAttempt], retries: usize) -> bool {
313        if retries >= self.options.max_retries {
314            return false;
315        }
316        let mut retryable = false;
317        for attempt in attempts {
318            match attempt {
319                WriteAttempt::Success(_) => {}
320                WriteAttempt::Failed {
321                    request: Some(_), ..
322                } => retryable = true,
323                WriteAttempt::Failed { request: None, .. } => return false,
324            }
325        }
326        retryable
327    }
328
329    async fn retry(
330        &self,
331        mut attempts: Vec<WriteAttempt>,
332        histogram: Option<&Histogram>,
333        cancellation: Option<&CancellationToken>,
334    ) -> Result<RegionResponse> {
335        let mut response = RegionResponse::new(0);
336        let mut backoff = ExponentialBuilder::default()
337            .with_min_delay(self.options.initial_backoff)
338            .with_max_delay(self.options.max_backoff)
339            .with_factor(self.options.backoff_multiplier)
340            .with_max_times(self.options.max_retries)
341            .with_jitter()
342            .build();
343        let mut retries = 0;
344        while self.should_retry(&attempts, retries) {
345            let Some(delay) = backoff.next() else {
346                break;
347            };
348            if let Some(cancellation) = cancellation {
349                tokio::select! {
350                    biased;
351                    _ = cancellation.cancelled() => break,
352                    _ = tokio::time::sleep(delay) => {}
353                }
354            } else {
355                tokio::time::sleep(delay).await;
356            }
357            if cancellation.is_some_and(CancellationToken::is_cancelled) {
358                break;
359            }
360            let failures = merge_attempts(attempts, &mut response)?;
361            let mut first_error = None;
362            let mut refresh = Vec::new();
363            let mut pending = Vec::new();
364            for (write, error) in failures {
365                if error.retry_hint() == RetryHint::RetryAfterRouteRefresh {
366                    refresh.push(write);
367                } else {
368                    pending.push(write);
369                }
370                first_error.get_or_insert(error);
371            }
372            if !refresh.is_empty() {
373                pending.extend(self.reroute_failed_writes(refresh).await?);
374            }
375            if let Some(error) =
376                first_error.filter(|_| cancellation.is_some_and(CancellationToken::is_cancelled))
377            {
378                return Err(error);
379            }
380            attempts =
381                future::join_all(pending.into_iter().map(|write| self.send(write, histogram)))
382                    .await;
383            retries += 1;
384        }
385        match merge_attempts(attempts, &mut response)?.into_iter().next() {
386            Some((_, error)) => Err(error),
387            None => Ok(response),
388        }
389    }
390
391    async fn reroute_failed_writes(&self, writes: Vec<RoutedWrite>) -> Result<Vec<RoutedWrite>> {
392        let regions = writes
393            .iter()
394            .map(|write| write_region_id(&write.request))
395            .collect::<Result<Vec<_>>>()?;
396        self.invalidate_routes(&regions).await?;
397        let mut table_names = HashMap::new();
398        for table_id in regions
399            .iter()
400            .map(|region| region.table_id())
401            .collect::<HashSet<_>>()
402        {
403            // Raw metadata is only used to recover the name of the failed request's table.
404            let table_info = self
405                .catalog_manager
406                .table_info_by_id(table_id)
407                .await
408                .map_err(BoxedError::new)
409                .context(ExternalSnafu)?
410                .with_context(|| TableInfoNotFoundSnafu {
411                    table: table_id.to_string(),
412                })?;
413            table_names.insert(
414                table_id,
415                TableName::new(
416                    &table_info.catalog_name,
417                    &table_info.schema_name,
418                    &table_info.name,
419                ),
420            );
421        }
422        // TableId invalidation refreshes metadata and routes, but TableCache is keyed by name.
423        // Use the refreshed names in case a table was renamed while the write was pending.
424        let caches = table_names
425            .values()
426            .cloned()
427            .map(CacheIdent::TableName)
428            .collect::<Vec<_>>();
429        self.cache_invalidator
430            .invalidate(&Context::default(), &caches)
431            .await?;
432        let mut partitioners = HashMap::new();
433        for (table_id, name) in table_names {
434            let table = self
435                .catalog_manager
436                .table(
437                    &name.catalog_name,
438                    &name.schema_name,
439                    &name.table_name,
440                    None,
441                )
442                .await
443                .map_err(BoxedError::new)
444                .context(ExternalSnafu)?
445                .with_context(|| TableInfoNotFoundSnafu {
446                    table: name.to_string(),
447                })?;
448            let table_info = table.table_info();
449            ensure!(
450                table_info.table_id() == table_id,
451                UnexpectedSnafu {
452                    err_msg: format!(
453                        "Table {name} changed ID from {table_id} to {} while retrying writes",
454                        table_info.table_id()
455                    ),
456                }
457            );
458            let partitioner = TablePartitioner::try_new(&self.partition_manager, &table_info)
459                .await
460                .map_err(BoxedError::new)
461                .context(ExternalSnafu)?;
462            partitioners.insert(table_id, (table_info, partitioner));
463        }
464        let mut rebuilt_requests = Vec::new();
465        for write in writes {
466            let request = Arc::try_unwrap(write.request).map_err(|_| {
467                UnexpectedSnafu {
468                    err_msg: "Write request is still shared after its RPC completed",
469                }
470                .build()
471            })?;
472            let table_id = write_region_id(&request)?.table_id();
473            let (table_info, partitioner) =
474                partitioners.get(&table_id).context(UnexpectedSnafu {
475                    err_msg: "Refreshed write has no partition snapshot",
476                })?;
477            let requests = partitioner
478                .rebuild_write_requests(request, table_info)
479                .map_err(BoxedError::new)
480                .context(ExternalSnafu)?;
481            rebuilt_requests.extend(requests);
482        }
483        self.prepare(rebuilt_requests).await
484    }
485
486    async fn prepare(&self, requests: Vec<RegionRequest>) -> Result<Vec<RoutedWrite>> {
487        let regions = requests
488            .iter()
489            .map(write_region_id)
490            .collect::<Result<Vec<_>>>()?;
491        let peers = resolve_peers(&self.partition_manager, &regions)
492            .await
493            .map_err(BoxedError::new)
494            .context(ExternalSnafu)?;
495        requests
496            .into_iter()
497            .zip(regions)
498            .map(|(request, region)| {
499                Ok(RoutedWrite {
500                    peer: get_peer(&peers, region)?,
501                    request: Arc::new(request),
502                })
503            })
504            .collect()
505    }
506
507    async fn invalidate_routes(&self, regions: &[RegionId]) -> Result<()> {
508        let mut table_ids = regions
509            .iter()
510            .map(|region| region.table_id())
511            .collect::<HashSet<_>>();
512        let mut physical_ids = HashSet::new();
513        for table_id in &table_ids {
514            let (physical_id, _) = self
515                .partition_manager
516                .find_physical_table_route_with_id(*table_id)
517                .await
518                .map_err(BoxedError::new)
519                .context(ExternalSnafu)?;
520            physical_ids.insert(physical_id);
521        }
522        table_ids.extend(physical_ids);
523        let caches = table_ids
524            .into_iter()
525            .map(CacheIdent::TableId)
526            .collect::<Vec<_>>();
527        self.cache_invalidator
528            .invalidate(&Context::default(), &caches)
529            .await
530    }
531}
532
533async fn resolve_peers(
534    partition_manager: &PartitionRuleManager,
535    regions: &[RegionId],
536) -> partition::error::Result<HashMap<RegionId, PeerRef>> {
537    let table_ids = regions
538        .iter()
539        .map(|region| region.table_id())
540        .collect::<HashSet<_>>();
541    let mut physical_tables = HashMap::new();
542    let mut routes = HashMap::new();
543    for table_id in table_ids {
544        let (physical_id, route) = partition_manager
545            .find_physical_table_route_with_id(table_id)
546            .await?;
547        physical_tables.insert(table_id, physical_id);
548        routes.entry(physical_id).or_insert(route);
549    }
550    let mut shared_peers: HashMap<&Peer, PeerRef> = HashMap::new();
551    let mut peers = HashMap::with_capacity(regions.len());
552    for region in regions {
553        if peers.contains_key(region) {
554            continue;
555        }
556        let peer = physical_tables
557            .get(&region.table_id())
558            .and_then(|physical_id| routes.get(physical_id))
559            .and_then(|route| {
560                router::find_region_leader_ref(&route.region_routes, region.region_number())
561            })
562            .context(partition::error::FindLeaderSnafu {
563                region_id: *region,
564                table_id: region.table_id(),
565            })?;
566        let peer = shared_peers
567            .entry(peer)
568            .or_insert_with(|| Arc::new(peer.clone()));
569        peers.insert(*region, Arc::clone(peer));
570    }
571    Ok(peers)
572}
573
574fn get_peer(peers: &HashMap<RegionId, PeerRef>, region: RegionId) -> Result<PeerRef> {
575    peers.get(&region).cloned().context(UnexpectedSnafu {
576        err_msg: "Planned write has no resolved peer",
577    })
578}
579
580fn merge_attempts(
581    attempts: Vec<WriteAttempt>,
582    response: &mut RegionResponse,
583) -> Result<Vec<(RoutedWrite, Error)>> {
584    let mut failures = Vec::new();
585    let mut terminal_error = None;
586    for attempt in attempts {
587        match attempt {
588            WriteAttempt::Success(result) => merge_response(response, result),
589            WriteAttempt::Failed {
590                error,
591                request: Some(write),
592            } => failures.push((write, error)),
593            WriteAttempt::Failed {
594                error,
595                request: None,
596            } => {
597                terminal_error.get_or_insert(error);
598            }
599        }
600    }
601    // A terminal error takes precedence over retained retryable failures.
602    match terminal_error {
603        Some(error) => Err(error),
604        None => Ok(failures),
605    }
606}
607
608fn merge_response(response: &mut RegionResponse, result: RegionResponse) {
609    response.affected_rows += result.affected_rows;
610    response.extensions.extend(result.extensions);
611    response.metadata = result.metadata;
612}
613
614fn split_request(request: RegionRequest) -> Vec<RegionRequest> {
615    let RegionRequest { header, body } = request;
616    let bodies: Vec<Body> = match body {
617        Some(Body::Inserts(inserts)) => inserts
618            .requests
619            .into_iter()
620            .map(|request| {
621                Body::Inserts(InsertRequests {
622                    requests: vec![request],
623                })
624            })
625            .collect(),
626        Some(Body::Deletes(deletes)) => deletes
627            .requests
628            .into_iter()
629            .map(|request| {
630                Body::Deletes(DeleteRequests {
631                    requests: vec![request],
632                })
633            })
634            .collect(),
635        body => return vec![RegionRequest { header, body }],
636    };
637    bodies
638        .into_iter()
639        .map(|body| RegionRequest {
640            header: header.clone(),
641            body: Some(body),
642        })
643        .collect()
644}
645
646fn write_region_id(request: &RegionRequest) -> Result<RegionId> {
647    let region_id = match &request.body {
648        Some(Body::Inserts(inserts)) => match inserts.requests.as_slice() {
649            [request] => Some(request.region_id),
650            _ => None,
651        },
652        Some(Body::Deletes(deletes)) => match deletes.requests.as_slice() {
653            [request] => Some(request.region_id),
654            _ => None,
655        },
656        Some(Body::BulkInsert(bulk)) => Some(bulk.region_id),
657        _ => None,
658    };
659    region_id.map(RegionId::from_u64).ok_or_else(|| {
660        UnexpectedSnafu {
661            err_msg: "Expected a single-region write request",
662        }
663        .build()
664    })
665}
operator/write_dispatcher.rs

operator/
write_dispatcher.rs
