From b90f5604b76ce57353419e8e87c8d3a92f5d1f6f Mon Sep 17 00:00:00 2001 From: WenyXu Date: Mon, 12 Oct 2026 13:01:14 +0000 Subject: [PATCH] style(operator): align dispatch imports and test utilities Signed-off-by: WenyXu --- src/operator/src/bulk_insert.rs | 3 +- src/operator/src/write_dispatcher.rs | 55 ++++++++------- src/operator/src/write_dispatcher/tests.rs | 80 ++++++++-------------- 3 files changed, 57 insertions(+), 81 deletions(-) diff --git a/src/operator/src/bulk_insert.rs b/src/operator/src/bulk_insert.rs index 779136c63ae..d5aa7168b44 100644 --- a/src/operator/src/bulk_insert.rs +++ b/src/operator/src/bulk_insert.rs @@ -14,6 +14,7 @@ use std::collections::HashSet; +use api::region::RegionResponse; use api::v1::flow::{DirtyWindowRequest, DirtyWindowRequests}; use api::v1::region::{ BulkInsertRequest, RegionRequest, RegionRequestHeader, bulk_insert_request, region_request, @@ -245,7 +246,7 @@ impl Inserter { } else { None }; - let handle: common_runtime::JoinHandle> = + let handle: common_runtime::JoinHandle> = common_runtime::spawn_global(async move { let (header, payload) = if mask.select_all() { // SAFETY: raw data must be present, we can avoid re-encoding. diff --git a/src/operator/src/write_dispatcher.rs b/src/operator/src/write_dispatcher.rs index 4acda2701b5..4af4fbfafdf 100644 --- a/src/operator/src/write_dispatcher.rs +++ b/src/operator/src/write_dispatcher.rs @@ -15,13 +15,13 @@ #[cfg(test)] mod tests; -use std::collections::{HashMap, HashSet}; +use std::collections::{HashMap, HashSet, hash_map::Entry}; use std::sync::Arc; use api::region::RegionResponse; use api::v1::region::{DeleteRequests, InsertRequests, RegionRequest, region_request::Body}; use common_error::ext::BoxedError; -use common_meta::error::{self, Result}; +use common_meta::error::{ExternalSnafu, Result, UnexpectedSnafu}; use common_meta::node_manager::NodeManagerRef; use common_meta::peer::Peer; use common_meta::rpc::router; @@ -31,7 +31,7 @@ use prometheus::Histogram; use snafu::{OptionExt, ResultExt}; use store_api::storage::RegionId; -use crate::error as operator_error; +use crate::error::{self, FindRegionLeaderSnafu, RequestRegionSnafu}; /// Dispatches already partitioned insert, delete, and bulk requests to their peers. pub struct WriteDispatcher { @@ -69,18 +69,18 @@ impl WriteDispatcher { pub(crate) async fn prepare( &self, requests: Vec, - ) -> operator_error::Result> { + ) -> error::Result> { let requests = requests.into_iter().map(split_request).collect::>(); let regions = requests .iter() .flatten() .map(write_region_id) .collect::>>() - .context(operator_error::RequestRegionSnafu)?; + .context(RequestRegionSnafu)?; let peers = self .resolve_peers(®ions) .await - .context(operator_error::FindRegionLeaderSnafu)?; + .context(FindRegionLeaderSnafu)?; let mut writes = Vec::new(); // Preserve each original envelope and peer batching when retries are disabled. for requests in requests { @@ -88,30 +88,28 @@ impl WriteDispatcher { for request in requests { let peer = get_peer( &peers, - write_region_id(&request).context(operator_error::RequestRegionSnafu)?, + write_region_id(&request).context(RequestRegionSnafu)?, ) - .context(operator_error::RequestRegionSnafu)?; + .context(RequestRegionSnafu)?; match by_peer.entry(peer) { - std::collections::hash_map::Entry::Vacant(entry) => { + Entry::Vacant(entry) => { entry.insert(request); } - std::collections::hash_map::Entry::Occupied(mut entry) => { - match (&mut entry.get_mut().body, request.body) { - (Some(Body::Inserts(existing)), Some(Body::Inserts(next))) => { - existing.requests.extend(next.requests) - } - (Some(Body::Deletes(existing)), Some(Body::Deletes(next))) => { - existing.requests.extend(next.requests) - } - _ => { - return error::UnexpectedSnafu { - err_msg: "Cannot combine different write kinds", - } - .fail() - .context(operator_error::RequestRegionSnafu); - } + Entry::Occupied(mut entry) => match (&mut entry.get_mut().body, request.body) { + (Some(Body::Inserts(existing)), Some(Body::Inserts(next))) => { + existing.requests.extend(next.requests) } - } + (Some(Body::Deletes(existing)), Some(Body::Deletes(next))) => { + existing.requests.extend(next.requests) + } + _ => { + return UnexpectedSnafu { + err_msg: "Cannot combine different write kinds", + } + .fail() + .context(RequestRegionSnafu); + } + }, } } writes.extend(by_peer); @@ -177,6 +175,7 @@ impl WriteDispatcher { } Ok(peers) } + async fn dispatch_inner( &self, requests: Vec, @@ -186,13 +185,13 @@ impl WriteDispatcher { .prepare(requests) .await .map_err(BoxedError::new) - .context(error::ExternalSnafu)?; + .context(ExternalSnafu)?; self.dispatch_prepared(writes, histogram).await } } fn get_peer(peers: &HashMap, region: RegionId) -> Result { - peers.get(®ion).cloned().context(error::UnexpectedSnafu { + peers.get(®ion).cloned().context(UnexpectedSnafu { err_msg: "Planned write has no resolved peer", }) } @@ -249,7 +248,7 @@ fn write_region_id(request: &RegionRequest) -> Result { _ => None, }; region_id.map(RegionId::from_u64).ok_or_else(|| { - error::UnexpectedSnafu { + UnexpectedSnafu { err_msg: "Expected a single-region write request", } .build() diff --git a/src/operator/src/write_dispatcher/tests.rs b/src/operator/src/write_dispatcher/tests.rs index 92b4de16bbc..8e5c59f8b68 100644 --- a/src/operator/src/write_dispatcher/tests.rs +++ b/src/operator/src/write_dispatcher/tests.rs @@ -31,16 +31,14 @@ use common_grpc::channel_manager::ChannelConfig; use common_meta::cache as meta_cache; use common_meta::ddl::test_util::datanode_handler::DatanodeWatcher; use common_meta::key::table_route::TableRouteValue; -use common_meta::node_manager::{ - Datanode, DatanodeManager, DatanodeRef, FlownodeManager, FlownodeRef, -}; -use common_meta::test_util::MockDatanodeManager; +use common_meta::test_util::{MockDatanodeHandler, MockDatanodeManager}; use common_query::request::QueryRequest; use common_recordbatch::SendableRecordBatchStream; use partition::manager::PartitionRuleManager; use tokio::sync::{Barrier, Notify}; use super::*; +use crate::error::Error; use crate::test_util; fn build_insert_request(regions: &[u32]) -> RegionRequest { @@ -64,7 +62,7 @@ async fn test_resolve_peers_batches_logical_and_physical_regions() { let partition_manager = test_util::create_partition_rule_manager(backend.clone()).await; common_meta::key::TableMetadataManager::new(backend) .create_table_metadata( - crate::test_util::new_test_table_info(42, "dispatcher_logical", [1, 2, 3].into_iter()), + test_util::new_test_table_info(42, "dispatcher_logical", [1, 2, 3].into_iter()), TableRouteValue::logical(1), Default::default(), ) @@ -194,22 +192,12 @@ async fn test_dispatch_preserves_write_payloads_and_envelopes() { .prepare(vec![build_insert_request(&[1, 999])]) .await .unwrap_err(); - assert!(matches!( - error, - operator_error::Error::FindRegionLeader { .. } - )); + assert!(matches!(error, Error::FindRegionLeader { .. })); assert!(received.try_recv().is_err()); } -struct GatedNodeManager { - entered: Arc, - rejected: Arc, - completed: Arc, - release: Arc, -} - -struct GatedDatanode { - reject: bool, +#[derive(Clone)] +struct GatedDatanodeHandler { entered: Arc, rejected: Arc, completed: Arc, @@ -217,32 +205,12 @@ struct GatedDatanode { } #[async_trait::async_trait] -impl DatanodeManager for GatedNodeManager { - async fn datanode(&self, peer: &Peer) -> DatanodeRef { - Arc::new(GatedDatanode { - reject: peer.id == 1, - entered: self.entered.clone(), - rejected: self.rejected.clone(), - completed: self.completed.clone(), - release: self.release.clone(), - }) - } -} - -#[async_trait::async_trait] -impl FlownodeManager for GatedNodeManager { - async fn flownode(&self, _peer: &Peer) -> FlownodeRef { - unimplemented!() - } -} - -#[async_trait::async_trait] -impl Datanode for GatedDatanode { - async fn handle(&self, _request: RegionRequest) -> Result { +impl MockDatanodeHandler for GatedDatanodeHandler { + async fn handle(&self, peer: &Peer, _request: RegionRequest) -> Result { self.entered.wait().await; - if self.reject { + if peer.id == 1 { self.rejected.store(true, Ordering::SeqCst); - return error::UnexpectedSnafu { + return UnexpectedSnafu { err_msg: "first region rejected", } .fail(); @@ -252,7 +220,11 @@ impl Datanode for GatedDatanode { Ok(RegionResponse::new(1)) } - async fn handle_query(&self, _request: QueryRequest) -> Result { + async fn handle_query( + &self, + _peer: &Peer, + _request: QueryRequest, + ) -> Result { unimplemented!() } } @@ -261,13 +233,16 @@ impl Datanode for GatedDatanode { async fn test_dispatch_sends_concurrently_and_waits_for_siblings_after_error() { let backend = test_util::prepare_mocked_backend().await; let partition_manager = test_util::create_partition_rule_manager(backend).await; - let nodes = Arc::new(GatedNodeManager { + let handler = GatedDatanodeHandler { entered: Arc::new(Barrier::new(2)), rejected: Arc::default(), completed: Arc::default(), release: Arc::default(), - }); - let dispatcher = WriteDispatcher::new(nodes.clone(), partition_manager); + }; + let dispatcher = WriteDispatcher::new( + Arc::new(MockDatanodeManager::new(handler.clone())), + partition_manager, + ); let dispatch = dispatcher.dispatch(vec![build_insert_request(&[1, 2])]); tokio::pin!(dispatch); tokio::time::timeout( @@ -277,7 +252,7 @@ async fn test_dispatch_sends_concurrently_and_waits_for_siblings_after_error() { dispatch.as_mut().poll(cx).is_pending(), "dispatch returned while a sibling remained blocked" ); - if nodes.rejected.load(Ordering::SeqCst) { + if handler.rejected.load(Ordering::SeqCst) { Poll::Ready(()) } else { Poll::Pending @@ -286,23 +261,24 @@ async fn test_dispatch_sends_concurrently_and_waits_for_siblings_after_error() { ) .await .expect("both region sends must reach the barrier concurrently"); - assert!(!nodes.completed.load(Ordering::SeqCst)); - nodes.release.notify_one(); + assert!(!handler.completed.load(Ordering::SeqCst)); + handler.release.notify_one(); let error = tokio::time::timeout(Duration::from_secs(1), dispatch) .await .unwrap() .unwrap_err(); assert!(error.to_string().contains("first region rejected")); - assert!(nodes.completed.load(Ordering::SeqCst)); + assert!(handler.completed.load(Ordering::SeqCst)); } #[tokio::test] async fn test_resolve_peers_cache_lookup_ablation() { + common_telemetry::init_default_ut_logging(); let backend = test_util::prepare_mocked_backend().await; test_util::create_partition_rule_manager(backend.clone()).await; common_meta::key::TableMetadataManager::new(backend.clone()) .create_table_metadata( - crate::test_util::new_test_table_info(42, "resolution_ablation", [1, 2, 3].into_iter()), + test_util::new_test_table_info(42, "resolution_ablation", [1, 2, 3].into_iter()), TableRouteValue::logical(1), Default::default(), ) @@ -363,7 +339,7 @@ async fn test_resolve_peers_cache_lookup_ablation() { _ => unreachable!(), }; assert_eq!(gets, expected_gets); - eprintln!( + common_telemetry::info!( "route_resolution_ablation run={run} mode={mode} regions={} cache_gets={gets} elapsed_ns={}", regions.len(), elapsed.as_nanos()