raftclient: count term and index in estimated size (tikv#11493)

* add test case

Signed-off-by: Jay Lee <[email protected]>

* count term and index tag

If there are many entries in a message, the estimated size of message
can be way smaller than the actual size. This PR fixes the error by
also counting index and term in estimation.

It also remove the hard limit as the estimation is closed enough.

Close tikv#9714.

Signed-off-by: Jay Lee <[email protected]>

Co-authored-by: Ti Chi Robot <[email protected]>

Author: BusyJay

src/server/raft_client.rs

@@ -174,7 +174,9 @@ impl Buffer for BatchMessageBuffer {
         let mut msg_size = msg.start_key.len()
             + msg.end_key.len()
             + msg.get_message().context.len()
-            + msg.extra_ctx.len();
+            + msg.extra_ctx.len()
+            // index: 3, term: 2, data tag and size: 3, entry tag and size: 3
+            + 11 * msg.get_message().get_entries().len();
         for entry in msg.get_message().get_entries() {
             msg_size += entry.data.len();
         }
@@ -568,7 +570,6 @@ where
 
         let cb = ChannelBuilder::new(self.builder.env.clone())
             .stream_initial_window_size(self.builder.cfg.grpc_stream_initial_window_size.0 as i32)
-            .max_send_message_len(self.builder.cfg.max_grpc_send_msg_len)
             .keepalive_time(self.builder.cfg.grpc_keepalive_time.0)
             .keepalive_timeout(self.builder.cfg.grpc_keepalive_timeout.0)
             .default_compression_algorithm(self.builder.cfg.grpc_compression_algorithm())

tests/integrations/server/raft_client.rs

@@ -217,6 +217,45 @@ fn test_batch_size_limit() {
     assert_eq!(msg_count.load(Ordering::SeqCst), 10);
 }
 
+/// In edge case that the estimated size may be inaccurate, we need to ensure connection
+/// will not be broken in this case.
+#[test]
+fn test_batch_size_edge_limit() {
+    let msg_count = Arc::new(AtomicUsize::new(0));
+    let batch_msg_count = Arc::new(AtomicUsize::new(0));
+    let service = MockKvForRaft::new(Arc::clone(&msg_count), Arc::clone(&batch_msg_count), true);
+    let (mock_server, port) = create_mock_server(service, 60200, 60300).unwrap();
+
+    let mut raft_client = get_raft_client_by_port(port);
+
+    // Put them in buffer so sibling messages will be likely be batched during sending.
+    let mut msgs = Vec::with_capacity(5);
+    for _ in 0..5 {
+        let mut raft_m = RaftMessage::default();
+        // Magic number, this can make estimated size about 4940000, hence two messages will be
+        // batched together, but the total size will be way largher than 10MiB as there are many
+        // indexes and terms.
+        for _ in 0..38000 {
+            let mut e = Entry::default();
+            e.set_term(1);
+            e.set_index(256);
+            e.set_data(vec![b'a'; 130].into());
+            raft_m.mut_message().mut_entries().push(e);
+        }
+        msgs.push(raft_m);
+    }
+    for m in msgs {
+        raft_client.send(m).unwrap();
+    }
+    raft_client.flush();
+
+    check_msg_count(10000, &msg_count, 5);
+    // The final received message count should be 5 exactly.
+    drop(raft_client);
+    drop(mock_server);
+    assert_eq!(msg_count.load(Ordering::SeqCst), 5);
+}
+
 // Try to create a mock server with `service`. The server will be binded wiht a random
 // port chosen between [`min_port`, `max_port`]. Return `None` if no port is available.
 fn create_mock_server<T>(service: T, min_port: u16, max_port: u16) -> Option<(Server, u16)>