raftstore: restrict the total write size of each apply round

components/raftstore/src/store/config.rs

@@ -116,6 +116,7 @@ pub struct Config {
     #[online_config(skip)]
     pub notify_capacity: usize,
     pub messages_per_tick: usize,
+    pub messages_size_per_tick: usize,
 
     /// When a peer is not active for max_peer_down_duration,
     /// the peer is considered to be down and is reported to PD.
@@ -346,6 +347,7 @@ impl Default for Config {
             snap_mgr_gc_tick_interval: ReadableDuration::minutes(1),
             snap_gc_timeout: ReadableDuration::hours(4),
             messages_per_tick: 4096,
+            messages_size_per_tick: 32768,
             max_peer_down_duration: ReadableDuration::minutes(10),
             max_leader_missing_duration: ReadableDuration::hours(2),
             abnormal_leader_missing_duration: ReadableDuration::minutes(10),
@@ -838,6 +840,9 @@ impl Config {
         CONFIG_RAFTSTORE_GAUGE
             .with_label_values(&["messages_per_tick"])
             .set(self.messages_per_tick as f64);
+        CONFIG_RAFTSTORE_GAUGE
+            .with_label_values(&["messages_size_per_tick"])
+            .set(self.messages_size_per_tick as f64);
 
         CONFIG_RAFTSTORE_GAUGE
             .with_label_values(&["max_peer_down_duration"])

components/raftstore/src/store/fsm/apply.rs

@@ -3386,6 +3386,13 @@ where
             merge_from_snapshot,
         })
     }
+
+    pub fn entries_size(&self) -> usize {
+        match self {
+            Msg::Apply { apply, .. } => apply.entries_size,
+            _ => 0,
+        }
+    }
 }
 
 impl<EK> Debug for Msg<EK>
@@ -4043,6 +4050,7 @@ where
     msg_buf: Vec<Msg<EK>>,
     apply_ctx: ApplyContext<EK>,
     messages_per_tick: usize,
+    messages_size_per_tick: usize,
     cfg_tracker: Tracker<Config>,
 
     trace_event: TraceEvent,
@@ -4068,6 +4076,7 @@ where
                 }
                 _ => {}
             }
+            self.messages_size_per_tick = incoming.messages_size_per_tick;
             update_cfg(&incoming.apply_batch_system);
         }
     }
@@ -4112,9 +4121,15 @@ where
             normal.delegate.id() == 1003,
             |_| { HandleResult::KeepProcessing }
         );
