From 1aebfa56ad4678cf08121e264cbd623c0a7b6582 Mon Sep 17 00:00:00 2001
From: Ping Yu <yuping@pingcap.com>
Date: Fri, 11 Apr 2025 08:32:58 +0800
Subject: [PATCH] store: Make "grpc_max_decoding_message_size" configurable
 (#485)

Signed-off-by: Ping Yu <yuping@pingcap.com>
---
 src/config.rs       | 10 ++++++++++
 src/pd/client.rs    |  8 +++++++-
 src/store/client.rs |  6 +++++-
 3 files changed, 22 insertions(+), 2 deletions(-)

diff --git a/src/config.rs b/src/config.rs
index 1be273c..99451e4 100644
--- a/src/config.rs
+++ b/src/config.rs
@@ -19,9 +19,11 @@ pub struct Config {
     pub cert_path: Option<PathBuf>,
     pub key_path: Option<PathBuf>,
     pub timeout: Duration,
+    pub grpc_max_decoding_message_size: usize,
 }
 
 const DEFAULT_REQUEST_TIMEOUT: Duration = Duration::from_secs(2);
+const DEFAULT_GRPC_MAX_DECODING_MESSAGE_SIZE: usize = 4 * 1024 * 1024; // 4MB
 
 impl Default for Config {
     fn default() -> Self {
@@ -30,6 +32,7 @@ impl Default for Config {
             cert_path: None,
             key_path: None,
             timeout: DEFAULT_REQUEST_TIMEOUT,
+            grpc_max_decoding_message_size: DEFAULT_GRPC_MAX_DECODING_MESSAGE_SIZE,
         }
     }
 }
@@ -83,4 +86,11 @@ impl Config {
         self.timeout = timeout;
         self
     }
+
+    /// Set the maximum decoding message size for gRPC.
+    #[must_use]
+    pub fn with_grpc_max_decoding_message_size(mut self, size: usize) -> Self {
+        self.grpc_max_decoding_message_size = size;
+        self
+    }
 }
diff --git a/src/pd/client.rs b/src/pd/client.rs
index a2469b6..69060ec 100644
--- a/src/pd/client.rs
+++ b/src/pd/client.rs
@@ -306,7 +306,13 @@ impl<Cod: Codec> PdRpcClient<Cod, TikvConnect, Cluster> {
     ) -> Result<PdRpcClient<Cod>> {
         PdRpcClient::new(
             config.clone(),
-            |security_mgr| TikvConnect::new(security_mgr, config.timeout),
+            |security_mgr| {
+                TikvConnect::new(
+                    security_mgr,
+                    config.timeout,
+                    config.grpc_max_decoding_message_size,
+                )
+            },
             |security_mgr| RetryClient::connect(pd_endpoints, security_mgr, config.timeout),
             enable_mvcc_codec,
             codec,
diff --git a/src/store/client.rs b/src/store/client.rs
index 363d413..88880c7 100644
--- a/src/store/client.rs
+++ b/src/store/client.rs
@@ -25,6 +25,7 @@ pub trait KvConnect: Sized + Send + Sync + 'static {
 pub struct TikvConnect {
     security_mgr: Arc<SecurityManager>,
     timeout: Duration,
+    grpc_max_decoding_message_size: usize,
 }
 
 #[async_trait]
@@ -33,7 +34,10 @@ impl KvConnect for TikvConnect {
 
     async fn connect(&self, address: &str) -> Result<KvRpcClient> {
         self.security_mgr
-            .connect(address, TikvClient::new)
+            .connect(address, move |channel| {
+                TikvClient::new(channel)
+                    .max_decoding_message_size(self.grpc_max_decoding_message_size)
+            })
             .await
             .map(|c| KvRpcClient::new(c, self.timeout))
     }