store: Make "grpc_max_decoding_message_size" configurable (#485)

Signed-off-by: Ping Yu <[email protected]>

Author: pingyu

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
+    }
 }

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,

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))
     }