-        while self.msg_buf.len() < self.messages_per_tick {
+        let mut total_size = 0;
+        while self.msg_buf.len() < self.messages_per_tick
+            && total_size < self.messages_size_per_tick
+        {
             match normal.receiver.try_recv() {
-                Ok(msg) => self.msg_buf.push(msg),
+                Ok(msg) => {
+                    total_size += msg.entries_size();
+                    self.msg_buf.push(msg);
+                }
                 Err(TryRecvError::Empty) => {
                     handle_result = HandleResult::stop_at(0, false);
                     break;
@@ -4211,6 +4226,7 @@ where
                 priority,
             ),
             messages_per_tick: cfg.messages_per_tick,
+            messages_size_per_tick: cfg.messages_size_per_tick,
             cfg_tracker: self.cfg.clone().tracker(self.tag.clone()),
             trace_event: Default::default(),
         }

components/raftstore/src/store/peer.rs

@@ -2955,55 +2955,74 @@ where
                 self.raft_group.skip_bcast_commit(true);
                 self.last_urgent_proposal_idx = u64::MAX;
             }
-            let cbs = if !self.proposals.is_empty() {
-                let current_term = self.term();
-                let cbs = committed_entries
-                    .iter()
-                    .filter_map(|e| {
-                        self.proposals
-                            .find_proposal(e.get_term(), e.get_index(), current_term)
-                    })
-                    .map(|mut p| {
-                        if p.must_pass_epoch_check {
-                            // In this case the apply can be guaranteed to be successful. Invoke the
-                            // on_committed callback if necessary.
-                            p.cb.invoke_committed();
-                        }
-                        p
-                    })
-                    .collect();
+        }
+        let mut batch_size = 0;
+        let mut entry_batch = Vec::with_capacity(committed_entries.len());
+        let length = committed_entries.len();
+        let has_entry = !committed_entries.is_empty();
+        let has_proposal = !self.proposals.is_empty();
+        for (i, entry) in committed_entries.into_iter().enumerate() {
+            batch_size += entry.get_data().len();
+            entry_batch.push(entry);
+            if batch_size >= ctx.cfg.messages_size_per_tick
+                || entry_batch.len() >= ctx.cfg.messages_per_tick
+                || i == length - 1
+            {
+                let entries = mem::take(&mut entry_batch);
+                batch_size = 0;
+                let cbs = if !self.proposals.is_empty() {
+                    let current_term = self.term();
+                    entries
+                        .iter()
+                        .filter_map(|e| {
+                            self.proposals
+                                .find_proposal(e.get_term(), e.get_index(), current_term)
+                        })
+                        .map(|mut p| {
+                            if p.must_pass_epoch_check {
+                                // In this case the apply can be guaranteed to be successful. Invoke
+                                // the on_committed callback if necessary.
+                                p.cb.invoke_committed();
+                            }
+                            p
+                        })
+                        .collect()
+                } else {
+                    vec![]
+                };
+                // Note that the `commit_index` and `commit_term` here may be used to
+                // forward the commit index. So it must be less than or equal to persist
+                // index.
+                let commit_index = cmp::min(
+                    self.raft_group.raft.raft_log.committed,
+                    self.raft_group.raft.raft_log.persisted,
+                );
+                let commit_term = self.get_store().term(commit_index).unwrap();
+                let mut apply = Apply::new(
+                    self.peer_id(),
+                    self.region_id,
+                    self.term(),
+                    commit_index,
+                    commit_term,
+                    entries,
+                    cbs,
+                    self.region_buckets.as_ref().map(|b| b.meta.clone()),
+                );
+                apply.on_schedule(&ctx.raft_metrics);
+                self.mut_store()
+                    .trace_cached_entries(apply.entries[0].clone());
+                ctx.apply_router
+                    .schedule_task(self.region_id, ApplyTask::apply(apply));
+            }
+        }
+        if has_entry {
+            if has_proposal {
                 self.proposals.gc();
-                cbs
-            } else {
-                vec![]
-            };
-            // Note that the `commit_index` and `commit_term` here may be used to
-            // forward the commit index. So it must be less than or equal to persist
-            // index.
-            let commit_index = cmp::min(
-                self.raft_group.raft.raft_log.committed,
-                self.raft_group.raft.raft_log.persisted,
-            );
-            let commit_term = self.get_store().term(commit_index).unwrap();
-            let mut apply = Apply::new(
-                self.peer_id(),
-                self.region_id,
-                self.term(),
-                commit_index,
-                commit_term,
-                committed_entries,
-                cbs,
-                self.region_buckets.as_ref().map(|b| b.meta.clone()),
-            );
-            apply.on_schedule(&ctx.raft_metrics);
-            self.mut_store()
-                .trace_cached_entries(apply.entries[0].clone());
+            }
             if needs_evict_entry_cache(ctx.cfg.evict_cache_on_memory_ratio) {
                 // Compact all cached entries instead of half evict.
                 self.mut_store().evict_entry_cache(false);
             }
-            ctx.apply_router
-                .schedule_task(self.region_id, ApplyTask::apply(apply));
         }
         fail_point!("after_send_to_apply_1003", self.peer_id() == 1003, |_| {});
     }

tests/integrations/config/dynamic/raftstore.rs

@@ -157,6 +157,7 @@ fn test_update_raftstore_config() {
     // dispatch updated config
     let change = new_changes(vec![
         ("raftstore.messages-per-tick", "12345"),
+        ("raftstore.messages-size-per-tick", "123456"),
         ("raftstore.raft-log-gc-threshold", "54321"),
         ("raftstore.raft-max-size-per-msg", "128MiB"),
         ("raftstore.apply-max-batch-size", "1234"),
@@ -169,6 +170,7 @@ fn test_update_raftstore_config() {
     // config should be updated
     let mut raft_store = config.raft_store;
     raft_store.messages_per_tick = 12345;
+    raft_store.messages_size_per_tick = 123456;
     raft_store.raft_log_gc_threshold = 54321;
     raft_store.apply_batch_system.max_batch_size = Some(1234);
     raft_store.store_batch_system.max_batch_size = Some(4321);

tests/integrations/config/mod.rs

@@ -199,6 +199,7 @@ fn test_serde_custom_tikv_config() {
         snap_mgr_gc_tick_interval: ReadableDuration::minutes(12),
         snap_gc_timeout: ReadableDuration::hours(12),
         messages_per_tick: 12_345,
+        messages_size_per_tick: 123_456,
         max_peer_down_duration: ReadableDuration::minutes(12),
         max_leader_missing_duration: ReadableDuration::hours(12),
         abnormal_leader_missing_duration: ReadableDuration::hours(6),

tests/integrations/config/test-custom.toml

@@ -173,6 +173,7 @@ lock-cf-compact-interval = "12m"
 lock-cf-compact-bytes-threshold = "123MB"
 notify-capacity = 12345
 messages-per-tick = 12345
+messages-size-per-tick = 123456
 max-peer-down-duration = "12m"
 max-leader-missing-duration = "12h"
 abnormal-leader-missing-duration = "6h"