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

operator/
write_dispatcher.rs
