1mod 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
47pub struct WriteDispatcher {
49 node_manager: NodeManagerRef,
50 partition_manager: PartitionRuleManagerRef,
51 retry: Option<WriteRetry>,
52}
53
54pub type WriteDispatcherRef = Arc<WriteDispatcher>;
56
57pub(crate) enum WriteAttempt {
59 Success(RegionResponse),
60 Failed {
61 error: Error,
62 request: Option<RoutedWrite>,
65 },
66}
67
68impl WriteDispatcher {
69 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 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 pub async fn dispatch(&self, requests: Vec<RegionRequest>) -> Result<RegionResponse> {
101 self.dispatch_inner(requests, None).await
102 }
103
104 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 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(®ions)
133 .await
134 .context(FindRegionLeaderSnafu)?;
135 let mut writes = Vec::new();
136 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 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 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 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 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 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 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
353struct WriteRetry {
355 node_manager: NodeManagerRef,
356 partition_manager: PartitionRuleManagerRef,
357 options: WriteRetryOptions,
358 catalog_manager: CatalogManagerRef,
359 cache_invalidator: CacheInvalidatorRef,
360}
361
362pub(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(®ions).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 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 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, ®ions)
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(®ion.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(®ion).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 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}