diff --git a/src/common/grpc/src/channel_manager.rs b/src/common/grpc/src/channel_manager.rs index 1be0e22d1a..01da70558a 100644 --- a/src/common/grpc/src/channel_manager.rs +++ b/src/common/grpc/src/channel_manager.rs @@ -383,6 +383,21 @@ pub struct ChannelConfig { pub accept_compression: bool, } +/// Configures a generated tonic client with compression and message size limits. +#[macro_export] +macro_rules! configure_tonic_client { + ($client:expr, $channel_manager:expr $(,)?) => {{ + let channel_manager = &$channel_manager; + let config = channel_manager.config(); + $client + .accept_compressed(::tonic::codec::CompressionEncoding::Gzip) + .accept_compressed(::tonic::codec::CompressionEncoding::Zstd) + .send_compressed(::tonic::codec::CompressionEncoding::Zstd) + .max_decoding_message_size(config.max_recv_message_size.as_bytes() as usize) + .max_encoding_message_size(config.max_send_message_size.as_bytes() as usize) + }}; +} + impl Default for ChannelConfig { fn default() -> Self { Self { @@ -710,6 +725,75 @@ mod tests { ); } + #[derive(Default)] + struct FakeTonicClient { + accepted_gzip: bool, + accepted_zstd: bool, + sent_gzip: bool, + sent_zstd: bool, + max_decoding_message_size: usize, + max_encoding_message_size: usize, + } + + impl FakeTonicClient { + fn accept_compressed(mut self, encoding: ::tonic::codec::CompressionEncoding) -> Self { + match encoding { + ::tonic::codec::CompressionEncoding::Gzip => self.accepted_gzip = true, + ::tonic::codec::CompressionEncoding::Zstd => self.accepted_zstd = true, + _ => unreachable!(), + } + self + } + + fn send_compressed(mut self, encoding: ::tonic::codec::CompressionEncoding) -> Self { + match encoding { + ::tonic::codec::CompressionEncoding::Gzip => self.sent_gzip = true, + ::tonic::codec::CompressionEncoding::Zstd => self.sent_zstd = true, + _ => unreachable!(), + } + self + } + + fn max_decoding_message_size(mut self, size: usize) -> Self { + self.max_decoding_message_size = size; + self + } + + fn max_encoding_message_size(mut self, size: usize) -> Self { + self.max_encoding_message_size = size; + self + } + } + + #[test] + fn test_configure_tonic_client() { + let recv_message_size = ReadableSize::mb(2); + let send_message_size = ReadableSize::mb(3); + let channel_manager = ChannelManager::with_config( + ChannelConfig { + max_recv_message_size: recv_message_size, + max_send_message_size: send_message_size, + ..ChannelConfig::new() + }, + None, + ); + + let client = crate::configure_tonic_client!(FakeTonicClient::default(), channel_manager,); + + assert!(client.accepted_gzip); + assert!(client.accepted_zstd); + assert!(!client.sent_gzip); + assert!(client.sent_zstd); + assert_eq!( + recv_message_size.as_bytes() as usize, + client.max_decoding_message_size + ); + assert_eq!( + send_message_size.as_bytes() as usize, + client.max_encoding_message_size + ); + } + #[test] fn test_build_endpoint() { let config = ChannelConfig::new() diff --git a/src/meta-client/src/client.rs b/src/meta-client/src/client.rs index de41caf19e..4fc6020fff 100644 --- a/src/meta-client/src/client.rs +++ b/src/meta-client/src/client.rs @@ -891,9 +891,11 @@ impl MetaClient { #[cfg(test)] mod tests { use std::sync::atomic::{AtomicUsize, Ordering}; + use std::sync::{Arc, Mutex}; use api::v1::meta::{HeartbeatRequest, Peer}; - use common_meta::kv_backend::{KvBackendRef, ResettableKvBackendRef}; + use common_base::readable_size::ReadableSize; + use common_meta::kv_backend::{KvBackend, KvBackendRef, ResettableKvBackendRef, TxnService}; use rand::Rng; use super::*; @@ -919,6 +921,27 @@ mod tests { } } + async fn new_with_grpc_message_sizes( + ns: impl Into, + server_max_recv_message_size: ReadableSize, + server_max_send_message_size: ReadableSize, + client_max_recv_message_size: ReadableSize, + client_max_send_message_size: ReadableSize, + ) -> Self { + let (client, meta_ctx) = mocks::mock_client_with_memstore_and_grpc_message_sizes( + server_max_recv_message_size, + server_max_send_message_size, + client_max_recv_message_size, + client_max_send_message_size, + ) + .await; + Self { + ns: ns.into(), + client, + meta_ctx, + } + } + fn key(&self, name: &str) -> Vec { format!("{}-{}-{}", TEST_KEY_PREFIX, self.ns, name).into_bytes() } @@ -1331,20 +1354,77 @@ mod tests { } } - fn mock_decoder(_kv: KeyValue) -> MetaResult<()> { - Ok(()) + fn mock_decoder(kv: KeyValue) -> MetaResult> { + Ok(kv.value) } - #[tokio::test] - async fn test_cluster_client_adaptive_range() { - let tx = new_client("test_cluster_client").await; + struct RecordingKvBackend { + inner: KvBackendRef, + limits: Arc>>, + } + + #[async_trait::async_trait] + impl TxnService for RecordingKvBackend { + type Error = meta_error::Error; + } + + #[async_trait::async_trait] + impl KvBackend for RecordingKvBackend { + fn name(&self) -> &str { + "RecordingKvBackend" + } + + fn as_any(&self) -> &dyn std::any::Any { + self + } + + async fn range(&self, req: RangeRequest) -> MetaResult { + self.limits.lock().unwrap().push(req.limit); + self.inner.range(req).await + } + + async fn put(&self, req: PutRequest) -> MetaResult { + self.inner.put(req).await + } + + async fn batch_put(&self, req: BatchPutRequest) -> MetaResult { + self.inner.batch_put(req).await + } + + async fn batch_get(&self, req: BatchGetRequest) -> MetaResult { + self.inner.batch_get(req).await + } + + async fn delete_range(&self, req: DeleteRangeRequest) -> MetaResult { + self.inner.delete_range(req).await + } + + async fn batch_delete(&self, req: BatchDeleteRequest) -> MetaResult { + self.inner.batch_delete(req).await + } + } + + async fn adaptive_range_with_message_sizes( + server_max_recv_message_size: ReadableSize, + server_max_send_message_size: ReadableSize, + client_max_recv_message_size: ReadableSize, + client_max_send_message_size: ReadableSize, + ) -> (Vec>, Vec>, Vec) { + let tx = TestClient::new_with_grpc_message_sizes( + "test_cluster_client", + server_max_recv_message_size, + server_max_send_message_size, + client_max_recv_message_size, + client_max_send_message_size, + ) + .await; let in_memory = tx.in_memory().unwrap(); let cluster_client = tx.client.cluster_client().unwrap(); let mut rng = rand::rng(); - // Generates rough 10MB data, which is larger than the default grpc message size limit. - for i in 0..10 { - let data: Vec = (0..1024 * 1024).map(|_| rng.random::()).collect(); + let mut expected = Vec::new(); + for i in 0..4 { + let data: Vec = (0..256 * 1024).map(|_| rng.random::()).collect(); in_memory .put( PutRequest::new() @@ -1353,13 +1433,48 @@ mod tests { ) .await .unwrap(); + expected.push(data); } let req = RangeRequest::new().with_prefix(b"__prefix/"); + let limits = Arc::new(Mutex::new(Vec::new())); + let recording_backend = RecordingKvBackend { + inner: Arc::new(cluster_client), + limits: limits.clone(), + }; let stream = - PaginationStream::new(Arc::new(cluster_client), req, 10, mock_decoder).into_stream(); + PaginationStream::new(Arc::new(recording_backend), req, 4, mock_decoder).into_stream(); let res = stream.try_collect::>().await.unwrap(); - assert_eq!(10, res.len()); + let limits = limits.lock().unwrap().clone(); + (expected, res, limits) + } + + #[tokio::test] + async fn test_cluster_client_adaptive_range() { + let (expected, res, limits) = adaptive_range_with_message_sizes( + ReadableSize::mb(2), + ReadableSize::mb(16), + ReadableSize::mb(1), + ReadableSize::mb(2), + ) + .await; + + assert_eq!(expected, res); + assert_eq!(vec![4, 2, 2], limits); + } + + #[tokio::test] + async fn test_cluster_client_adaptive_range_server_limit() { + let (expected, res, limits) = adaptive_range_with_message_sizes( + ReadableSize::mb(2), + ReadableSize::mb(1), + ReadableSize::mb(16), + ReadableSize::mb(2), + ) + .await; + + assert_eq!(expected, res); + assert_eq!(vec![4, 2, 2], limits); } } diff --git a/src/meta-client/src/client/ask_leader.rs b/src/meta-client/src/client/ask_leader.rs index 9a41559e14..ff665ef9d4 100644 --- a/src/meta-client/src/client/ask_leader.rs +++ b/src/meta-client/src/client/ask_leader.rs @@ -186,10 +186,13 @@ impl AskLeader { } fn create_asker(&self, addr: impl AsRef) -> Result> { - Ok(HeartbeatClient::new( - self.channel_manager - .get(addr) - .context(error::CreateChannelSnafu)?, + Ok(common_grpc::configure_tonic_client!( + HeartbeatClient::new( + self.channel_manager + .get(addr) + .context(error::CreateChannelSnafu)?, + ), + self.channel_manager, )) } diff --git a/src/meta-client/src/client/cluster.rs b/src/meta-client/src/client/cluster.rs index 976626b3a0..56d745284d 100644 --- a/src/meta-client/src/client/cluster.rs +++ b/src/meta-client/src/client/cluster.rs @@ -34,7 +34,6 @@ use common_telemetry::{error, info, warn}; use snafu::{ResultExt, ensure}; use tokio::sync::RwLock; use tonic::Status; -use tonic::codec::CompressionEncoding; use tonic::transport::Channel; use crate::client::{LeaderProviderRef, util}; @@ -155,10 +154,10 @@ impl Inner { fn make_client(&self, addr: impl AsRef) -> Result> { let channel = self.channel_manager.get(addr).context(CreateChannelSnafu)?; - Ok(ClusterClient::new(channel) - .accept_compressed(CompressionEncoding::Gzip) - .accept_compressed(CompressionEncoding::Zstd) - .send_compressed(CompressionEncoding::Zstd)) + Ok(common_grpc::configure_tonic_client!( + ClusterClient::new(channel), + self.channel_manager, + )) } #[inline] diff --git a/src/meta-client/src/client/config.rs b/src/meta-client/src/client/config.rs index a2bcb00904..bd81066681 100644 --- a/src/meta-client/src/client/config.rs +++ b/src/meta-client/src/client/config.rs @@ -21,7 +21,6 @@ use common_meta::util; use common_telemetry::tracing_context::TracingContext; use snafu::{OptionExt, ResultExt, ensure}; use tokio::sync::RwLock; -use tonic::codec::CompressionEncoding; use tonic::transport::Channel; use crate::client::{Id, LeaderProviderRef}; @@ -132,10 +131,10 @@ impl Inner { .get(addr) .context(error::CreateChannelSnafu)?; - Ok(ConfigClient::new(channel) - .accept_compressed(CompressionEncoding::Zstd) - .accept_compressed(CompressionEncoding::Gzip) - .send_compressed(CompressionEncoding::Zstd)) + Ok(common_grpc::configure_tonic_client!( + ConfigClient::new(channel), + self.channel_manager, + )) } #[inline] diff --git a/src/meta-client/src/client/heartbeat.rs b/src/meta-client/src/client/heartbeat.rs index add7bd54de..96fc398556 100644 --- a/src/meta-client/src/client/heartbeat.rs +++ b/src/meta-client/src/client/heartbeat.rs @@ -27,7 +27,6 @@ use snafu::{OptionExt, ResultExt, ensure}; use tokio::sync::{RwLock, mpsc}; use tokio_stream::wrappers::ReceiverStream; use tonic::Streaming; -use tonic::codec::CompressionEncoding; use tonic::transport::Channel; use crate::client::{Id, LeaderProviderRef}; @@ -282,16 +281,10 @@ impl Inner { .get(addr) .context(error::CreateChannelSnafu)?; - let config = self.channel_manager.config(); - let max_decoding_message_size = config.max_recv_message_size.as_bytes() as usize; - let max_encoding_message_size = config.max_send_message_size.as_bytes() as usize; - - Ok(HeartbeatClient::new(channel) - .accept_compressed(CompressionEncoding::Zstd) - .accept_compressed(CompressionEncoding::Gzip) - .send_compressed(CompressionEncoding::Zstd) - .max_decoding_message_size(max_decoding_message_size) - .max_encoding_message_size(max_encoding_message_size)) + Ok(common_grpc::configure_tonic_client!( + HeartbeatClient::new(channel), + self.channel_manager, + )) } #[inline] diff --git a/src/meta-client/src/client/procedure.rs b/src/meta-client/src/client/procedure.rs index 5ac6710c28..d7d5ccfaff 100644 --- a/src/meta-client/src/client/procedure.rs +++ b/src/meta-client/src/client/procedure.rs @@ -36,7 +36,6 @@ use common_telemetry::tracing_context::TracingContext; use common_telemetry::{error, info, warn}; use snafu::{ResultExt, ensure}; use tokio::sync::RwLock; -use tonic::codec::CompressionEncoding; use tonic::transport::Channel; use tonic::{Request, Status}; @@ -155,10 +154,10 @@ impl Inner { .get(addr) .context(error::CreateChannelSnafu)?; - Ok(ProcedureServiceClient::new(channel) - .accept_compressed(CompressionEncoding::Gzip) - .accept_compressed(CompressionEncoding::Zstd) - .send_compressed(CompressionEncoding::Zstd)) + Ok(common_grpc::configure_tonic_client!( + ProcedureServiceClient::new(channel), + self.channel_manager, + )) } #[inline] @@ -672,10 +671,14 @@ mod tests { let server = tonic::transport::Server::builder() .add_service( - HeartbeatServer::new(heartbeat).accept_compressed(CompressionEncoding::Zstd), + HeartbeatServer::new(heartbeat) + .accept_compressed(CompressionEncoding::Gzip) + .accept_compressed(CompressionEncoding::Zstd), ) .add_service( - ProcedureServiceServer::new(procedure).accept_compressed(CompressionEncoding::Zstd), + ProcedureServiceServer::new(procedure) + .accept_compressed(CompressionEncoding::Gzip) + .accept_compressed(CompressionEncoding::Zstd), ) .serve_with_incoming(TcpListenerStream::new(listener)); let server_handle = tokio::spawn(server); diff --git a/src/meta-client/src/client/store.rs b/src/meta-client/src/client/store.rs index 2aee8d8323..3f9bb98862 100644 --- a/src/meta-client/src/client/store.rs +++ b/src/meta-client/src/client/store.rs @@ -25,7 +25,6 @@ use common_grpc::channel_manager::ChannelManager; use common_telemetry::tracing_context::TracingContext; use snafu::{OptionExt, ResultExt, ensure}; use tokio::sync::RwLock; -use tonic::codec::CompressionEncoding; use tonic::transport::Channel; use crate::client::{Id, load_balance as lb}; @@ -275,17 +274,10 @@ impl Inner { .get(addr) .context(error::CreateChannelSnafu)?; - let max_decoding_message_size = self - .channel_manager - .config() - .max_recv_message_size - .as_bytes() as usize; - - Ok(StoreClient::new(channel) - .accept_compressed(CompressionEncoding::Gzip) - .accept_compressed(CompressionEncoding::Zstd) - .send_compressed(CompressionEncoding::Zstd) - .max_decoding_message_size(max_decoding_message_size)) + Ok(common_grpc::configure_tonic_client!( + StoreClient::new(channel), + self.channel_manager, + )) } #[inline] diff --git a/src/meta-client/src/error.rs b/src/meta-client/src/error.rs index d227896e34..5121bba847 100644 --- a/src/meta-client/src/error.rs +++ b/src/meta-client/src/error.rs @@ -176,7 +176,7 @@ impl Error { matches!( self, Error::MetaServer { - tonic_code: tonic::Code::OutOfRange, + tonic_code: tonic::Code::OutOfRange | tonic::Code::ResourceExhausted, .. } ) @@ -242,4 +242,28 @@ mod tests { assert_eq!(err.retry_hint(), RetryHint::Retryable); } + + #[test] + fn test_is_exceeded_size_limit_for_out_of_range() { + let err = Error::from(tonic::Status::new(tonic::Code::OutOfRange, "any message")); + + assert!(err.is_exceeded_size_limit()); + } + + #[test] + fn test_is_exceeded_size_limit_for_resource_exhausted() { + let err = Error::from(tonic::Status::new( + tonic::Code::ResourceExhausted, + "arbitrary message", + )); + + assert!(err.is_exceeded_size_limit()); + } + + #[test] + fn test_is_exceeded_size_limit_for_non_size_code() { + let err = Error::from(tonic::Status::new(tonic::Code::Internal, "message")); + + assert!(!err.is_exceeded_size_limit()); + } } diff --git a/src/meta-client/src/mocks.rs b/src/meta-client/src/mocks.rs index 2c80091c58..6e1daa2efc 100644 --- a/src/meta-client/src/mocks.rs +++ b/src/meta-client/src/mocks.rs @@ -12,8 +12,13 @@ // See the License for the specific language governing permissions and // limitations under the License. -use common_grpc::channel_manager::ChannelManager; +use std::sync::Arc; + +use common_base::readable_size::ReadableSize; +use common_grpc::channel_manager::{ChannelConfig, ChannelManager}; +use common_meta::kv_backend::memory::MemoryKvBackend; use common_meta::kv_backend::{KvBackendRef, ResettableKvBackendRef}; +use meta_srv::metasrv::MetasrvOptions; use meta_srv::mocks as server_mock; use meta_srv::mocks::MockInfo; @@ -41,6 +46,50 @@ pub async fn mock_client_with_memstore() -> (MetaClient, MockMetaContext) { ) } +pub async fn mock_client_with_memstore_and_grpc_message_sizes( + server_max_recv_message_size: ReadableSize, + server_max_send_message_size: ReadableSize, + client_max_recv_message_size: ReadableSize, + client_max_send_message_size: ReadableSize, +) -> (MetaClient, MockMetaContext) { + let mut opts = MetasrvOptions::default(); + opts.grpc.server_addr = "127.0.0.1:3002".to_string(); + opts.grpc.max_recv_message_size = server_max_recv_message_size; + opts.grpc.max_send_message_size = server_max_send_message_size; + + let client_channel_config = ChannelConfig { + max_recv_message_size: client_max_recv_message_size, + max_send_message_size: client_max_send_message_size, + ..ChannelConfig::new() + }; + + let kv_backend = Arc::new(MemoryKvBackend::new()); + let in_memory = Arc::new(MemoryKvBackend::new()); + let MockInfo { + server_addr, + channel_manager, + kv_backend, + in_memory, + .. + } = server_mock::mock_with_client_channel_config( + opts, + kv_backend, + None, + None, + Some(in_memory), + client_channel_config, + ) + .await; + + ( + mock_client_by(server_addr, channel_manager).await, + MockMetaContext { + kv_backend, + in_memory, + }, + ) +} + #[allow(dead_code)] pub async fn mock_client_with_etcdstore(addr: &str) -> (MetaClient, MockMetaContext) { let MockInfo { diff --git a/src/meta-srv/src/bootstrap.rs b/src/meta-srv/src/bootstrap.rs index c451ca7c2d..e9607ae78d 100644 --- a/src/meta-srv/src/bootstrap.rs +++ b/src/meta-srv/src/bootstrap.rs @@ -233,13 +233,15 @@ pub async fn bootstrap_metasrv_with_router( #[macro_export] macro_rules! add_compressed_service { - ($builder:expr, $server:expr) => { + ($builder:expr, $server:expr, $grpc_config:expr) => { $builder.add_service( $server .accept_compressed(CompressionEncoding::Gzip) .accept_compressed(CompressionEncoding::Zstd) .send_compressed(CompressionEncoding::Gzip) - .send_compressed(CompressionEncoding::Zstd), + .send_compressed(CompressionEncoding::Zstd) + .max_decoding_message_size($grpc_config.max_recv_message_size) + .max_encoding_message_size($grpc_config.max_send_message_size), ) }; } @@ -252,18 +254,25 @@ pub fn router(metasrv: Arc) -> Router { .http2_keepalive_interval(Some(metasrv.options().grpc.http2_keep_alive_interval)) .http2_keepalive_timeout(Some(metasrv.options().grpc.http2_keep_alive_timeout)); let grpc_config = metasrv.options().grpc.as_config(); - let heartbeat_server = HeartbeatServer::from_arc(metasrv.clone()) - .accept_compressed(CompressionEncoding::Gzip) - .accept_compressed(CompressionEncoding::Zstd) - .send_compressed(CompressionEncoding::Gzip) - .send_compressed(CompressionEncoding::Zstd) - .max_decoding_message_size(grpc_config.max_recv_message_size) - .max_encoding_message_size(grpc_config.max_send_message_size); - let router = router.add_service(heartbeat_server); - let router = add_compressed_service!(router, StoreServer::from_arc(metasrv.clone())); - let router = add_compressed_service!(router, ClusterServer::from_arc(metasrv.clone())); - let router = add_compressed_service!(router, ProcedureServiceServer::from_arc(metasrv.clone())); - let router = add_compressed_service!(router, ConfigServer::from_arc(metasrv.clone())); + let router = add_compressed_service!( + router, + HeartbeatServer::from_arc(metasrv.clone()), + grpc_config + ); + let router = + add_compressed_service!(router, StoreServer::from_arc(metasrv.clone()), grpc_config); + let router = add_compressed_service!( + router, + ClusterServer::from_arc(metasrv.clone()), + grpc_config + ); + let router = add_compressed_service!( + router, + ProcedureServiceServer::from_arc(metasrv.clone()), + grpc_config + ); + let router = + add_compressed_service!(router, ConfigServer::from_arc(metasrv.clone()), grpc_config); router.add_service(admin::make_admin_service(metasrv)) } diff --git a/src/meta-srv/src/cluster.rs b/src/meta-srv/src/cluster.rs index ef3ba07702..688dc3f333 100644 --- a/src/meta-srv/src/cluster.rs +++ b/src/meta-srv/src/cluster.rs @@ -81,10 +81,7 @@ impl KvBackend for MetaPeerClient { let retry_interval_ms = self.retry_interval_ms; for _ in 0..max_retry_count { - match self - .remote_range(req.key.clone(), req.range_end.clone(), req.keys_only) - .await - { + match self.remote_range(req.clone()).await { Ok(res) => return Ok(res), Err(e) => { if need_retry(&e) { @@ -238,12 +235,7 @@ impl MetaPeerClient { to_stat_kv_map(res.kvs) } - async fn remote_range( - &self, - key: Vec, - range_end: Vec, - keys_only: bool, - ) -> Result { + async fn remote_range(&self, req: RangeRequest) -> Result { // Safety: when self.is_leader() == false, election must not empty. let election = self.election.as_ref().unwrap(); @@ -254,14 +246,11 @@ impl MetaPeerClient { .get(&leader_addr) .context(error::CreateChannelSnafu)?; - let request = tonic::Request::new(PbRangeRequest { - key, - range_end, - keys_only, - ..Default::default() - }); + let request = tonic::Request::new(PbRangeRequest::from(req)); - let response: PbRangeResponse = ClusterClient::new(channel) + let mut client = + common_grpc::configure_tonic_client!(ClusterClient::new(channel), self.channel_manager); + let response: PbRangeResponse = client .range(request) .await .context(error::RangeSnafu)? @@ -291,7 +280,9 @@ impl MetaPeerClient { ..Default::default() }); - let response: PbBatchGetResponse = ClusterClient::new(channel) + let mut client = + common_grpc::configure_tonic_client!(ClusterClient::new(channel), self.channel_manager); + let response: PbBatchGetResponse = client .batch_get(request) .await .context(error::BatchGetSnafu)? @@ -365,13 +356,168 @@ fn need_retry(error: &error::Error) -> bool { #[cfg(test)] mod tests { - use api::v1::meta::{Error, ErrorCode, ResponseHeader}; - use common_meta::datanode::{DatanodeStatKey, DatanodeStatValue, Stat}; - use common_meta::rpc::KeyValue; + use std::sync::Arc; + use std::sync::atomic::{AtomicI64, Ordering}; - use super::{Context, check_resp_header, to_stat_kv_map}; + use api::v1::meta::cluster_server::{Cluster, ClusterServer}; + use api::v1::meta::{ + BatchGetRequest as PbBatchGetRequest, BatchGetResponse as PbBatchGetResponse, Error, + ErrorCode, MetasrvPeersRequest, MetasrvPeersResponse, RangeRequest as PbRangeRequest, + RangeResponse as PbRangeResponse, ResponseHeader, + }; + use common_grpc::channel_manager::ChannelManager; + use common_meta::datanode::{DatanodeStatKey, DatanodeStatValue, Stat}; + use common_meta::election::{Election, LeaderChangeMessage, LeaderValue, MetasrvNodeInfo}; + use common_meta::kv_backend::KvBackend; + use common_meta::kv_backend::memory::MemoryKvBackend; + use common_meta::rpc::KeyValue; + use common_meta::rpc::store::RangeRequest; + use hyper_util::rt::TokioIo; + use tonic::{Request, Response, Status}; + use tower::service_fn; + + use super::{Context, MetaPeerClientBuilder, check_resp_header, to_stat_kv_map}; use crate::error; + struct FollowerElection { + leader_addr: String, + } + + #[async_trait::async_trait] + impl Election for FollowerElection { + type Leader = LeaderValue; + + fn is_leader(&self) -> bool { + false + } + + fn in_leader_infancy(&self) -> bool { + false + } + + async fn register_candidate(&self, _: &MetasrvNodeInfo) -> common_meta::error::Result<()> { + Ok(()) + } + + async fn all_candidates(&self) -> common_meta::error::Result> { + Ok(vec![]) + } + + async fn campaign(&self) -> common_meta::error::Result<()> { + Ok(()) + } + + async fn leader(&self) -> common_meta::error::Result { + Ok(LeaderValue(self.leader_addr.clone())) + } + + async fn resign(&self) -> common_meta::error::Result<()> { + Ok(()) + } + + fn subscribe_leader_change(&self) -> tokio::sync::broadcast::Receiver { + let (_, receiver) = tokio::sync::broadcast::channel(1); + receiver + } + } + + struct RangeServer { + requested_limit: Arc, + } + + #[async_trait::async_trait] + impl Cluster for RangeServer { + async fn batch_get( + &self, + _: Request, + ) -> std::result::Result, Status> { + Err(Status::unimplemented("batch_get is not used in this test")) + } + + async fn range( + &self, + request: Request, + ) -> std::result::Result, Status> { + self.requested_limit + .store(request.into_inner().limit, Ordering::Relaxed); + Ok(Response::new(PbRangeResponse { + header: Some(ResponseHeader::success()), + kvs: vec![api::v1::meta::KeyValue { + key: b"key".to_vec(), + value: b"value".to_vec(), + }], + more: true, + })) + } + + async fn metasrv_peers( + &self, + _: Request, + ) -> std::result::Result, Status> { + Err(Status::unimplemented( + "metasrv_peers is not used in this test", + )) + } + } + + #[tokio::test] + async fn test_follower_range_forwards_limit() { + let requested_limit = Arc::new(AtomicI64::new(0)); + let range_server = RangeServer { + requested_limit: requested_limit.clone(), + }; + let (client, server) = tokio::io::duplex(1024); + let _server_handle = tokio::spawn(async move { + tonic::transport::Server::builder() + .add_service( + ClusterServer::new(range_server) + .accept_compressed(tonic::codec::CompressionEncoding::Zstd) + .send_compressed(tonic::codec::CompressionEncoding::Zstd), + ) + .serve_with_incoming(futures::stream::iter([Ok::<_, std::io::Error>(server)])) + .await + }); + + let channel_manager = ChannelManager::new(); + let mut client = Some(client); + channel_manager + .reset_with_connector( + "leader:0", + service_fn(move |_| { + let client = client.take(); + async move { + client + .map(TokioIo::new) + .ok_or_else(|| std::io::Error::other("client already taken")) + } + }), + ) + .unwrap(); + + let follower = MetaPeerClientBuilder::default() + .election(Some(Arc::new(FollowerElection { + leader_addr: "leader:0".to_string(), + }))) + .in_memory(Arc::new(MemoryKvBackend::new())) + .channel_manager(channel_manager) + .max_retry_count(1) + .build() + .unwrap(); + + let response = follower + .range(RangeRequest { + key: b"key".to_vec(), + limit: 1, + ..Default::default() + }) + .await + .unwrap(); + + assert_eq!(1, requested_limit.load(Ordering::Relaxed)); + assert_eq!(1, response.kvs.len()); + assert!(response.more); + } + #[test] fn test_to_stat_kv_map() { let stat_key = DatanodeStatKey { node_id: 100 }; diff --git a/src/meta-srv/src/mocks.rs b/src/meta-srv/src/mocks.rs index 26b1687e58..1a6fa30674 100644 --- a/src/meta-srv/src/mocks.rs +++ b/src/meta-srv/src/mocks.rs @@ -87,12 +87,52 @@ pub async fn mock( selector: Option, datanode_clients: Option>, in_memory: Option, +) -> MockInfo { + mock_inner( + opts, + kv_backend, + selector, + datanode_clients, + in_memory, + None, + ) + .await +} + +pub async fn mock_with_client_channel_config( + opts: MetasrvOptions, + kv_backend: KvBackendRef, + selector: Option, + datanode_clients: Option>, + in_memory: Option, + client_channel_config: ChannelConfig, +) -> MockInfo { + mock_inner( + opts, + kv_backend, + selector, + datanode_clients, + in_memory, + Some(client_channel_config), + ) + .await +} + +async fn mock_inner( + opts: MetasrvOptions, + kv_backend: KvBackendRef, + selector: Option, + datanode_clients: Option>, + in_memory: Option, + client_channel_config: Option, ) -> MockInfo { let server_addr = opts.grpc.server_addr.clone(); let table_metadata_manager = Arc::new(TableMetadataManager::new(kv_backend.clone())); table_metadata_manager.init().await.unwrap(); + let grpc_config = opts.grpc.as_config(); + let grpc_options = opts.grpc.clone(); let builder = MetasrvBuilder::new() .options(opts) .kv_backend(kv_backend.clone()); @@ -121,22 +161,43 @@ pub async fn mock( let _handle = tokio::spawn(async move { let mut router = tonic::transport::Server::builder(); - let router = add_compressed_service!(router, HeartbeatServer::from_arc(service.clone())); - let router = add_compressed_service!(router, StoreServer::from_arc(service.clone())); + let router = add_compressed_service!( + router, + HeartbeatServer::from_arc(service.clone()), + grpc_config + ); let router = - add_compressed_service!(router, ProcedureServiceServer::from_arc(service.clone())); - let router = add_compressed_service!(router, ClusterServer::from_arc(service.clone())); - let router = add_compressed_service!(router, ConfigServer::from_arc(service.clone())); + add_compressed_service!(router, StoreServer::from_arc(service.clone()), grpc_config); + let router = add_compressed_service!( + router, + ProcedureServiceServer::from_arc(service.clone()), + grpc_config + ); + let router = add_compressed_service!( + router, + ClusterServer::from_arc(service.clone()), + grpc_config + ); + let router = + add_compressed_service!(router, ConfigServer::from_arc(service.clone()), grpc_config); router .serve_with_incoming(futures::stream::iter(vec![Ok::<_, std::io::Error>(server)])) .await }); - let config = ChannelConfig::new() - // Use an long timeout to prevent test failures due to slow operations (e.g., when testing with S3). - .timeout(Some(Duration::from_secs(60))) - .connect_timeout(Duration::from_secs(10)) - .tcp_nodelay(true); + // Keep the mock client's codec limits aligned with the server by default. + let config = client_channel_config.unwrap_or_else(|| { + let config = ChannelConfig::new() + // Use an long timeout to prevent test failures due to slow operations (e.g., when testing with S3). + .timeout(Some(Duration::from_secs(60))) + .connect_timeout(Duration::from_secs(10)) + .tcp_nodelay(true); + ChannelConfig { + max_recv_message_size: grpc_options.max_recv_message_size, + max_send_message_size: grpc_options.max_send_message_size, + ..config + } + }); let channel_manager = ChannelManager::with_config(config, None); // Move client to an option so we can _move_ the inner value